hellofinch 1 سال پیش
والد
کامیت
21a68506b1
5فایلهای تغییر یافته به همراه64 افزوده شده و 39 حذف شده
  1. 1 1
      pdf2zh/converter.py
  2. 3 1
      pdf2zh/high_level.py
  3. 3 3
      pdf2zh/pdf2zh.py
  4. 6 6
      pdf2zh/pdfinterp.py
  5. 51 28
      pdf2zh/translator.py

+ 1 - 1
pdf2zh/converter.py

@@ -1,4 +1,4 @@
-from typing import Dict,List
+from typing import Dict, List
 
 from pdfminer.pdfinterp import PDFGraphicState, PDFResourceManager
 from pdfminer.pdffont import PDFCIDFont

+ 3 - 1
pdf2zh/high_level.py

@@ -293,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'), prompt=kwarg["prompt"], **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")

+ 3 - 3
pdf2zh/pdf2zh.py

@@ -177,9 +177,9 @@ def main(args: Optional[List[str]] = None) -> int:
 
     if parsed_args.prompt:
         try:
-            with open(parsed_args.prompt,'r',encoding='utf-8') as file:
-                content=file.read()
-            parsed_args.prompt=Template(content)
+            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.")
 

+ 6 - 6
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] = ""
 

+ 51 - 28
pdf2zh/translator.py

@@ -17,6 +17,7 @@ from tencentcloud.tmt.v20180321.models import TextTranslateResponse
 
 import json
 
+
 def remove_control_characters(s):
     return "".join(ch for ch in s if unicodedata.category(ch)[0] != "C")
 
@@ -50,10 +51,10 @@ class BaseTranslator:
 
     def prompt(self, text, prompt):
         if prompt:
-            context={
-                "lang_in":self.lang_in,
-                "lang_out":self.lang_out,
-                "text":text,
+            context = {
+                "lang_in": self.lang_in,
+                "lang_out": self.lang_out,
+                "text": text,
             }
             return eval(prompt.safe_substitute(context))
         else:
@@ -200,23 +201,23 @@ class OllamaTranslator(BaseTranslator):
         "OLLAMA_MODEL": "gemma2",
     }
 
-    def __init__(self, lang_in, lang_out, model, envs=None,prompt=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
+        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])
+        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,self.prompttext),
+            messages=self.prompt(text, self.prompttext),
         )
         return response["message"]["content"].strip()
 
@@ -231,7 +232,14 @@ class OpenAITranslator(BaseTranslator):
     }
 
     def __init__(
-        self, lang_in, lang_out, model, base_url=None, api_key=None, envs=None,prompt=None
+        self,
+        lang_in,
+        lang_out,
+        model,
+        base_url=None,
+        api_key=None,
+        envs=None,
+        prompt=None,
     ):
         self.set_envs(envs)
         if not model:
@@ -239,13 +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
+        self.prompttext = prompt
 
     def translate(self, text) -> str:
         response = self.client.chat.completions.create(
             model=self.model,
             **self.options,
-            messages=self.prompt(text,self.prompttext),
+            messages=self.prompt(text, self.prompttext),
         )
         return response.choices[0].message.content.strip()
 
@@ -259,7 +267,14 @@ class AzureOpenAITranslator(BaseTranslator):
     }
 
     def __init__(
-        self, lang_in, lang_out, model, base_url=None, api_key=None, envs=None,prompt=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"]
@@ -273,13 +288,13 @@ class AzureOpenAITranslator(BaseTranslator):
             api_version="2024-06-01",
             api_key=api_key,
         )
-        self.prompttext=prompt
+        self.prompttext = prompt
 
     def translate(self, text) -> str:
         response = self.client.chat.completions.create(
             model=self.model,
             **self.options,
-            messages=self.prompt(text,self.prompttext),
+            messages=self.prompt(text, self.prompttext),
         )
         return response.choices[0].message.content.strip()
 
@@ -293,7 +308,14 @@ class ModelScopeTranslator(OpenAITranslator):
     }
 
     def __init__(
-        self, lang_in, lang_out, model, base_url=None, api_key=None, envs=None,prompt=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"
@@ -301,7 +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
+        self.prompttext = prompt
 
 
 class ZhipuTranslator(OpenAITranslator):
@@ -312,21 +334,21 @@ class ZhipuTranslator(OpenAITranslator):
         "ZHIPU_MODEL": "glm-4-flash",
     }
 
-    def __init__(self, lang_in, lang_out, model, envs=None,prompt=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
+        self.prompttext = prompt
 
     def translate(self, text) -> str:
         try:
             response = self.client.chat.completions.create(
                 model=self.model,
                 **self.options,
-                messages=self.prompt(text,self.prompttext),
+                messages=self.prompt(text, self.prompttext),
             )
         except openai.BadRequestError as e:
             if (
@@ -346,14 +368,14 @@ class SiliconTranslator(OpenAITranslator):
         "SILICON_MODEL": "Qwen/Qwen2.5-7B-Instruct",
     }
 
-    def __init__(self, lang_in, lang_out, model, envs=None,prompt=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
+        self.prompttext = prompt
 
 
 class GeminiTranslator(OpenAITranslator):
@@ -364,14 +386,15 @@ class GeminiTranslator(OpenAITranslator):
         "GEMINI_MODEL": "gemini-1.5-flash",
     }
 
-    def __init__(self, lang_in, lang_out, model, envs=None,prompt=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
+        self.prompttext = prompt
+
 
 class AzureTranslator(BaseTranslator):
     # https://github.com/Azure/azure-sdk-for-python
@@ -436,7 +459,7 @@ class AnythingLLMTranslator(BaseTranslator):
         "AnythingLLM_APIKEY": None,
     }
 
-    def __init__(self, lang_out, lang_in, model, envs=None,prompt=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"]
@@ -446,10 +469,10 @@ class AnythingLLMTranslator(BaseTranslator):
             "Authorization": f"Bearer {self.api_key}",
             "Content-Type": "application/json",
         }
-        self.prompttext=prompt
+        self.prompttext = prompt
 
     def translate(self, text):
-        messages = self.prompt(text,self.prompttext)
+        messages = self.prompt(text, self.prompttext)
         payload = {
             "message": messages,
             "mode": "chat",