Improve NOUGAT pdf plugin

Add an API version of NOUGAT plugin
Add advanced argument support to NOUGAT plugin

Adapt new text breakdown function

bugfix

config.py

@@ -217,6 +217,10 @@ GROBID_URLS = [
 ]
 
 
+# NOUGAT_API主机地址
+NOUGAT_URLS = ["http://localhost:8503"] # 此处填写NOUGAT_API的主机地址
+
+
 # 是否允许通过自然语言描述修改本页的配置，该功能具有一定的危险性，默认关闭
 ALLOW_RESET_CONFIG = False
 

crazy_functional.py

@@ -549,13 +549,24 @@ def get_crazy_functions():
         print('Load function plugin failed')
 
     try:
-        from crazy_functions.批量翻译PDF文档_NOUGAT import 批量翻译PDF文档
+        from crazy_functions.批量翻译PDF文档_NOUGAT import 批量翻译PDF文档, 批量翻译PDF文档_API
         function_plugins.update({
             "精准翻译PDF文档（NOUGAT）": {
                 "Group": "学术",
                 "Color": "stop",
                 "AsButton": False,
+                "AdvancedArgs": True,  # 调用时，唤起高级参数输入区（默认False）
+                "ArgsReminder": "在这里输入自定义参数, 支持的参数有: --batchsize BATCHSIZE, --model MODEL_TAG, --recompute, --full-precision, --no-markdown --no-skipping, --pages PAGES/-p PAGES", # 高级参数输入区的显示提示
                 "Function": HotReload(批量翻译PDF文档)
+            },
+            "精准翻译PDF文档（NOUGAT_API）": {
+                "Group": "学术",
+                "Color": "stop",
+                "AsButton": False,
+                "AdvancedArgs": True,  # 调用时，唤起高级参数输入区（默认False）
+                "ArgsReminder": "在这里输入自定义参数, 支持的参数有: --batchsize BATCHSIZE, --recompute, --no-markdown --no-skipping, --pages PAGES/-p PAGES (官方版本的API仅支持--pages参数)",
+                # 高级参数输入区的显示提示
+                "Function": HotReload(批量翻译PDF文档_API)
             }
         })
     except:

crazy_functions/crazy_utils.py

@@ -545,7 +545,20 @@ def get_files_from_everything(txt, type): # type='.md'
 @Singleton
 class nougat_interface():
     def __init__(self):
+        def model_check(model_tag):
+            if model_tag in ['0.1.0-small', '0.1.0-base']: return model_tag
+            return '0.1.0-small'
+
+        import argparse
         self.threadLock = threading.Lock()
+        self.arg_parser = argparse.ArgumentParser()
+        self.arg_parser.add_argument('--batchsize', type=int)
+        self.arg_parser.add_argument('--model', type=model_check)
+        self.arg_parser.add_argument('--recompute', action='store_true')
+        self.arg_parser.add_argument('--full-precision', action='store_true')
+        self.arg_parser.add_argument('--no-markdown', action='store_true')
+        self.arg_parser.add_argument('--no-skipping', action='store_true')
+        self.arg_parser.add_argument('--pages', type=str)
 
     def nougat_with_timeout(self, command, cwd, timeout=3600):
         import subprocess
@@ -563,7 +576,7 @@ class nougat_interface():
         return True
 
 
-    def NOUGAT_parse_pdf(self, fp, chatbot, history):
+    def NOUGAT_parse_pdf(self, fp, chatbot, history, advanced_cfg=''):
         from toolbox import update_ui_lastest_msg
 
         yield from update_ui_lastest_msg("正在解析论文, 请稍候。进度：正在排队, 等待线程锁...", 
@@ -576,7 +589,10 @@ class nougat_interface():
 
         yield from update_ui_lastest_msg("正在解析论文, 请稍候。进度：正在加载NOUGAT... （提示：首次运行需要花费较长时间下载NOUGAT参数）", 
                                          chatbot=chatbot, history=history, delay=0)
-        self.nougat_with_timeout(f'nougat --out "{os.path.abspath(dst)}" "{os.path.abspath(fp)}"', os.getcwd(), timeout=3600)
+        self.nougat_with_timeout(
+            f'nougat --out "{os.path.abspath(dst)}" "{os.path.abspath(fp)}" {self.parse_argument(advanced_cfg)}',
+            os.getcwd(), timeout=3600
+        )
         res = glob.glob(os.path.join(dst,'*.mmd'))
         if len(res) == 0:
             self.threadLock.release()
@@ -585,6 +601,87 @@ class nougat_interface():
         return res[0]
 
 
+    def NOUGAT_API_parse_pdf(self, fp, chatbot, history, nougat_url, advanced_cfg=''):
+        from toolbox import update_ui_lastest_msg
+
+        yield from update_ui_lastest_msg("正在解析论文, 请稍候。",
+                                         chatbot=chatbot, history=history, delay=0)
+
+        import requests
+        from toolbox import get_log_folder, gen_time_str
+        dst = os.path.join(get_log_folder(plugin_name='nougat'), gen_time_str())
+        os.makedirs(dst)
+
+        ret = requests.post(
+            f'{nougat_url}/predict{self.parse_api_argument(advanced_cfg)}',
+            files={"file": open(fp, "rb")}
+        )
+        if ret.status_code != 200:
+            raise RuntimeError("Nougat解析论文失败。")
+
+        with open(os.path.join(dst, '*.mmd'), 'w') as f:
+            f.write(ret.json())
+        return os.path.join(dst, '*.mmd')
+
+
+    def parse_argument(self, argument_string):
+        args, _ = self.arg_parser.parse_known_args(argument_string.split())
+        reduce_args = []
+        for k, v in args.__dict__.items():
+            if (v is not None) and (v is not False):
+                reduce_args.append('--' + k.replace('_', '-'))
+            if not isinstance(v, bool) and v is not None:
+                reduce_args.append(str(v))
+
+        return ' '.join(reduce_args)
+
+
+    def parse_api_argument(self, argument_string):
+        def parse_pages(pages_string):
+            if pages_string.count(',') > 0:
+                pages_list = pages_string.split(',')
+                page_start = pages_list[0].split('-')[0] if '-' in pages_list[0] else pages_list[0]
+                page_end = pages_list[-1].split('-')[-1] if '-' in pages_list[-1] else pages_list[-1]
+            else:
+                if '-' in pages_string:
+                    page_start = pages_string.split('-')[0]
+                    page_end = pages_string.split('-')[-1]
+                else:
+                    page_start = page_end = int(pages_string)
+
+            return page_start, page_end
+
+        args, _ = self.arg_parser.parse_known_args(argument_string.split())
+        reduce_args = []
+        for k, v in args.__dict__.items():
+            arg_pair = ''
+            if (v is not None) and (v is not False):
+                if k == 'pages':
+                    page_start, page_end = parse_pages(v)
+                    arg_pair = f'start={page_start}&stop={page_end}'
+                elif k not in ['model', 'full_precision']:
+                    arg_pair = f'{k}={int(v)}'
+            if arg_pair:
+                reduce_args.append(arg_pair)
+
+        return '?' + '&'.join(reduce_args)
+
+    @staticmethod
+    def get_avail_nougat_url():
+        import random
+        import requests
+        NOUGAT_URLS = get_conf('NOUGAT_URLS')
+        if len(NOUGAT_URLS) == 0: return None
+        try:
+            _nougat_url = random.choice(NOUGAT_URLS)  # 随机负载均衡
+            if _nougat_url.endswith('/'): _nougat_url = _nougat_url.rstrip('/')
+            ret = requests.get(_nougat_url + '/')
+            if ret.status_code == 200:
+                return _nougat_url
+            else:
+                return None
+        except:
+            return None
 
 
 def try_install_deps(deps, reload_m=[]):

crazy_functions/批量总结PDF文档_初步.py

@@ -29,14 +29,9 @@ def 解析PDF(file_manifest, project_folder, llm_kwargs, plugin_kwargs, chatbot,
 
             TOKEN_LIMIT_PER_FRAGMENT = 2500
 
-            from .crazy_utils import breakdown_txt_to_satisfy_token_limit_for_pdf
+            from crazy_functions.pdf_fns.breakdown_txt import breakdown_text_to_satisfy_token_limit
-            from request_llms.bridge_all import model_info
+            paper_fragments = breakdown_text_to_satisfy_token_limit(txt=file_content, limit=TOKEN_LIMIT_PER_FRAGMENT, llm_model=llm_kwargs['llm_model'])
-            enc = model_info["gpt-3.5-turbo"]['tokenizer']
+            page_one_fragments = breakdown_text_to_satisfy_token_limit(txt=str(page_one), limit=TOKEN_LIMIT_PER_FRAGMENT // 4, llm_model=llm_kwargs['llm_model'])
-            def get_token_num(txt): return len(enc.encode(txt, disallowed_special=()))
-            paper_fragments = breakdown_txt_to_satisfy_token_limit_for_pdf(
-                txt=file_content,  get_token_fn=get_token_num, limit=TOKEN_LIMIT_PER_FRAGMENT)
-            page_one_fragments = breakdown_txt_to_satisfy_token_limit_for_pdf(
-                txt=str(page_one), get_token_fn=get_token_num, limit=TOKEN_LIMIT_PER_FRAGMENT//4)
             # 为了更好的效果，我们剥离Introduction之后的部分（如果有）
             paper_meta = page_one_fragments[0].split('introduction')[0].split('Introduction')[0].split('INTRODUCTION')[0]
 

crazy_functions/批量翻译PDF文档_NOUGAT.py

@@ -54,7 +54,7 @@ def 批量翻译PDF文档(txt, llm_kwargs, plugin_kwargs, chatbot, history, syst
     # 基本信息：功能、贡献者
     chatbot.append([
         "函数插件功能？",
-        "批量翻译PDF文档。函数插件贡献者: Binary-Husky"])
+        "批量翻译PDF文档。函数插件贡献者: Binary-Husky，Joshua Reed"])
     yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
 
     # 清空历史，以免输入溢出
@@ -104,11 +104,13 @@ def 解析PDF_基于NOUGAT(file_manifest, project_folder, llm_kwargs, plugin_kwa
     DST_LANG = "中文"
     from crazy_functions.crazy_utils import nougat_interface
     from crazy_functions.pdf_fns.report_gen_html import construct_html
+    if ("advanced_arg" in plugin_kwargs) and (plugin_kwargs["advanced_arg"] == ""): plugin_kwargs.pop("advanced_arg")
+    advanced_cfg = plugin_kwargs.get("advanced_arg", '')
     nougat_handle = nougat_interface()
     for index, fp in enumerate(file_manifest):
         if fp.endswith('pdf'):
             chatbot.append(["当前进度：", f"正在解析论文，请稍候。（第一次运行时，需要花费较长时间下载NOUGAT参数）"]); yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
-            fpp = yield from nougat_handle.NOUGAT_parse_pdf(fp, chatbot, history)
+            fpp = yield from nougat_handle.NOUGAT_parse_pdf(fp, chatbot, history, advanced_cfg=advanced_cfg)
             promote_file_to_downloadzone(fpp, rename_file=os.path.basename(fpp)+'.nougat.mmd', chatbot=chatbot)
         else:
             chatbot.append(["当前论文无需解析：", fp]); yield from update_ui(      chatbot=chatbot, history=history)
@@ -123,3 +125,87 @@ def 解析PDF_基于NOUGAT(file_manifest, project_folder, llm_kwargs, plugin_kwa
     yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
 
 
+@CatchException
+def 批量翻译PDF文档_API(txt, llm_kwargs, plugin_kwargs, chatbot, history, system_prompt, web_port):
+
+    disable_auto_promotion(chatbot)
+    # 基本信息：功能、贡献者
+    chatbot.append([
+        "函数插件功能？",
+        "使用NOUGAT_API批量翻译PDF文档。函数插件贡献者: Binary-Husky，Joshua Reed。\n"
+        + "官方版本API仅支持页码范围选择，若要支持更多参数，请移步https://github.com/leike0813/nougat",
+    ])
+    yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
+
+    # 清空历史，以免输入溢出
+    history = []
+
+    from .crazy_utils import get_files_from_everything
+    success, file_manifest, project_folder = get_files_from_everything(txt, type='.pdf')
+    if len(file_manifest) > 0:
+        # 尝试导入依赖，如果缺少依赖，则给出安装建议
+        try:
+            import tiktoken
+        except:
+            report_exception(chatbot, history,
+                             a=f"解析项目: {txt}",
+                             b=f"导入软件依赖失败。使用该模块需要额外依赖，安装方法```pip install --upgrade tiktoken```。")
+            yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
+            return
+    success_mmd, file_manifest_mmd, _ = get_files_from_everything(txt, type='.mmd')
+    success = success or success_mmd
+    file_manifest += file_manifest_mmd
+    chatbot.append(["文件列表：", ", ".join([e.split('/')[-1] for e in file_manifest])]);
+    yield from update_ui(chatbot=chatbot, history=history)
+    # 检测输入参数，如没有给定输入参数，直接退出
+    if not success:
+        if txt == "": txt = '空空如也的输入栏'
+
+    # 如果没找到任何文件
+    if len(file_manifest) == 0:
+        report_exception(chatbot, history,
+                         a=f"解析项目: {txt}", b=f"找不到任何.pdf拓展名的文件: {txt}")
+        yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
+        return
+
+    # 开始正式执行任务
+    yield from 解析PDF_基于NOUGAT_API(file_manifest, project_folder, llm_kwargs, plugin_kwargs, chatbot, history, system_prompt)
+
+
+def 解析PDF_基于NOUGAT_API(file_manifest, project_folder, llm_kwargs, plugin_kwargs, chatbot, history, system_prompt):
+    import copy
+    import tiktoken
+    TOKEN_LIMIT_PER_FRAGMENT = 1024
+    generated_conclusion_files = []
+    generated_html_files = []
+    DST_LANG = "中文"
+    from crazy_functions.crazy_utils import nougat_interface
+    from crazy_functions.pdf_fns.report_gen_html import construct_html
+    if ("advanced_arg" in plugin_kwargs) and (plugin_kwargs["advanced_arg"] == ""): plugin_kwargs.pop("advanced_arg")
+    advanced_cfg = plugin_kwargs.get("advanced_arg", '')
+    nougat_handle = nougat_interface()
+    chatbot.append(["当前进度：", f"正在检查NOUGAT服务可用性..."]);
+    yield from update_ui(chatbot=chatbot, history=history)  # 刷新界面
+    nougat_url = nougat_handle.get_avail_nougat_url()
+    if nougat_url is None:
+        report_exception(chatbot, history,
+                         a=f"检查结果：", b="NOUGAT服务不可用，请检查config中的NOUGAT_URL")
+        yield from update_ui(chatbot=chatbot, history=history)  # 刷新界面
+        return
+
+    for index, fp in enumerate(file_manifest):
+        if fp.endswith('pdf'):
+            chatbot.append(["当前进度：", f"正在解析论文，请稍候。"]); yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
+            fpp = yield from nougat_handle.NOUGAT_API_parse_pdf(fp, chatbot, history, nougat_url, advanced_cfg=advanced_cfg)
+            promote_file_to_downloadzone(fpp, rename_file=os.path.basename(fpp)+'.nougat.mmd', chatbot=chatbot)
+        else:
+            chatbot.append(["当前论文无需解析：", fp]); yield from update_ui(chatbot=chatbot, history=history)
+            fpp = fp
+        with open(fpp, 'r', encoding='utf8') as f:
+            article_content = f.readlines()
+        article_dict = markdown_to_dict(article_content)
+        logging.info(article_dict)
+        yield from translate_pdf(article_dict, llm_kwargs, chatbot, fp, generated_conclusion_files, TOKEN_LIMIT_PER_FRAGMENT, DST_LANG)
+
+    chatbot.append(("给出输出文件清单", str(generated_conclusion_files + generated_html_files)))
+    yield from update_ui(chatbot=chatbot, history=history) # 刷新界面