Update app.py

app.py

@@ -15,6 +15,7 @@ from langchain.document_loaders.pdf import PyMuPDFLoader
 import os
 #import fitz
 #import tempfile
+from langchain.chains.summarize import load_summarize_chain
 
 img = Image.open('image/nexio_logo1.png')
 st.set_page_config(page_title="PDF Chatbot App",page_icon=img,layout="centered")
@@ -59,21 +60,26 @@ def main():
         
         # Accept user question
         query = st.text_input("Ask questions about your PDF file:")
-
+        
         if query:
             
             #PATH = 'model/'
             #llm = AutoModelForCausalLM.from_pretrained("CohereForAI/aya-101")
             # llm = AutoModelForCausalLM.from_pretrained(PATH,local_files_only=True)
             llm = huggingface_hub.HuggingFaceHub(repo_id="google/flan-t5-small",
-                                                  model_kwargs={"temperature":1.0, "max_length":500})
-            docs = vector_store.similarity_search(query=query, k=3)
-            st.write(docs)
-            chain = load_qa_chain(llm=llm, chain_type="stuff")
-            response = chain.run(input_documents=docs, question=query)
-            #retriever=vector_store.as_retriever()
-            #chain = RetrievalQA.from_chain_type(llm=llm,chain_type="stuff",retriever=retriever)
-            #response = chain.run(chain)
+                                                  model_kwargs={"temperature":1.0, "max_length":256})
+            if query == 'Summarize':
+                docs = pdf_reader.load_and_split()
+                chain = load_summarize_chain(llm, chain_type="map_reduce")
+                response = chain.run(docs)
+            else:
+                docs = vector_store.similarity_search(query=query, k=2)
+                #st.write(docs)
+                chain = load_qa_chain(llm=llm, chain_type="stuff")
+                response = chain.run(input_documents=docs, question=query)
+                #retriever=vector_store.as_retriever()
+                #chain = RetrievalQA.from_chain_type(llm=llm,chain_type="stuff",retriever=retriever)
+                #response = chain.run(chain)
             st.write(response)