添加总结按钮

ChuanhuChatbot.py

@@ -97,6 +97,7 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
                     )
                     index_files = gr.Files(label=i18n("上传"), type="file")
                     two_column = gr.Checkbox(label=i18n("双栏pdf"), value=advance_docs["pdf"].get("two_column", False))
+                    summarize_btn = gr.Button(i18n("总结"))
                     # TODO: 公式ocr
                     # formula_ocr = gr.Checkbox(label=i18n("识别公式"), value=advance_docs["pdf"].get("formula_ocr", False))
 
@@ -333,6 +334,7 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
     submitBtn.click(**get_usage_args)
 
     index_files.change(handle_file_upload, [current_model, index_files, chatbot, language_select_dropdown], [index_files, chatbot, status_display])
+    summarize_btn.click(handle_summarize_index, [current_model, index_files, chatbot, language_select_dropdown], [chatbot, status_display])
 
     emptyBtn.click(
         reset,

modules/models/base_model.py

@@ -263,22 +263,28 @@ class BaseLLMModel:
         if files:
             index = construct_index(self.api_key, file_src=files)
             status = i18n("索引构建完成")
-            # Summarize the document
-            # logging.info(i18n("生成内容总结中……"))
-            # os.environ["OPENAI_API_KEY"] = self.api_key
-            # from langchain.chains.summarize import load_summarize_chain
-            # from langchain.prompts import PromptTemplate
-            # from langchain.chat_models import ChatOpenAI
-            # from langchain.callbacks import StdOutCallbackHandler
-            # prompt_template = "Write a concise summary of the following:\n\n{text}\n\nCONCISE SUMMARY IN " + language + ":"
-            # PROMPT = PromptTemplate(template=prompt_template, input_variables=["text"])
-            # llm = ChatOpenAI()
-            # chain = load_summarize_chain(llm, chain_type="map_reduce", return_intermediate_steps=True, map_prompt=PROMPT, combine_prompt=PROMPT)
-            # summary = chain({"input_documents": list(index.docstore.__dict__["_dict"].values())}, return_only_outputs=True)["output_text"]
-            # print(i18n("总结") + f": {summary}")
-            # chatbot.append([i18n("上传了")+str(len(files))+"个文件", summary])
         return gr.Files.update(), chatbot, status
 
+    def summarize_index(self, files, chatbot, language):
+        status = gr.Markdown.update()
+        if files:
+            index = construct_index(self.api_key, file_src=files)
+            status = i18n("总结完成")
+            logging.info(i18n("生成内容总结中……"))
+            os.environ["OPENAI_API_KEY"] = self.api_key
+            from langchain.chains.summarize import load_summarize_chain
+            from langchain.prompts import PromptTemplate
+            from langchain.chat_models import ChatOpenAI
+            from langchain.callbacks import StdOutCallbackHandler
+            prompt_template = "Write a concise summary of the following:\n\n{text}\n\nCONCISE SUMMARY IN " + language + ":"
+            PROMPT = PromptTemplate(template=prompt_template, input_variables=["text"])
+            llm = ChatOpenAI()
+            chain = load_summarize_chain(llm, chain_type="map_reduce", return_intermediate_steps=True, map_prompt=PROMPT, combine_prompt=PROMPT)
+            summary = chain({"input_documents": list(index.docstore.__dict__["_dict"].values())}, return_only_outputs=True)["output_text"]
+            print(i18n("总结") + f": {summary}")
+            chatbot.append([i18n("上传了")+str(len(files))+"个文件", summary])
+        return chatbot, status
+
     def prepare_inputs(self, real_inputs, use_websearch, files, reply_language, chatbot):
         fake_inputs = None
         display_append = []

modules/utils.py

@@ -116,6 +116,9 @@ def set_single_turn(current_model, *args):
 def handle_file_upload(current_model, *args):
     return current_model.handle_file_upload(*args)
 
+def handle_summarize_index(current_model, *args):
+    return current_model.summarize_index(*args)
+
 def like(current_model, *args):
     return current_model.like(*args)