From c60a7452bfcaa22588b9fe5bdae3b7fc94b8927b Mon Sep 17 00:00:00 2001 From: leike0813 Date: Wed, 20 Dec 2023 08:57:27 +0800 Subject: [PATCH] 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 | 4 + crazy_functional.py | 13 ++- crazy_functions/crazy_utils.py | 101 +++++++++++++++++++++- crazy_functions/批量总结PDF文档_初步.py | 11 +-- crazy_functions/批量翻译PDF文档_NOUGAT.py | 90 ++++++++++++++++++- 5 files changed, 206 insertions(+), 13 deletions(-) diff --git a/config.py b/config.py index 17dac34..dc8ef9b 100644 --- a/config.py +++ b/config.py @@ -217,6 +217,10 @@ GROBID_URLS = [ ] +# NOUGAT_API主机地址 +NOUGAT_URLS = ["http://localhost:8503"] # 此处填写NOUGAT_API的主机地址 + + # 是否允许通过自然语言描述修改本页的配置，该功能具有一定的危险性，默认关闭 ALLOW_RESET_CONFIG = False diff --git a/crazy_functional.py b/crazy_functional.py index 3275f79..d6a9dd9 100644 --- a/crazy_functional.py +++ b/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: diff --git a/crazy_functions/crazy_utils.py b/crazy_functions/crazy_utils.py index 731da1a..a583e25 100644 --- a/crazy_functions/crazy_utils.py +++ b/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=[]): diff --git a/crazy_functions/批量总结PDF文档_初步.py b/crazy_functions/批量总结PDF文档_初步.py index 5628fa1..52fd47b 100644 --- a/crazy_functions/批量总结PDF文档_初步.py +++ b/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 request_llms.bridge_all import model_info - enc = model_info["gpt-3.5-turbo"]['tokenizer'] - 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) + from crazy_functions.pdf_fns.breakdown_txt import breakdown_text_to_satisfy_token_limit + paper_fragments = breakdown_text_to_satisfy_token_limit(txt=file_content, limit=TOKEN_LIMIT_PER_FRAGMENT, llm_model=llm_kwargs['llm_model']) + 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']) # 为了更好的效果，我们剥离Introduction之后的部分（如果有） paper_meta = page_one_fragments[0].split('introduction')[0].split('Introduction')[0].split('INTRODUCTION')[0] diff --git a/crazy_functions/批量翻译PDF文档_NOUGAT.py b/crazy_functions/批量翻译PDF文档_NOUGAT.py index 97170d0..3b841a1 100644 --- a/crazy_functions/批量翻译PDF文档_NOUGAT.py +++ b/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) # 刷新界面 \ No newline at end of file