BabelDOC 0.2.1__tar.gz → 0.2.2__tar.gz

{babeldoc-0.2.1 → babeldoc-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: BabelDOC
-Version: 0.2.1
+Version: 0.2.2
 Summary: Yet Another Document Translator
 Project-URL: Homepage, https://github.com/funstory-ai/BabelDOC
 Project-URL: Issues, https://github.com/funstory-ai/BabelDOC/issues
@@ -26,9 +26,11 @@ Requires-Dist: pdfminer-six>=20240706
 Requires-Dist: peewee>=3.17.8
 Requires-Dist: pydantic>=2.10.6
 Requires-Dist: pymupdf>=1.25.1
+Requires-Dist: python-levenshtein>=0.27.1
 Requires-Dist: rich>=13.9.4
 Requires-Dist: scikit-image>=0.25.2
 Requires-Dist: tenacity>=9.0.0
+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

babeldoc-0.2.2/babeldoc/__init__.py

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

{babeldoc-0.2.1 → babeldoc-0.2.2}/babeldoc/assets/assets.py

@@ -14,6 +14,7 @@ from babeldoc.assets.embedding_assets_metadata import (
 from babeldoc.assets.embedding_assets_metadata import EMBEDDING_FONT_METADATA
 from babeldoc.assets.embedding_assets_metadata import FONT_METADATA_URL
 from babeldoc.assets.embedding_assets_metadata import FONT_URL_BY_UPSTREAM
+from babeldoc.assets.embedding_assets_metadata import TIKTOKEN_CACHES
 from babeldoc.const import get_cache_file_path
 from tenacity import retry
 from tenacity import stop_after_attempt
@@ -292,6 +293,9 @@ async def download_all_fonts_async(client: httpx.AsyncClient | None = None):
 
 async def async_warmup():
     logger.info("Downloading all assets...")
+    from tiktoken import encoding_for_model
+
+    _ = encoding_for_model("gpt-4o")
     async with httpx.AsyncClient() as client:
         onnx_task = asyncio.create_task(get_doclayout_onnx_model_path_async(client))
         font_tasks = asyncio.create_task(download_all_fonts_async(client))
@@ -306,6 +310,7 @@ def generate_all_assets_file_list():
     result = {}
     result["fonts"] = []
     result["models"] = []
+    result["tiktoken"] = []
     for font_file_name in EMBEDDING_FONT_METADATA:
         result["fonts"].append(
             {
@@ -313,6 +318,13 @@ def generate_all_assets_file_list():
                 "sha3_256": EMBEDDING_FONT_METADATA[font_file_name]["sha3_256"],
             }
         )
+    for tiktoken_file, sha3_256 in TIKTOKEN_CACHES.items():
+        result["tiktoken"].append(
+            {
+                "name": tiktoken_file,
+                "sha3_256": sha3_256,
+            }
+        )
     result["models"].append(
         {
             "name": "doclayout_yolo_docstructbench_imgsz1024.onnx",

{babeldoc-0.2.1 → babeldoc-0.2.2}/babeldoc/assets/embedding_assets_metadata.py

@@ -2,6 +2,9 @@ DOCLAYOUT_YOLO_DOCSTRUCTBENCH_IMGSZ1024ONNX_SHA3_256 = (
     "60be061226930524958b5465c8c04af3d7c03bcb0beb66454f5da9f792e3cf2a"
 )
 
+TIKTOKEN_CACHES = {
+    "fb374d419588a4632f3f557e76b4b70aebbca790": "cb04bcda5782cfbbe77f2f991d92c0ea785d9496ef1137c91dfc3c8c324528d6"
+}
 
 FONT_METADATA_URL = {
     "github": "https://raw.githubusercontent.com/funstory-ai/BabelDOC-Assets/refs/heads/main/font_metadata.json",

{babeldoc-0.2.1 → babeldoc-0.2.2}/babeldoc/const.py

@@ -1,8 +1,9 @@
+import os
 import shutil
 import subprocess
 from pathlib import Path
 
-__version__ = "0.2.1"
+__version__ = "0.2.2"
 
 CACHE_FOLDER = Path.home() / ".cache" / "babeldoc"
 
@@ -35,3 +36,7 @@ try:
     )
 except (OSError, FileNotFoundError, subprocess.CalledProcessError):
     WATERMARK_VERSION = f"v{__version__}"
+
+TIKTOKEN_CACHE_FOLDER = CACHE_FOLDER / "tiktoken"
+TIKTOKEN_CACHE_FOLDER.mkdir(parents=True, exist_ok=True)
+os.environ["TIKTOKEN_CACHE_DIR"] = str(TIKTOKEN_CACHE_FOLDER)

{babeldoc-0.2.1 → babeldoc-0.2.2}/babeldoc/document_il/il_version_1.py

@@ -750,6 +750,12 @@ class PdfParagraph:
             "type": "Attribute",
         },
     )
+    layout_label: str | None = field(
+        default=None,
+        metadata={
+            "type": "Attribute",
+        },
+    )
 
 
 @dataclass

{babeldoc-0.2.1 → babeldoc-0.2.2}/babeldoc/document_il/il_version_1.rnc

@@ -114,6 +114,7 @@ PDFParagraph =
     attribute vertical { xsd:boolean }?,
     attribute FirstLineIndent { xsd:boolean }?,
     attribute debug_id { xsd:string }?,
+    attribute layout_label { xsd:string }?,
     Box,
     PDFStyle,
     PDFParagraphComposition*

{babeldoc-0.2.1 → babeldoc-0.2.2}/babeldoc/document_il/il_version_1.rng

@@ -334,6 +334,11 @@
           <data type="string"/>
         </attribute>
       </optional>
+      <optional>
+        <attribute name="layout_label">
+          <data type="string"/>
+        </attribute>
+      </optional>
       <ref name="Box"/>
       <ref name="PDFStyle"/>
       <zeroOrMore>

{babeldoc-0.2.1 → babeldoc-0.2.2}/babeldoc/document_il/il_version_1.xsd

@@ -187,6 +187,7 @@
       <xs:attribute name="vertical" type="xs:boolean"/>
       <xs:attribute name="FirstLineIndent" type="xs:boolean"/>
       <xs:attribute name="debug_id" type="xs:string"/>
+      <xs:attribute name="layout_label" type="xs:string"/>
     </xs:complexType>
   </xs:element>
   <xs:element name="pdfParagraphComposition">

{babeldoc-0.2.1 → babeldoc-0.2.2}/babeldoc/document_il/midend/il_translator.py

@@ -524,6 +524,56 @@ class ILTranslator:
 
         return result
 
+    def pre_translate_paragraph(
+        self,
+        paragraph: PdfParagraph,
+        tracker: ParagraphTranslateTracker,
+        page_font_map: dict[str, PdfFont],
+        xobj_font_map: dict[int, dict[str, PdfFont]],
+    ):
+        """Pre-translation processing: prepare text for translation."""
+        if paragraph.vertical:
+            return None, None
+        tracker.set_pdf_unicode(paragraph.unicode)
+        if paragraph.xobj_id in xobj_font_map:
+            page_font_map = xobj_font_map[paragraph.xobj_id]
+        translate_input = self.get_translate_input(paragraph, page_font_map)
+        if not translate_input:
+            return None, None
+        tracker.set_input(translate_input.unicode)
+        text = translate_input.unicode
+        if len(text) < self.translation_config.min_text_length:
+            logger.debug(
+                f"Text too short to translate, skip. Text: {text}. Paragraph id: {paragraph.debug_id}."
+            )
+            return None, None
+        return text, translate_input
+
+    def post_translate_paragraph(
+        self,
+        paragraph: PdfParagraph,
+        tracker: ParagraphTranslateTracker,
+        translate_input,
+        translated_text: str,
+    ):
+        """Post-translation processing: update paragraph with translated text."""
+        tracker.set_output(translated_text)
+        if translated_text == translate_input.unicode:
+            return False
+        paragraph.unicode = translated_text
+        paragraph.pdf_paragraph_composition = self.parse_translate_output(
+            translate_input, translated_text
+        )
+        for composition in paragraph.pdf_paragraph_composition:
+            if (
+                composition.pdf_same_style_unicode_characters
+                and composition.pdf_same_style_unicode_characters.pdf_style is None
+            ):
+                composition.pdf_same_style_unicode_characters.pdf_style = (
+                    paragraph.pdf_style
+                )
+        return True
+
     def translate_paragraph(
         self,
         paragraph: PdfParagraph,
@@ -532,51 +582,25 @@ class ILTranslator:
         page_font_map: dict[str, PdfFont] = None,
         xobj_font_map: dict[int, dict[str, PdfFont]] = None,
     ):
+        """Translate a paragraph using pre and post processing functions."""
         self.translation_config.raise_if_cancelled()
         with PbarContext(pbar):
             try:
-                if paragraph.vertical:
-                    return
-
-                tracker.set_pdf_unicode(paragraph.unicode)
-                if paragraph.xobj_id in xobj_font_map:
-                    page_font_map = xobj_font_map[paragraph.xobj_id]
-                translate_input = self.get_translate_input(paragraph, page_font_map)
-                if not translate_input:
-                    return
-
-                tracker.set_input(translate_input.unicode)
-
-                text = translate_input.unicode
-
-                if len(text) < self.translation_config.min_text_length:
-                    logger.debug(
-                        f"Text too short to translate, skip. Text: {text}. Paragraph id: {paragraph.debug_id}.",
-                    )
+                # Pre-translation processing
+                text, translate_input = self.pre_translate_paragraph(
+                    paragraph, tracker, page_font_map, xobj_font_map
+                )
+                if text is None:
                     return
 
+                # Perform translation
                 translated_text = self.translate_engine.translate(text)
                 translated_text = re.sub(r"[. 。…，]{20,}", ".", translated_text)
 
-                tracker.set_output(translated_text)
-
-                if translated_text == text:
-                    return
-
-                paragraph.unicode = translated_text
-                paragraph.pdf_paragraph_composition = self.parse_translate_output(
-                    translate_input,
-                    translated_text,
+                # Post-translation processing
+                self.post_translate_paragraph(
+                    paragraph, tracker, translate_input, translated_text
                 )
-                for composition in paragraph.pdf_paragraph_composition:
-                    if (
-                        composition.pdf_same_style_unicode_characters
-                        and composition.pdf_same_style_unicode_characters.pdf_style
-                        is None
-                    ):
-                        composition.pdf_same_style_unicode_characters.pdf_style = (
-                            paragraph.pdf_style
-                        )
             except Exception as e:
                 logger.exception(
                     f"Error translating paragraph. Paragraph: {paragraph}. Error: {e}. ",

babeldoc-0.2.2/babeldoc/document_il/midend/il_translator_llm_only.py

@@ -0,0 +1,375 @@
+import concurrent.futures
+import json
+import logging
+import re
+from pathlib import Path
+
+import Levenshtein
+import tiktoken
+from tqdm import tqdm
+
+from babeldoc.document_il import Document
+from babeldoc.document_il import Page
+from babeldoc.document_il import PdfFont
+from babeldoc.document_il import PdfParagraph
+from babeldoc.document_il.midend.il_translator import DocumentTranslateTracker
+from babeldoc.document_il.midend.il_translator import ILTranslator
+from babeldoc.document_il.midend.il_translator import PageTranslateTracker
+from babeldoc.document_il.translator.translator import BaseTranslator
+from babeldoc.document_il.utils.fontmap import FontMapper
+from babeldoc.translation_config import TranslationConfig
+
+enc = tiktoken.encoding_for_model("gpt-4o")
+
+logger = logging.getLogger(__name__)
+
+
+class BatchParagraph:
+    def __init__(
+        self, paragraphs: list[PdfParagraph], page_tracker: PageTranslateTracker
+    ):
+        self.paragraphs = paragraphs
+        self.trackers = [page_tracker.new_paragraph() for _ in paragraphs]
+
+
+class ILTranslatorLLMOnly:
+    stage_name = "Translate Paragraphs"
+
+    def __init__(
+        self,
+        translate_engine: BaseTranslator,
+        translation_config: TranslationConfig,
+    ):
+        self.translate_engine = translate_engine
+        self.translation_config = translation_config
+        self.font_mapper = FontMapper(translation_config)
+
+        self.il_translator = ILTranslator(
+            translate_engine=translate_engine,
+            translation_config=translation_config,
+        )
+
+        try:
+            self.translate_engine.do_llm_translate(None)
+        except NotImplementedError as e:
+            raise ValueError("LLM translator not supported") from e
+
+    def find_title_paragraph(self, docs: Document) -> PdfParagraph | None:
+        """Find the first paragraph with layout_label 'title' in the document.
+
+        Args:
+            docs: The document to search in
+
+        Returns:
+            The first title paragraph found, or None if no title paragraph exists
+        """
+        for page in docs.page:
+            for paragraph in page.pdf_paragraph:
+                if paragraph.layout_label == "title":
+                    logger.info(f"Found title paragraph: {paragraph.unicode}")
+                    return paragraph
+        return None
+
+    def translate(self, docs: Document) -> None:
+        tracker = DocumentTranslateTracker()
+
+        # Try to find the first title paragraph
+        title_paragraph = self.find_title_paragraph(docs)
+
+        # count total paragraph
+        total = sum(
+            [
+                len([p for p in page.pdf_paragraph if p.debug_id is not None])
+                for page in docs.page
+            ]
+        )
+        with self.translation_config.progress_monitor.stage_start(
+            self.stage_name,
+            total,
+        ) as pbar:
+            with concurrent.futures.ThreadPoolExecutor(
+                max_workers=self.translation_config.qps,
+            ) as executor2:
+                with concurrent.futures.ThreadPoolExecutor(
+                    max_workers=self.translation_config.qps,
+                ) as executor:
+                    for page in docs.page:
+                        self.process_page(
+                            page,
+                            executor,
+                            pbar,
+                            tracker.new_page(),
+                            title_paragraph,
+                            executor2,
+                        )
+
+        path = self.translation_config.get_working_file_path("translate_tracking.json")
+
+        if self.translation_config.debug:
+            logger.debug(f"save translate tracking to {path}")
+            with Path(path).open("w", encoding="utf-8") as f:
+                f.write(tracker.to_json())
+
+    def process_page(
+        self,
+        page: Page,
+        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()
+        page_font_map = {}
+        for font in page.pdf_font:
+            page_font_map[font.font_id] = font
+        page_xobj_font_map = {}
+        for xobj in page.pdf_xobject:
+            page_xobj_font_map[xobj.xobj_id] = page_font_map.copy()
+            for font in xobj.pdf_font:
+                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:
+            if paragraph.debug_id is None:
+                continue
+            # self.translate_paragraph(paragraph, pbar,tracker.new_paragraph(), page_font_map, page_xobj_font_map)
+            total_unicode_counts += len(paragraph.unicode)
+            paragraphs.append(paragraph)
+            if paragraph.layout_label == "title":
+                local_title_paragraph = paragraph
+
+            if total_unicode_counts > 1200 or len(paragraphs) > 4:
+                executor.submit(
+                    self.translate_paragraph,
+                    BatchParagraph(paragraphs, tracker),
+                    pbar,
+                    page_font_map,
+                    page_xobj_font_map,
+                    title_paragraph,
+                    local_title_paragraph,
+                    executor2,
+                )
+                paragraphs = []
+                total_unicode_counts = 0
+
+        if paragraphs:
+            executor.submit(
+                self.translate_paragraph,
+                BatchParagraph(paragraphs, tracker),
+                pbar,
+                page_font_map,
+                page_xobj_font_map,
+                title_paragraph,
+                local_title_paragraph,
+                executor2,
+            )
+
+    def translate_paragraph(
+        self,
+        batch_paragraph: BatchParagraph,
+        pbar: tqdm | None = None,
+        page_font_map: dict[str, PdfFont] = None,
+        xobj_font_map: dict[int, dict[str, PdfFont]] = None,
+        title_paragraph: PdfParagraph | None = None,
+        local_title_paragraph: PdfParagraph | None = None,
+        executor=None,
+    ):
+        """Translate a paragraph using pre and post processing functions."""
+        self.translation_config.raise_if_cancelled()
+        try:
+            inputs = []
+
+            for i in range(len(batch_paragraph.paragraphs)):
+                paragraph = batch_paragraph.paragraphs[i]
+                tracker = batch_paragraph.trackers[i]
+                text, translate_input = self.il_translator.pre_translate_paragraph(
+                    paragraph, tracker, page_font_map, xobj_font_map
+                )
+                if text is None:
+                    pbar.advance(1)
+                    continue
+                inputs.append((text, translate_input, paragraph, tracker))
+            if not inputs:
+                return
+            json_format_input = []
+
+            for id_, input_text in enumerate(inputs):
+                json_format_input.append(
+                    {
+                        "id": id_,
+                        "input": input_text[0],
+                        "layout_label": input_text[2].layout_label,
+                    }
+                )
+            json_format_input = json.dumps(
+                json_format_input, ensure_ascii=False, indent=2
+            )
+            llm_input = [
+                "You are a professional, authentic machine translation engine."
+            ]
+
+            if title_paragraph:
+                llm_input.append(
+                    f"The first title in the full text: {title_paragraph.unicode}"
+                )
+            if (
+                local_title_paragraph
+                and local_title_paragraph.debug_id != title_paragraph.debug_id
+            ):
+                llm_input.append(
+                    f"The most similar title in the full text: {local_title_paragraph.unicode}"
+                )
+            # Create a structured prompt template for LLM translation
+            prompt_template = (
+                f"""
+    You will be given a JSON formatted input containing entries with "id" and "input" fields. Here is the input:
+    
+    ```json
+    {json_format_input}
+    ```
+    
+    For each entry in the JSON, translate the contents of the "input" field into {self.translation_config.lang_out}.
+    Write the translation back into the "output" field for that entry.
+    
+    """
+                + """
+    Here is an example of the expected format:
+    
+    <example>
+    ```json
+    Input:
+    {
+        "id": 1,
+        "input": "Source",
+        "layout_label": "plain text"
+    }
+    ```
+    Output:
+    ```json
+    {
+        "id": 1,
+        "output": "Translation"
+    }
+    ```
+    </example>
+    
+    Please return the translated json directly without wrapping ```json``` tag or include any additional information.
+    """
+            )
+            llm_input.append(prompt_template)
+
+            final_input = "\n".join(llm_input).strip()
+            llm_output = self.translate_engine.llm_translate(final_input)
+            llm_output = llm_output.strip()
+
+            llm_output = self._clean_json_output(llm_output)
+
+            parsed_output = json.loads(llm_output)
+
+            translation_results = {item["id"]: item["output"] for item in parsed_output}
+
+            if len(translation_results) != len(inputs):
+                raise Exception(
+                    f"Translation results length mismatch. Expected: {len(inputs)}, Got: {len(translation_results)}"
+                )
+
+            for id_, output in translation_results.items():
+                should_fallback = True
+                try:
+                    if not isinstance(output, str):
+                        logger.warning(
+                            f"Translation result is not a string. Output: {output}"
+                        )
+                        continue
+
+                    id_ = int(id_)  # Ensure id is an integer
+                    if id_ >= len(inputs):
+                        logger.warning(f"Invalid id {id_}, skipping")
+                        continue
+
+                    # Clean up any excessive punctuation in the translated text
+                    translated_text = re.sub(r"[. 。…，]{20,}", ".", output)
+
+                    # Get the original input for this translation
+                    translate_input = inputs[id_][1]
+
+                    input_unicode = inputs[id_][2].unicode
+                    output_unicode = translated_text
+
+                    input_token_count = len(enc.encode(input_unicode))
+                    output_token_count = len(enc.encode(output_unicode))
+
+                    if not (0.3 < output_token_count / input_token_count < 3):
+                        logger.warning(
+                            f"Translation result is too long or too short. Input: {input_token_count}, Output: {output_token_count}"
+                        )
+                        continue
+
+                    edit_distance = Levenshtein.distance(input_unicode, output_unicode)
+                    if edit_distance < 5 and input_token_count > 20:
+                        logger.warning(
+                            f"Translation result edit distance is too small. distance: {edit_distance}, input: {input_unicode}, output: {output_unicode}"
+                        )
+                        continue
+                    # Apply the translation to the paragraph
+                    self.il_translator.post_translate_paragraph(
+                        inputs[id_][2],
+                        inputs[id_][3],
+                        translate_input,
+                        translated_text,
+                    )
+                    if pbar:
+                        pbar.advance(1)
+                    should_fallback = False
+                except Exception as e:
+                    logger.exception(f"Error translating paragraph. Error: {e}.")
+                    # Ignore error and continue
+                    continue
+                finally:
+                    if should_fallback:
+                        logger.warning(
+                            f"Fallback to simple translation. paragraph id: {inputs[id_][2].debug_id}"
+                        )
+                        executor.submit(
+                            self.il_translator.translate_paragraph,
+                            inputs[id_][2],
+                            pbar,
+                            inputs[id_][3],
+                            page_font_map,
+                            xobj_font_map,
+                        )
+
+        except Exception as e:
+            logger.warning(f"Error {e} during translation. try fallback")
+            for i in range(len(batch_paragraph.paragraphs)):
+                paragraph = batch_paragraph.paragraphs[i]
+                tracker = batch_paragraph.trackers[i]
+                if paragraph.debug_id is None:
+                    continue
+                executor.submit(
+                    self.il_translator.translate_paragraph,
+                    paragraph,
+                    pbar,
+                    tracker,
+                    page_font_map,
+                    xobj_font_map,
+                )
+
+    def _clean_json_output(self, llm_output: str) -> str:
+        # Clean up JSON output by removing common wrapper tags
+        llm_output = llm_output.strip()
+        if llm_output.startswith("<json>"):
+            llm_output = llm_output[6:]
+        if llm_output.endswith("</json>"):
+            llm_output = llm_output[:-7]
+        if llm_output.startswith("```json"):
+            llm_output = llm_output[7:]
+        if llm_output.startswith("```"):
+            llm_output = llm_output[3:]
+        if llm_output.endswith("```"):
+            llm_output = llm_output[:-3]
+        return llm_output.strip()

{babeldoc-0.2.1 → babeldoc-0.2.2}/babeldoc/document_il/midend/paragraph_finder.py

@@ -177,6 +177,9 @@ class ParagraphFinder:
                         current_paragraph = PdfParagraph(
                             pdf_paragraph_composition=[line],
                             debug_id=generate_base58_id(),
+                            layout_label=char_layout.name
+                            if not current_layout
+                            else current_layout.name,
                         )
                         paragraphs.append(current_paragraph)
                     else:
@@ -212,6 +215,7 @@ class ParagraphFinder:
                         current_paragraph = PdfParagraph(
                             pdf_paragraph_composition=[line],
                             debug_id=generate_base58_id(),
+                            layout_label=current_layout.name,
                         )
                         self.update_paragraph_data(current_paragraph)
                         paragraphs.append(current_paragraph)
@@ -228,6 +232,7 @@ class ParagraphFinder:
                 current_paragraph = PdfParagraph(
                     pdf_paragraph_composition=[line],
                     debug_id=generate_base58_id(),
+                    layout_label=current_layout.name,
                 )
                 paragraphs.append(current_paragraph)
             else:
@@ -412,6 +417,7 @@ class ParagraphFinder:
                         ),
                         unicode="",
                         debug_id=generate_base58_id(),
+                        layout_label=prev_composition.layout_label,
                     )
                     # 更新原段落
                     paragraph.pdf_paragraph_composition = (
@@ -440,6 +446,7 @@ class ParagraphFinder:
                         ),
                         unicode="",
                         debug_id=generate_base58_id(),
+                        layout_label=prev_composition.layout_label,
                     )
                     # 更新原段落
                     paragraph.pdf_paragraph_composition = (