BabelDOC 0.2.7__tar.gz → 0.2.9__tar.gz

{babeldoc-0.2.7 → babeldoc-0.2.9}/PKG-INFO

@@ -1,9 +1,11 @@
 Metadata-Version: 2.4
 Name: BabelDOC
-Version: 0.2.7
+Version: 0.2.9
 Summary: Yet Another Document Translator
 Project-URL: Homepage, https://github.com/funstory-ai/BabelDOC
 Project-URL: Issues, https://github.com/funstory-ai/BabelDOC/issues
+Author-email: awwaawwa <aw@funstory.ai>
+Maintainer-email: awwaawwa <aw@funstory.ai>
 License-Expression: AGPL-3.0
 License-File: LICENSE
 Keywords: PDF
@@ -20,10 +22,11 @@ Requires-Dist: numpy>=2.0.2
 Requires-Dist: onnx>=1.17.0
 Requires-Dist: onnxruntime>=1.16.1
 Requires-Dist: openai>=1.59.3
-Requires-Dist: opencv-python>=4.10.0.84
+Requires-Dist: opencv-python-headless>=4.10.0.84
 Requires-Dist: orjson>=3.10.14
 Requires-Dist: pdfminer-six>=20240706
 Requires-Dist: peewee>=3.17.8
+Requires-Dist: psutil>=7.0.0
 Requires-Dist: pydantic>=2.10.6
 Requires-Dist: pymupdf>=1.25.1
 Requires-Dist: python-levenshtein>=0.27.1
@@ -34,6 +37,10 @@ Requires-Dist: tiktoken>=0.9.0
 Requires-Dist: toml>=0.10.2
 Requires-Dist: tqdm>=4.67.1
 Requires-Dist: xsdata[cli,lxml,soap]>=24.12
+Provides-Extra: cuda
+Requires-Dist: onnxruntime-gpu>=1.16.1; extra == 'cuda'
+Provides-Extra: directml
+Requires-Dist: onnxruntime-directml>=1.16.1; extra == 'directml'
 Description-Content-Type: text/markdown
 
 <!-- # Yet Another Document Translator -->
@@ -172,6 +179,7 @@ uv run babeldoc --files example.pdf --files example2.pdf --openai --openai-model
 - `--enhance-compatibility`: Enable all compatibility enhancement options (equivalent to --skip-clean --dual-translate-first --disable-rich-text-translate)
 - `--use-alternating-pages-dual`: Use alternating pages mode for dual PDF. When enabled, original and translated pages are arranged in alternate order. When disabled (default), original and translated pages are shown side by side on the same page.
 - `--watermark-output-mode`: Control watermark output mode: 'watermarked' (default) adds watermark to translated PDF, 'no_watermark' doesn't add watermark, 'both' outputs both versions.
+- `--max-pages-per-part`: Maximum number of pages per part for split translation. If not set, no splitting will be performed.
 - `--no-watermark`: [DEPRECATED] Use --watermark-output-mode=no_watermark instead.
 
 > [!TIP]
@@ -179,6 +187,7 @@ uv run babeldoc --files example.pdf --files example2.pdf --openai --openai-model
 > - `--disable-rich-text-translate` can also help with compatibility by simplifying translation input
 > - However, using `--skip-clean` will result in larger file sizes
 > - If you encounter any compatibility issues, try using `--enhance-compatibility` first
+> - Use `--max-pages-per-part` for large documents to split them into smaller parts for translation and automatically merge them back.
 
 ### Translation Service Options
 
@@ -234,26 +243,6 @@ uv run babeldoc --files example.pdf --files example2.pdf --openai --openai-model
 
 Example Configuration:
 
-```toml
-[babeldoc]
-debug = true
-lang-in = "en-US"
-lang-out = "zh-CN"
-qps = 20
-# this is a comment
-# pages = 4
-openai = true
-openai-model = "SOME_ALSOME_MODEL"
-openai-base-url = "https://example.example/v1"
-openai-api-key = "[KEY]"
-# Offline assets management
-# generate-offline-assets = "/path/to/output/dir"
-# restore-offline-assets = "/path/to/offline_assets_package.zip"
-# All other options can also be set in the configuration file.
-```
-
-For a more comprehensive configuration example with offline assets management:
-
 ```toml
 [babeldoc]
 # Basic settings
@@ -263,12 +252,6 @@ lang-out = "zh-CN"
 qps = 10
 output = "/path/to/output/dir"
 
-# Translation service
-openai = true
-openai-model = "gpt-4o-mini"
-openai-base-url = "https://api.openai.com/v1"
-openai-api-key = "your-api-key-here"
-
 # PDF processing options
 split-short-lines = false
 short-line-split-factor = 0.8
@@ -277,8 +260,15 @@ dual-translate-first = false
 disable-rich-text-translate = false
 use-alternating-pages-dual = false
 watermark-output-mode = "watermarked"  # Choices: "watermarked", "no_watermark", "both"
+max-pages-per-part = 50  # Automatically split the document for translation and merge it back.
 # no-watermark = false  # DEPRECATED: Use watermark-output-mode instead
 
+# Translation service
+openai = true
+openai-model = "gpt-4o-mini"
+openai-base-url = "https://api.openai.com/v1"
+openai-api-key = "your-api-key-here"
+
 # Output control
 no-dual = false
 no-mono = false
@@ -293,6 +283,12 @@ report-interval = 0.5
 
 ## Python API
 
+> [!TIP]
+>
+> 1. Before pdf2zh 2.0 is released, you can temporarily use BabelDOC's Python API. However, after pdf2zh 2.0 is released, please directly use pdf2zh's Python API.
+>
+> 2. This project's Python API does not guarantee any compatibility. However, the Python API from pdf2zh will guarantee a certain level of compatibility.
+
 You can refer to the example in [main.py](https://github.com/funstory-ai/yadt/blob/main/babeldoc/main.py) to use BabelDOC's Python API.
 
 Please note:

{babeldoc-0.2.7 → babeldoc-0.2.9}/README.md

@@ -134,6 +134,7 @@ uv run babeldoc --files example.pdf --files example2.pdf --openai --openai-model
 - `--enhance-compatibility`: Enable all compatibility enhancement options (equivalent to --skip-clean --dual-translate-first --disable-rich-text-translate)
 - `--use-alternating-pages-dual`: Use alternating pages mode for dual PDF. When enabled, original and translated pages are arranged in alternate order. When disabled (default), original and translated pages are shown side by side on the same page.
 - `--watermark-output-mode`: Control watermark output mode: 'watermarked' (default) adds watermark to translated PDF, 'no_watermark' doesn't add watermark, 'both' outputs both versions.
+- `--max-pages-per-part`: Maximum number of pages per part for split translation. If not set, no splitting will be performed.
 - `--no-watermark`: [DEPRECATED] Use --watermark-output-mode=no_watermark instead.
 
 > [!TIP]
@@ -141,6 +142,7 @@ uv run babeldoc --files example.pdf --files example2.pdf --openai --openai-model
 > - `--disable-rich-text-translate` can also help with compatibility by simplifying translation input
 > - However, using `--skip-clean` will result in larger file sizes
 > - If you encounter any compatibility issues, try using `--enhance-compatibility` first
+> - Use `--max-pages-per-part` for large documents to split them into smaller parts for translation and automatically merge them back.
 
 ### Translation Service Options
 
@@ -196,26 +198,6 @@ uv run babeldoc --files example.pdf --files example2.pdf --openai --openai-model
 
 Example Configuration:
 
-```toml
-[babeldoc]
-debug = true
-lang-in = "en-US"
-lang-out = "zh-CN"
-qps = 20
-# this is a comment
-# pages = 4
-openai = true
-openai-model = "SOME_ALSOME_MODEL"
-openai-base-url = "https://example.example/v1"
-openai-api-key = "[KEY]"
-# Offline assets management
-# generate-offline-assets = "/path/to/output/dir"
-# restore-offline-assets = "/path/to/offline_assets_package.zip"
-# All other options can also be set in the configuration file.
-```
-
-For a more comprehensive configuration example with offline assets management:
-
 ```toml
 [babeldoc]
 # Basic settings
@@ -225,12 +207,6 @@ lang-out = "zh-CN"
 qps = 10
 output = "/path/to/output/dir"
 
-# Translation service
-openai = true
-openai-model = "gpt-4o-mini"
-openai-base-url = "https://api.openai.com/v1"
-openai-api-key = "your-api-key-here"
-
 # PDF processing options
 split-short-lines = false
 short-line-split-factor = 0.8
@@ -239,8 +215,15 @@ dual-translate-first = false
 disable-rich-text-translate = false
 use-alternating-pages-dual = false
 watermark-output-mode = "watermarked"  # Choices: "watermarked", "no_watermark", "both"
+max-pages-per-part = 50  # Automatically split the document for translation and merge it back.
 # no-watermark = false  # DEPRECATED: Use watermark-output-mode instead
 
+# Translation service
+openai = true
+openai-model = "gpt-4o-mini"
+openai-base-url = "https://api.openai.com/v1"
+openai-api-key = "your-api-key-here"
+
 # Output control
 no-dual = false
 no-mono = false
@@ -255,6 +238,12 @@ report-interval = 0.5
 
 ## Python API
 
+> [!TIP]
+>
+> 1. Before pdf2zh 2.0 is released, you can temporarily use BabelDOC's Python API. However, after pdf2zh 2.0 is released, please directly use pdf2zh's Python API.
+>
+> 2. This project's Python API does not guarantee any compatibility. However, the Python API from pdf2zh will guarantee a certain level of compatibility.
+
 You can refer to the example in [main.py](https://github.com/funstory-ai/yadt/blob/main/babeldoc/main.py) to use BabelDOC's Python API.
 
 Please note:

babeldoc-0.2.9/babeldoc/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.2.9"

{babeldoc-0.2.7 → babeldoc-0.2.9}/babeldoc/const.py

@@ -3,7 +3,7 @@ import shutil
 import subprocess
 from pathlib import Path
 
-__version__ = "0.2.7"
+__version__ = "0.2.9"
 
 CACHE_FOLDER = Path.home() / ".cache" / "babeldoc"
 

{babeldoc-0.2.7 → babeldoc-0.2.9}/babeldoc/document_il/backend/pdf_creater.py

@@ -552,7 +552,7 @@ class PDFCreater:
         while process.is_alive():
             if time.time() - start_time > timeout:
                 logger.warning(
-                    f"PDF save with clean=True timeout after {timeout} seconds, terminating subprocess"
+                    f"PDF save with clean={clean} timeout after {timeout} seconds, terminating subprocess"
                 )
                 process.terminate()
                 try:
@@ -569,7 +569,7 @@ class PDFCreater:
                     logger.error(f"Error terminating PDF save process: {e}")
 
                 # Fallback to save without clean parameter
-                logger.info("Falling back to save without clean=True")
+                logger.info("Falling back to save with clean=False")
                 try:
                     pdf.save(
                         output_path,
@@ -598,7 +598,7 @@ class PDFCreater:
             and Path(temp_output).exists()
             and Path(temp_output).stat().st_size > 0
         ):
-            logger.info("PDF save with clean=True completed successfully")
+            logger.info(f"PDF save with clean={clean} completed successfully")
             # Copy the successfully created file to the target path
             try:
                 import shutil
@@ -614,7 +614,7 @@ class PDFCreater:
                 Path(temp_output).unlink()
         else:
             logger.warning(
-                f"PDF save with clean=True failed with exit code {exit_code} or produced empty file"
+                f"PDF save with clean={clean} failed with exit code {exit_code} or produced empty file"
             )
             # Fallback to save without clean parameter
             try:

{babeldoc-0.2.7 → babeldoc-0.2.9}/babeldoc/document_il/midend/il_translator.py

@@ -603,7 +603,7 @@ class ILTranslator:
                 )
             except Exception as e:
                 logger.exception(
-                    f"Error translating paragraph. Paragraph: {paragraph}. Error: {e}. ",
+                    f"Error translating paragraph. Paragraph: {paragraph.debug_id} ({paragraph.unicode}). Error: {e}. ",
                 )
                 # ignore error and continue
                 return

{babeldoc-0.2.7 → babeldoc-0.2.9}/babeldoc/document_il/midend/il_translator_llm_only.py

@@ -43,6 +43,9 @@ class ILTranslatorLLMOnly:
         self.translate_engine = translate_engine
         self.translation_config = translation_config
         self.font_mapper = FontMapper(translation_config)
+        self.shared_context_cross_split_part = (
+            translation_config.shared_context_cross_split_part
+        )
 
         self.il_translator = ILTranslator(
             translate_engine=translate_engine,
@@ -73,8 +76,15 @@ class ILTranslatorLLMOnly:
     def translate(self, docs: Document) -> None:
         tracker = DocumentTranslateTracker()
 
-        # Try to find the first title paragraph
-        title_paragraph = self.find_title_paragraph(docs)
+        if not self.translation_config.shared_context_cross_split_part.first_paragraph:
+            # Try to find the first title paragraph
+            title_paragraph = self.find_title_paragraph(docs)
+            self.translation_config.shared_context_cross_split_part.first_paragraph = (
+                title_paragraph
+            )
+            self.translation_config.shared_context_cross_split_part.recent_title_paragraph = title_paragraph
+            if title_paragraph:
+                logger.info(f"Found first title paragraph: {title_paragraph.unicode}")
 
         # count total paragraph
         total = sum(
@@ -105,7 +115,6 @@ class ILTranslatorLLMOnly:
                             executor,
                             pbar,
                             tracker.new_page(),
-                            title_paragraph,
                             executor2,
                         )
 
@@ -122,7 +131,6 @@ class ILTranslatorLLMOnly:
         executor: concurrent.futures.ThreadPoolExecutor,
         pbar: tqdm | None = None,
         tracker: PageTranslateTracker = None,
-        title_paragraph: PdfParagraph | None = None,
         executor2: concurrent.futures.ThreadPoolExecutor | None = None,
     ):
         self.translation_config.raise_if_cancelled()
@@ -136,7 +144,6 @@ class ILTranslatorLLMOnly:
                 page_xobj_font_map[xobj.xobj_id][font.font_id] = font
 
         paragraphs = []
-        local_title_paragraph = None
 
         total_unicode_counts = 0
         for paragraph in page.pdf_paragraph:
@@ -146,7 +153,7 @@ class ILTranslatorLLMOnly:
             total_unicode_counts += len(paragraph.unicode)
             paragraphs.append(paragraph)
             if paragraph.layout_label == "title":
-                local_title_paragraph = paragraph
+                self.shared_context_cross_split_part.recent_title_paragraph = paragraph
 
             if total_unicode_counts > 1200 or len(paragraphs) > 4:
                 executor.submit(
@@ -155,8 +162,8 @@ class ILTranslatorLLMOnly:
                     pbar,
                     page_font_map,
                     page_xobj_font_map,
-                    title_paragraph,
-                    local_title_paragraph,
+                    self.translation_config.shared_context_cross_split_part.first_paragraph,
+                    self.translation_config.shared_context_cross_split_part.recent_title_paragraph,
                     executor2,
                 )
                 paragraphs = []
@@ -169,8 +176,8 @@ class ILTranslatorLLMOnly:
                 pbar,
                 page_font_map,
                 page_xobj_font_map,
-                title_paragraph,
-                local_title_paragraph,
+                self.translation_config.shared_context_cross_split_part.first_paragraph,
+                self.translation_config.shared_context_cross_split_part.recent_title_paragraph,
                 executor2,
             )
 
@@ -328,9 +335,9 @@ class ILTranslatorLLMOnly:
                         translate_input,
                         translated_text,
                     )
+                    should_fallback = False
                     if pbar:
                         pbar.advance(1)
-                    should_fallback = False
                 except Exception as e:
                     logger.exception(f"Error translating paragraph. Error: {e}.")
                     # Ignore error and continue

{babeldoc-0.2.7 → babeldoc-0.2.9}/babeldoc/document_il/midend/styles_and_formulas.py

@@ -528,6 +528,7 @@ class StylesAndFormulas:
             r"|NimbusRomNo9L.+"
             r"|GlosaMath.+"
             r"|URWPalladioL.+"
+            r"|CMSS.+"
             r")$"
         )
         if self.translation_config.formular_font_pattern:

{babeldoc-0.2.7 → babeldoc-0.2.9}/babeldoc/document_il/translator/translator.py

@@ -62,6 +62,31 @@ def set_translate_rate_limiter(max_qps):
     _translate_rate_limiter.set_max_qps(max_qps)
 
 
+class AtomicInteger:
+    def __init__(self, value=0):
+        self._value = int(value)
+        self._lock = threading.Lock()
+
+    def inc(self, d=1):
+        with self._lock:
+            self._value += int(d)
+            return self._value
+
+    def dec(self, d=1):
+        return self.inc(-d)
+
+    @property
+    def value(self):
+        with self._lock:
+            return self._value
+
+    @value.setter
+    def value(self, v):
+        with self._lock:
+            self._value = int(v)
+            return self._value
+
+
 class BaseTranslator(ABC):
     # Due to cache limitations, name should be within 20 characters.
     # cache.py: translate_engine = CharField(max_length=20)
@@ -195,6 +220,9 @@ class OpenAITranslator(BaseTranslator):
         self.model = model
         self.add_cache_impact_parameters("model", self.model)
         self.add_cache_impact_parameters("prompt", self.prompt(""))
+        self.token_count = AtomicInteger()
+        self.prompt_token_count = AtomicInteger()
+        self.completion_token_count = AtomicInteger()
 
     @retry(
         retry=retry_if_exception_type(openai.RateLimitError),
@@ -211,6 +239,9 @@ class OpenAITranslator(BaseTranslator):
             **self.options,
             messages=self.prompt(text),
         )
+        self.token_count.inc(response.usage.total_tokens)
+        self.prompt_token_count.inc(response.usage.prompt_tokens)
+        self.completion_token_count.inc(response.usage.completion_tokens)
         return response.choices[0].message.content.strip()
 
     def prompt(self, text):
@@ -248,6 +279,9 @@ class OpenAITranslator(BaseTranslator):
                 },
             ],
         )
+        self.token_count.inc(response.usage.total_tokens)
+        self.prompt_token_count.inc(response.usage.prompt_tokens)
+        self.completion_token_count.inc(response.usage.completion_tokens)
         return response.choices[0].message.content.strip()
 
     def get_formular_placeholder(self, placeholder_id: int):

{babeldoc-0.2.7 → babeldoc-0.2.9}/babeldoc/document_il/utils/fontmap.py

@@ -172,10 +172,14 @@ class FontMapper:
 
         font_id = {}
         xreflen = doc_zh.xref_length()
+        total = xreflen - 1 + len(font_list) + len(il.page) + len(font_list)
         with self.translation_config.progress_monitor.stage_start(
             self.stage_name,
-            xreflen - 1 + len(font_list) + len(il.page) + len(font_list),
+            total,
         ) as pbar:
+            if not il.page:
+                pbar.advance(total)
+                return
             for font in font_list:
                 if font[0] in font_id:
                     continue

{babeldoc-0.2.7 → babeldoc-0.2.9}/babeldoc/docvision/doclayout.py

@@ -2,12 +2,22 @@ import abc
 import ast
 import logging
 import platform
+import re
 from collections.abc import Generator
 
 import cv2
 import numpy as np
-import onnx
-import onnxruntime
+
+try:
+    import onnx
+    import onnxruntime
+except ImportError as e:
+    if "DLL load failed" in str(e):
+        raise OSError(
+            "Microsoft Visual C++ Redistributable is not installed. "
+            "Download it at https://aka.ms/vs/17/release/vc_redist.x64.exe"
+        ) from e
+    raise
 import pymupdf
 
 import babeldoc.document_il.il_version_1
@@ -91,26 +101,6 @@ class YoloBox:
 # 检测操作系统类型
 os_name = platform.system()
 
-providers = []
-
-if os_name == "Darwin" and False:  # Temporarily disable CoreML due to some issues
-    providers.append(
-        (
-            "CoreMLExecutionProvider",
-            {
-                "ModelFormat": "MLProgram",
-                "MLComputeUnits": "ALL",
-                "RequireStaticInputShapes": "0",
-                "EnableOnSubgraphs": "0",
-            },
-        ),
-    )
-    # workaround for CoreML batch inference issues
-    max_batch_size = 1
-else:
-    max_batch_size = 1
-providers.append("CPUExecutionProvider")  # CPU 执行提供者作为通用后备选项
-
 
 class OnnxModel(DocLayoutModel):
     def __init__(self, model_path: str):
@@ -120,7 +110,13 @@ class OnnxModel(DocLayoutModel):
         metadata = {d.key: d.value for d in model.metadata_props}
         self._stride = ast.literal_eval(metadata["stride"])
         self._names = ast.literal_eval(metadata["names"])
+        providers = []
 
+        available_providers = onnxruntime.get_available_providers()
+        for provider in available_providers:
+            if re.match(r"dml|cuda|coreml|cpu", provider, re.IGNORECASE):
+                logger.info(f"Available Provider: {provider}")
+                providers.append(provider)
         self.model = onnxruntime.InferenceSession(
             model.SerializeToString(),
             providers=providers,
@@ -228,7 +224,7 @@ class OnnxModel(DocLayoutModel):
 
         total_images = len(image)
         results = []
-        batch_size = min(batch_size, max_batch_size)
+        batch_size = 1
 
         # Process images in batches
         for i in range(0, total_images, batch_size):
@@ -237,7 +233,7 @@ class OnnxModel(DocLayoutModel):
 
             # Calculate target size based on the maximum height in the batch
             max_height = max(img.shape[0] for img in batch_images)
-            target_imgsz = int(max_height / 32) * 32
+            target_imgsz = 1024
 
             # Preprocess batch
             processed_batch = []