Merge pull request #270 from hellofinch/main

添加cli的自定义prompt支持。更新readme

README.md

@@ -174,7 +174,8 @@ In the following table, we list all advanced options for reference:
 | `-o`  | Output dir | `pdf2zh example.pdf -o output` |
 | `-f`, `-c` | [Exceptions](#exceptions) | `pdf2zh example.pdf -f "(MS.*)"` |
 | `--share` | [Get gradio public link] | `pdf2zh -i --share` |
-| `-a` | [add authorization and custom login page] | `pdf2zh -i -a users.txt [auth.html]` |
+| `--authorized` | [add authorization and custom login page] | `pdf2zh -i --authorized users.txt [auth.html]` |
+| `--prompt` | [custom llm prompt] | `pdf2zh --prompt [prompt.txt]` |
 
 <h3 id="partial">Full / partial document translation</h3>
 
@@ -254,7 +255,34 @@ Use `-t` to specify how many threads to use in translation:
 ```bash
 pdf2zh example.pdf -t 1
 ```
+<h3 id="prompt">custom prompt</h3>
+Use `--prompt` to specify which prompt to use in llm:
+```bash
+pdf2zh example.pdf -pr prompt.txt
+```
+
+
+example prompt.txt
+```
+[
+    {
+        "role": "system",
+        "content": "You are a professional,authentic machine translation engine.",
+    },
+    {
+        "role": "user",
+        "content": "Translate the following markdown source text to ${lang_out}. Keep the formula notation {{v*}} unchanged. Output translation directly without any additional text.\nSource Text: ${text}\nTranslated Text:",
+    },
+]
+```
+
 
+In custom prompt file, there are three variables can be used.
+|**variables**|**comment**|
+|-|-|
+|`lang_in`|input language|
+|`lang_out`|output language|
+|`text`|text need to be translated|
 <h2 id="todo">API</h2>
 
 ### Python

README_ja-JP.md

@@ -175,7 +175,8 @@ Python環境を事前にインストールする必要はありません
 | `-o`  | 出力ディレクトリ | `pdf2zh example.pdf -o output` |
 | `-f`, `-c` | [例外](#exceptions) | `pdf2zh example.pdf -f "(MS.*)"` |
 | `--share` | [gradio公開リンクを取得] | `pdf2zh -i --share` |
-| `-a` | [ウェブ認証とカスタム認証ページの追加] | `pdf2zh -i -a users.txt [auth.html]` |
+| `--authorized` | [ウェブ認証とカスタム認証ページの追加] | `pdf2zh -i --authorized users.txt [auth.html]` |
+| `--prompt` | [カスタムビッグモデルのプロンプトを使用する] | `pdf2zh --prompt [prompt.txt]` |
 
 <h3 id="partial">全文または部分的なドキュメント翻訳</h3>
 
@@ -256,6 +257,35 @@ pdf2zh example.pdf -f "(CM[^R]|(MS|XY|MT|BL|RM|EU|LA|RS)[A-Z]|LINE|LCIRCLE|TeX-|
 pdf2zh example.pdf -t 1
 ```
 
+<h3 id="prompt">custom prompt</h3>
+(need Japenese translation)
+Use `--prompt` to specify which prompt to use in llm:
+```bash
+pdf2zh example.pdf -pr prompt.txt
+```
+
+
+example prompt.txt
+```
+[
+    {
+        "role": "system",
+        "content": "You are a professional,authentic machine translation engine.",
+    },
+    {
+        "role": "user",
+        "content": "Translate the following markdown source text to ${lang_out}. Keep the formula notation {{v*}} unchanged. Output translation directly without any additional text.\nSource Text: ${text}\nTranslated Text:",
+    },
+]
+```
+
+
+In custom prompt file, there are three variables can be used.
+|**variables**|**comment**|
+|-|-|
+|`lang_in`|input language|
+|`lang_out`|output language|
+|`text`|text need to be translated|
 <h2 id="todo">API</h2>
 
 ### Python

README_zh-CN.md

@@ -174,7 +174,8 @@ USE_MODELSCOPE=1 pdf2zh
 | `-o`  | 输出目录 | `pdf2zh example.pdf -o output` |
 | `-f`, `-c` | [例外规则](#exceptions) | `pdf2zh example.pdf -f "(MS.*)"` |
 | `--share` | [获取 gradio 公开链接] | `pdf2zh -i --share` |
-| `-a` | [添加网页认证和自定义认证页] | `pdf2zh -i -a users.txt [auth.html]` |
+| `--authorized` | [添加网页认证和自定义认证页] | `pdf2zh -i --authorized users.txt [auth.html]` |
+| `--prompt` | [使用自定义的大模型prompt] | `pdf2zh --prompt [prompt.txt]` |
 
 <h3 id="partial">全文或部分文档翻译</h3>
 
@@ -254,6 +255,37 @@ pdf2zh example.pdf -f "(CM[^R]|(MS|XY|MT|BL|RM|EU|LA|RS)[A-Z]|LINE|LCIRCLE|TeX-|
 ```bash
 pdf2zh example.pdf -t 1
 ```
+<h3 id="prompt">自定义大模型prompt</h3>
+
+使用 `--prompt` 指定使用大模型翻译时使用的 Prompt 文件。
+
+```bash
+pdf2zh example.pdf -pr prompt.txt
+```
+
+
+示例 `prompt.txt` 文件
+
+```
+[
+    {
+        "role": "system",
+        "content": "You are a professional,authentic machine translation engine.",
+    },
+    {
+        "role": "user",
+        "content": "Translate the following markdown source text to ${lang_out}. Keep the formula notation {{v*}} unchanged. Output translation directly without any additional text.\nSource Text: ${text}\nTranslated Text:",
+    },
+]
+```
+
+
+自定义 Prompt 文件中，可以使用三个内置变量用来传递参数。
+|**变量名**|**说明**|
+|-|-|
+|`lang_in`|输入的语言|
+|`lang_out`|输出的语言|
+|`text`|需要翻译的文本|
 
 <h2 id="todo">API</h2>
 

pdf2zh/converter.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, List
 
 from pdfminer.pdfinterp import PDFGraphicState, PDFResourceManager
 from pdfminer.pdffont import PDFCIDFont
@@ -136,6 +136,7 @@ class TranslateConverter(PDFConverterEx):
         resfont: str = "",
         noto: Font = None,
         envs: Dict = None,
+        prompt: List = None,
     ) -> None:
         super().__init__(rsrcmgr)
         self.vfont = vfont
@@ -151,7 +152,7 @@ class TranslateConverter(PDFConverterEx):
         for translator in [GoogleTranslator, BingTranslator, DeepLTranslator, DeepLXTranslator, OllamaTranslator, AzureOpenAITranslator,
                            OpenAITranslator, ZhipuTranslator, ModelScopeTranslator, SiliconTranslator, GeminiTranslator, AzureTranslator, TencentTranslator, DifyTranslator, AnythingLLMTranslator]:
             if service_name == translator.name:
-                self.translator = translator(lang_in, lang_out, service_model, envs=envs)
+                self.translator = translator(lang_in, lang_out, service_model, envs=envs,prompt=prompt)
         if not self.translator:
             raise ValueError("Unsupported translation service")
 

pdf2zh/high_level.py

@@ -103,6 +103,7 @@ def translate_patch(
         resfont,
         noto,
         kwarg.get("envs", {}),
+        kwarg.get("prompt", []),
     )
 
     assert device is not None
@@ -226,7 +227,7 @@ def translate_stream(
 
     fp = io.BytesIO()
     doc_zh.save(fp)
-    obj_patch: dict = translate_patch(fp, envs=kwarg["envs"], **locals())
+    obj_patch: dict = translate_patch(fp, prompt=kwarg["prompt"], **locals())
 
     for obj_id, ops_new in obj_patch.items():
         # ops_old=doc_en.xref_stream(obj_id)
@@ -292,7 +293,9 @@ def translate(
 
         doc_raw = open(file, "rb")
         s_raw = doc_raw.read()
-        s_mono, s_dual = translate_stream(s_raw, envs=kwarg.get('envs'), **locals())
+        s_mono, s_dual = translate_stream(
+            s_raw, envs=kwarg.get("envs"), prompt=kwarg["prompt"], **locals()
+        )
         file_mono = Path(output) / f"{filename}-mono.pdf"
         file_dual = Path(output) / f"{filename}-dual.pdf"
         doc_mono = open(file_mono, "wb")

pdf2zh/pdf2zh.py

@@ -11,6 +11,7 @@ import logging
 from typing import List, Optional
 from pdf2zh import __version__, log
 from pdf2zh.high_level import translate
+from string import Template
 
 
 def create_parser() -> argparse.ArgumentParser:
@@ -117,12 +118,15 @@ def create_parser() -> argparse.ArgumentParser:
     )
     parse_params.add_argument(
         "--authorized",
-        "-a",
         type=str,
         nargs="+",
-        default=["./users.txt", "./auth.html"],
         help="user name and password.",
     )
+    parse_params.add_argument(
+        "--prompt",
+        type=str,
+        help="user custom prompt.",
+    )
 
     return parser
 
@@ -169,6 +173,14 @@ def main(args: Optional[List[str]] = None) -> int:
         celery_app.start(argv=sys.argv[2:])
         return 0
 
+    if parsed_args.prompt:
+        try:
+            with open(parsed_args.prompt, "r", encoding="utf-8") as file:
+                content = file.read()
+            parsed_args.prompt = Template(content)
+        except Exception as e:
+            raise ValueError("prompt error.")
+
     translate(**vars(parsed_args))
     return 0
 

pdf2zh/pdfinterp.py

@@ -236,9 +236,9 @@ class PDFPageInterpreterEx(PDFPageInterpreter):
                     pos_inv = -np.mat(ctm[4:]) * ctm_inv
                 a, b, c, d = ctm_inv.reshape(4).tolist()
                 e, f = pos_inv.tolist()[0]
-                self.obj_patch[self.xobjmap[xobjid].objid] = (
-                    f"q {ops_base}Q {a} {b} {c} {d} {e} {f} cm {ops_new}"
-                )
+                self.obj_patch[
+                    self.xobjmap[xobjid].objid
+                ] = f"q {ops_base}Q {a} {b} {c} {d} {e} {f} cm {ops_new}"
             except Exception:
                 pass
         elif subtype is LITERAL_IMAGE and "Width" in xobj and "Height" in xobj:
@@ -269,9 +269,9 @@ class PDFPageInterpreterEx(PDFPageInterpreter):
         self.device.fontmap = self.fontmap
         ops_new = self.device.end_page(page)
         # 上面渲染的时候会根据 cropbox 减掉页面偏移得到真实坐标，这里输出的时候需要用 cm 把页面偏移加回来
-        self.obj_patch[page.page_xref] = (
-            f"q {ops_base}Q 1 0 0 1 {x0} {y0} cm {ops_new}"  # ops_base 里可能有图，需要让 ops_new 里的文字覆盖在上面，使用 q/Q 重置位置矩阵
-        )
+        self.obj_patch[
+            page.page_xref
+        ] = f"q {ops_base}Q 1 0 0 1 {x0} {y0} cm {ops_new}"  # ops_base 里可能有图，需要让 ops_new 里的文字覆盖在上面，使用 q/Q 重置位置矩阵
         for obj in page.contents:
             self.obj_patch[obj.objid] = ""
 

pdf2zh/translator.py

@@ -49,17 +49,25 @@ class BaseTranslator:
     def translate(self, text):
         pass
 
-    def prompt(self, text):
-        return [
-            {
-                "role": "system",
-                "content": "You are a professional,authentic machine translation engine.",
-            },
-            {
-                "role": "user",
-                "content": f"Translate the following markdown source text to {self.lang_out}. Keep the formula notation {{v*}} unchanged. Output translation directly without any additional text.\nSource Text: {text}\nTranslated Text:",  # noqa: E501
-            },
-        ]
+    def prompt(self, text, prompt):
+        if prompt:
+            context = {
+                "lang_in": self.lang_in,
+                "lang_out": self.lang_out,
+                "text": text,
+            }
+            return eval(prompt.safe_substitute(context))
+        else:
+            return [
+                {
+                    "role": "system",
+                    "content": "You are a professional,authentic machine translation engine.",
+                },
+                {
+                    "role": "user",
+                    "content": f"Translate the following markdown source text to {self.lang_out}. Keep the formula notation {{v*}} unchanged. Output translation directly without any additional text.\nSource Text: {text}\nTranslated Text:",  # noqa: E501
+                },
+            ]
 
     def __str__(self):
         return f"{self.name} {self.lang_in} {self.lang_out} {self.model}"
@@ -145,7 +153,7 @@ class DeepLTranslator(BaseTranslator):
     }
     lang_map = {"zh": "zh-Hans"}
 
-    def __init__(self, lang_in, lang_out, model, envs=None):
+    def __init__(self, lang_in, lang_out, model, envs=None, **kwargs):
         self.set_envs(envs)
         super().__init__(lang_in, lang_out, model)
         auth_key = self.envs["DEEPL_AUTH_KEY"]
@@ -166,7 +174,7 @@ class DeepLXTranslator(BaseTranslator):
     }
     lang_map = {"zh": "zh-Hans"}
 
-    def __init__(self, lang_in, lang_out, model, envs=None):
+    def __init__(self, lang_in, lang_out, model, envs=None, **kwargs):
         self.set_envs(envs)
         super().__init__(lang_in, lang_out, model)
         self.endpoint = self.envs["DEEPLX_ENDPOINT"]
@@ -193,19 +201,23 @@ class OllamaTranslator(BaseTranslator):
         "OLLAMA_MODEL": "gemma2",
     }
 
-    def __init__(self, lang_in, lang_out, model, envs=None):
+    def __init__(self, lang_in, lang_out, model, envs=None, prompt=None):
         self.set_envs(envs)
         if not model:
             model = self.envs["OLLAMA_MODEL"]
         super().__init__(lang_in, lang_out, model)
         self.options = {"temperature": 0}  # 随机采样可能会打断公式标记
         self.client = ollama.Client()
+        self.prompttext = prompt
 
     def translate(self, text):
+        print(len(self.prompt(text, self.prompttext)))
+        print(self.prompt(text, self.prompttext)[0])
+        print(self.prompt(text, self.prompttext)[1])
         response = self.client.chat(
             model=self.model,
             options=self.options,
-            messages=self.prompt(text),
+            messages=self.prompt(text, self.prompttext),
         )
         return response["message"]["content"].strip()
 
@@ -220,7 +232,14 @@ class OpenAITranslator(BaseTranslator):
     }
 
     def __init__(
-        self, lang_in, lang_out, model, base_url=None, api_key=None, envs=None
+        self,
+        lang_in,
+        lang_out,
+        model,
+        base_url=None,
+        api_key=None,
+        envs=None,
+        prompt=None,
     ):
         self.set_envs(envs)
         if not model:
@@ -228,12 +247,13 @@ class OpenAITranslator(BaseTranslator):
         super().__init__(lang_in, lang_out, model)
         self.options = {"temperature": 0}  # 随机采样可能会打断公式标记
         self.client = openai.OpenAI(base_url=base_url, api_key=api_key)
+        self.prompttext = prompt
 
     def translate(self, text) -> str:
         response = self.client.chat.completions.create(
             model=self.model,
             **self.options,
-            messages=self.prompt(text),
+            messages=self.prompt(text, self.prompttext),
         )
         return response.choices[0].message.content.strip()
 
@@ -247,7 +267,14 @@ class AzureOpenAITranslator(BaseTranslator):
     }
 
     def __init__(
-        self, lang_in, lang_out, model, base_url=None, api_key=None, envs=None
+        self,
+        lang_in,
+        lang_out,
+        model,
+        base_url=None,
+        api_key=None,
+        envs=None,
+        prompt=None,
     ):
         self.set_envs(envs)
         base_url = self.envs["AZURE_OPENAI_BASE_URL"]
@@ -261,12 +288,13 @@ class AzureOpenAITranslator(BaseTranslator):
             api_version="2024-06-01",
             api_key=api_key,
         )
+        self.prompttext = prompt
 
     def translate(self, text) -> str:
         response = self.client.chat.completions.create(
             model=self.model,
             **self.options,
-            messages=self.prompt(text),
+            messages=self.prompt(text, self.prompttext),
         )
         return response.choices[0].message.content.strip()
 
@@ -280,7 +308,14 @@ class ModelScopeTranslator(OpenAITranslator):
     }
 
     def __init__(
-        self, lang_in, lang_out, model, base_url=None, api_key=None, envs=None
+        self,
+        lang_in,
+        lang_out,
+        model,
+        base_url=None,
+        api_key=None,
+        envs=None,
+        prompt=None,
     ):
         self.set_envs(envs)
         base_url = "https://api-inference.modelscope.cn/v1"
@@ -288,6 +323,7 @@ class ModelScopeTranslator(OpenAITranslator):
         if not model:
             model = self.envs["MODELSCOPE_MODEL"]
         super().__init__(lang_in, lang_out, model, base_url=base_url, api_key=api_key)
+        self.prompttext = prompt
 
 
 class ZhipuTranslator(OpenAITranslator):
@@ -298,20 +334,21 @@ class ZhipuTranslator(OpenAITranslator):
         "ZHIPU_MODEL": "glm-4-flash",
     }
 
-    def __init__(self, lang_in, lang_out, model, envs=None):
+    def __init__(self, lang_in, lang_out, model, envs=None, prompt=None):
         self.set_envs(envs)
         base_url = "https://open.bigmodel.cn/api/paas/v4"
         api_key = self.envs["ZHIPU_API_KEY"]
         if not model:
             model = self.envs["ZHIPU_MODEL"]
         super().__init__(lang_in, lang_out, model, base_url=base_url, api_key=api_key)
+        self.prompttext = prompt
 
     def translate(self, text) -> str:
         try:
             response = self.client.chat.completions.create(
                 model=self.model,
                 **self.options,
-                messages=self.prompt(text),
+                messages=self.prompt(text, self.prompttext),
             )
         except openai.BadRequestError as e:
             if (
@@ -331,13 +368,14 @@ class SiliconTranslator(OpenAITranslator):
         "SILICON_MODEL": "Qwen/Qwen2.5-7B-Instruct",
     }
 
-    def __init__(self, lang_in, lang_out, model, envs=None):
+    def __init__(self, lang_in, lang_out, model, envs=None, prompt=None):
         self.set_envs(envs)
         base_url = "https://api.siliconflow.cn/v1"
         api_key = self.envs["SILICON_API_KEY"]
         if not model:
             model = self.envs["SILICON_MODEL"]
         super().__init__(lang_in, lang_out, model, base_url=base_url, api_key=api_key)
+        self.prompttext = prompt
 
 
 class GeminiTranslator(OpenAITranslator):
@@ -348,13 +386,14 @@ class GeminiTranslator(OpenAITranslator):
         "GEMINI_MODEL": "gemini-1.5-flash",
     }
 
-    def __init__(self, lang_in, lang_out, model, envs=None):
+    def __init__(self, lang_in, lang_out, model, envs=None, prompt=None):
         self.set_envs(envs)
         base_url = "https://generativelanguage.googleapis.com/v1beta/openai/"
         api_key = self.envs["GEMINI_API_KEY"]
         if not model:
             model = self.envs["GEMINI_MODEL"]
         super().__init__(lang_in, lang_out, model, base_url=base_url, api_key=api_key)
+        self.prompttext = prompt
 
 
 class AzureTranslator(BaseTranslator):
@@ -366,7 +405,7 @@ class AzureTranslator(BaseTranslator):
     }
     lang_map = {"zh": "zh-Hans"}
 
-    def __init__(self, lang_in, lang_out, model, envs=None):
+    def __init__(self, lang_in, lang_out, model, envs=None, **kwargs):
         self.set_envs(envs)
         super().__init__(lang_in, lang_out, model)
         endpoint = self.envs["AZURE_ENDPOINT"]
@@ -397,7 +436,7 @@ class TencentTranslator(BaseTranslator):
         "TENCENTCLOUD_SECRET_KEY": None,
     }
 
-    def __init__(self, lang_in, lang_out, model, envs=None):
+    def __init__(self, lang_in, lang_out, model, envs=None, **kwargs):
         self.set_envs(envs)
         super().__init__(lang_in, lang_out, model)
         cred = credential.DefaultCredentialProvider().get_credential()
@@ -420,7 +459,7 @@ class AnythingLLMTranslator(BaseTranslator):
         "AnythingLLM_APIKEY": None,
     }
 
-    def __init__(self, lang_out, lang_in, model, envs=None):
+    def __init__(self, lang_out, lang_in, model, envs=None, prompt=None):
         self.set_envs(envs)
         super().__init__(lang_out, lang_in, model)
         self.api_url = self.envs["AnythingLLM_URL"]
@@ -430,9 +469,10 @@ class AnythingLLMTranslator(BaseTranslator):
             "Authorization": f"Bearer {self.api_key}",
             "Content-Type": "application/json",
         }
+        self.prompttext = prompt
 
     def translate(self, text):
-        messages = self.prompt(text)
+        messages = self.prompt(text, self.prompttext)
         payload = {
             "message": messages,
             "mode": "chat",
@@ -456,7 +496,7 @@ class DifyTranslator(BaseTranslator):
         "DIFY_API_KEY": None,  # 替换为实际 API 密钥
     }
 
-    def __init__(self, lang_out, lang_in, model, envs=None):
+    def __init__(self, lang_out, lang_in, model, envs=None, **kwargs):
         self.set_envs(envs)
         super().__init__(lang_out, lang_in, model)
         self.api_url = self.envs["DIFY_API_URL"]