功能优化: 添加双栏pdf识别选项到页面,并优化config文件中关于文档解析的设置

ChuanhuChatbot.py

@@ -78,6 +78,10 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
                         value=REPLY_LANGUAGES[0],
                     )
                     index_files = gr.Files(label="上传索引文件", type="file", multiple=True)
+                    two_column = gr.Checkbox(label="双栏pdf", value=advance_docs["pdf"].get("two_column", False))
+                    # TODO: 公式ocr
+                    # formula_ocr = gr.Checkbox(label="识别公式", value=advance_docs["pdf"].get("formula_ocr", False))
+                    updateDocConfigBtn = gr.Button("更新解析文件参数")
 
                 with gr.Tab(label="Prompt"):
                     systemPromptTxt = gr.Textbox(
@@ -295,6 +299,8 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
     )
     reduceTokenBtn.click(**get_usage_args)
 
+    updateDocConfigBtn.click(update_doc_config, [two_column], None)
+
     # ChatGPT
     keyTxt.change(submit_key, keyTxt, [user_api_key, status_display]).then(**get_usage_args)
     keyTxt.submit(**get_usage_args)

config_example.json

@@ -2,8 +2,11 @@
     "openai_api_key": "sk-xxxxxxxxxxxxxxxxxxxxxxxxx",
     "https_proxy": "http://127.0.0.1:1079",
     "http_proxy": "http://127.0.0.1:1079",
-    "advanced_pdf_kwargs": {
-        "two_column": true
+    "advance_docs": {
+        "pdf": {
+            "two_column": true,
+            "formula_ocr": true
+        }
     },
     "users": [
         ["root", "root"]

modules/chat_func.py

@@ -291,12 +291,14 @@ def predict(
         msg = "索引构建完成，获取回答中……"
         logging.info(msg)
         yield chatbot+[(inputs, "")], history, msg, all_token_counts
-        llm_predictor = LLMPredictor(llm=OpenAIChat(temperature=0, model_name=selected_model))
-        prompt_helper = PromptHelper(max_input_size = 4096, num_output = 5, max_chunk_overlap = 20, chunk_size_limit=600)
-        service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)
-        query_object = GPTVectorStoreIndexQuery(index.index_struct, service_context=service_context, similarity_top_k=5, vector_store=index._vector_store, docstore=index._docstore)
-        query_bundle = QueryBundle(inputs)
-        nodes = query_object.retrieve(query_bundle)
+        with retrieve_proxy():
+            llm_predictor = LLMPredictor(llm=OpenAIChat(temperature=0, model_name=selected_model))
+            prompt_helper = PromptHelper(max_input_size = 4096, num_output = 5, max_chunk_overlap = 20, chunk_size_limit=600)
+            from llama_index import ServiceContext
+            service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)
+            query_object = GPTVectorStoreIndexQuery(index.index_struct, service_context=service_context, similarity_top_k=5, vector_store=index._vector_store, docstore=index._docstore)
+            query_bundle = QueryBundle(inputs)
+            nodes = query_object.retrieve(query_bundle)
         reference_results = [n.node.text for n in nodes]
         reference_results = add_source_numbers(reference_results, use_source=False)
         display_reference = add_details(reference_results)

modules/config.py

@@ -1,3 +1,4 @@
+from collections import defaultdict
 from contextlib import contextmanager
 import os
 import logging
@@ -11,6 +12,8 @@ __all__ = [
     "dockerflag",
     "retrieve_proxy",
     "log_level",
+    "advance_docs",
+    "update_doc_config"
 ]
 
 # 添加一个统一的config文件，避免文件过多造成的疑惑（优先级最低）
@@ -109,5 +112,12 @@ def retrieve_proxy(proxy=None):
         os.environ["HTTP_PROXY"], os.environ["HTTPS_PROXY"] = old_var
 
 
-## 处理advance pdf
-advance_pdf = config.get("advance_pdf", {})
+## 处理advance docs
+advance_docs = defaultdict(lambda: defaultdict(dict))
+advance_docs.update(config.get("advance_docs", {}))
+def update_doc_config(two_column_pdf):
+    global advance_docs
+    if two_column_pdf:
+        advance_docs["pdf"]["two_column"] = True
+    
+    logging.info(f"更新后的文件参数为：{advance_docs}")

modules/llama_func.py

@@ -45,8 +45,9 @@ def get_documents(file_src):
             logging.debug("Loading PDF...")
             try:
                 from modules.pdf_func import parse_pdf
-                from modules.config import advance_pdf
-                text = parse_pdf(file.name, advance_pdf.get("two_column", False)).text
+                from modules.config import advance_docs
+                two_column = advance_docs["pdf"].get("two_column", False)
+                pdftext = parse_pdf(file.name, two_column).text
             except:
                 pdftext = ""
                 with open(file.name, 'rb') as pdfFileObj:
@@ -106,10 +107,11 @@ def construct_index(
         try:
             documents = get_documents(file_src)
             logging.info("构建索引中……")
-            service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper, chunk_size_limit=chunk_size_limit)
-            index = GPTSimpleVectorIndex.from_documents(
-                documents,  service_context=service_context
-            )
+            with retrieve_proxy():
+                service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper, chunk_size_limit=chunk_size_limit)
+                index = GPTSimpleVectorIndex.from_documents(
+                    documents,  service_context=service_context
+                )
             logging.debug("索引构建完成！")
             os.makedirs("./index", exist_ok=True)
             index.save_to_disk(f"./index/{index_name}.json")