Change to summarization.

main.py

@@ -5,6 +5,9 @@ from langchain.text_splitter import CharacterTextSplitter
 from langchain.embeddings import OpenAIEmbeddings
 from langchain.vectorstores import Chroma
 from langchain.document_loaders import PyPDFLoader
+from langchain import PromptTemplate
+from langchain.chains.summarize import load_summarize_chain
+import textwrap
 import panel as pn
 import PyPDF2
 
@@ -24,6 +27,8 @@ prompt = pn.widgets.TextEditor(
 )
 run_button = pn.widgets.Button(name="Run!")
 
+summary_button = pn.widgets.Button(name="Resumo!")
+
 select_k = pn.widgets.IntSlider(
     name="Number of relevant chunks", start=1, end=5, step=1, value=2
 )
@@ -41,6 +46,10 @@ widgets = pn.Row(
     ), width=600
 )
 
+summary_filed = pn.Row(
+    pn.Column(summary_button),
+    width=630
+)
 
 def is_valid_pdf(file_path):
     try:
@@ -74,12 +83,57 @@ def qa(file, query, chain_type, k):
     # create a chain to answer questions
 
     qa = RetrievalQA.from_chain_type(
-        llm=OpenAI(temperature=0), chain_type=chain_type, retriever=retriever, return_source_documents=True)
+        llm=OpenAI(model_name="gpt-3.5-turbo", temperature=0), chain_type=chain_type, retriever=retriever, return_source_documents=False)
     result = qa({"query": query})
     print(result['result'])
     return result
 
 
+def summary(file):
+    # load document
+    result = {}
+    if not is_valid_pdf(file):
+        result = {'error': 'Invalid PDF file.'}
+        return result
+
+    loader = PyPDFLoader(file)
+    documents = loader.load()
+
+    combine_template = """Write a summary of the following in Portuguese in 100 words:
+
+    {text} 
+
+    SUMMARY IN PORTUGUESE IN 100 WORDS:"""
+    COMBINE_TEMPLATE = PromptTemplate(template=combine_template, input_variables=["text"])
+
+    map_template = """Write a concise summary of the following in Portuguese in 40 words or less:
+
+    {text} 
+
+    CONCISE SUMMARY IN PORTUGUESE IN 40 WORDS OR LESS:"""
+    MAP_TEMPLATE = PromptTemplate(template=map_template, input_variables=["text"])
+
+    chain = load_summarize_chain(OpenAI(temperature=0),
+                                 chain_type="map_reduce",
+                                 return_intermediate_steps=True,
+                                 combine_prompt=COMBINE_TEMPLATE,
+                                 map_prompt=MAP_TEMPLATE)
+
+    output_summary = chain({"input_documents": documents}, return_only_outputs=True)
+    result['summary'] = textwrap.fill(output_summary['output_text'],
+                                 width=100,
+                                 break_long_words=False,
+                                 replace_whitespace=False)
+
+    output_steps = output_summary['intermediate_steps']
+    result['steps'] = textwrap.fill('\n'.join(output_steps),
+                                  width=100,
+                                  break_long_words=False,
+                                  replace_whitespace=False)
+    return result
+
+
+
 convos = []  # store all panel objects in a list
 
 
@@ -91,11 +145,11 @@ def qa_result(_):
 
     # save pdf file to a temp file
     if file_input.value is not None:
-        file_input.save("/.cache/temp.pdf")
+        file_input.save("temp.pdf")
 
         prompt_text = prompt.value
         if prompt_text:
-            result = qa(file="/.cache/temp.pdf", query=prompt_text, chain_type=select_chain_type.value,
+            result = qa(file="temp.pdf", query=prompt_text, chain_type=select_chain_type.value,
                         k=select_k.value)
             if result.get('error') is None:
                 convos.extend([
@@ -123,9 +177,43 @@ def qa_result(_):
         pn.state.notifications.error('Missing file.', duration=2000)
     return pn.Column(*convos, margin=15, width=575, min_height=400)
 
+def summary_result(_):
+    os.environ["OPENAI_API_KEY"] = openaikey.value
+    if not openaikey.value:
+        pn.state.notifications.error('Missing API key.', duration=2000)
+        return pn.Column(*convos, margin=15, width=575, min_height=400)
+
+    # save pdf file to a temp file
+    if file_input.value is not None:
+        file_input.save("temp.pdf")
+
+        result = summary(file="temp.pdf")
+        if result.get('error') is None:
+            convos.extend([
+                pn.Row(
+                    pn.panel("\U0001F60A", width=10),
+                     "Resumo geral: ",
+                    result['summary'],
+                    width=600
+                ),
+                pn.Row(
+                    pn.panel("\U0001F916", width=10),
+                    pn.Column(
+                        "Resumo por página:",
+                        result['steps']
+                    )
+                )
+            ])
+        else:
+            pn.state.notifications.error(result['error'], duration=2000)
+    else:
+        pn.state.notifications.error('Missing file.', duration=2000)
+    return pn.Column(*convos, margin=15, width=575, min_height=400)
+
 
 qa_interactive = pn.panel(
-    pn.bind(qa_result, run_button),
+    #pn.bind(qa_result, run_button),
+    pn.bind(summary_result, summary_button),
     loading_indicator=True,
 )
 
@@ -137,11 +225,12 @@ pn.Column(
     ## \U0001F4D3 Perguntas e respostas com um PDF 
     (original implementation: @sophiamyang)
 
-    1) Suba o PDF. 2) Entre com a OpenAI API key. 3) Digite a pergunta e clique "Run".
+    1) Suba o PDF. 2) Entre com a OpenAI API key. 3) Clique "Resumo!".
 
     """),
     pn.Row(file_input, openaikey),
+    summary_filed,
     output,
-    widgets
+    #widgets
 
 ).servable()