BabelDOC 0.2.20__tar.gz → 0.2.22__tar.gz

{babeldoc-0.2.20 → babeldoc-0.2.22}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: BabelDOC
-Version: 0.2.20
+Version: 0.2.22
 Summary: Yet Another Document Translator
 Project-URL: Homepage, https://github.com/funstory-ai/BabelDOC
 Project-URL: Issues, https://github.com/funstory-ai/BabelDOC/issues
@@ -399,6 +399,7 @@ Everyone interacting in YADT and its sub-projects' codebases, issue trackers, ch
 - [pdfminer](https://github.com/pdfminer/pdfminer.six)
 - [PyMuPDF](https://github.com/pymupdf/PyMuPDF)
 - [Asynchronize](https://github.com/multimeric/Asynchronize/tree/master?tab=readme-ov-file)
+- [PriorityThreadPoolExecutor](https://github.com/oleglpts/PriorityThreadPoolExecutor)
 
 <h2 id="star_hist">Star History</h2>
 

{babeldoc-0.2.20 → babeldoc-0.2.22}/README.md

@@ -353,6 +353,7 @@ Everyone interacting in YADT and its sub-projects' codebases, issue trackers, ch
 - [pdfminer](https://github.com/pdfminer/pdfminer.six)
 - [PyMuPDF](https://github.com/pymupdf/PyMuPDF)
 - [Asynchronize](https://github.com/multimeric/Asynchronize/tree/master?tab=readme-ov-file)
+- [PriorityThreadPoolExecutor](https://github.com/oleglpts/PriorityThreadPoolExecutor)
 
 <h2 id="star_hist">Star History</h2>
 

babeldoc-0.2.22/babeldoc/__init__.py

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

{babeldoc-0.2.20 → babeldoc-0.2.22}/babeldoc/const.py

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

{babeldoc-0.2.20 → babeldoc-0.2.22}/babeldoc/document_il/midend/il_translator.py

@@ -4,6 +4,7 @@ import logging
 import re
 from pathlib import Path
 
+import tiktoken
 from tqdm import tqdm
 
 from babeldoc.document_il import Document
@@ -119,11 +120,20 @@ class ILTranslator:
         self,
         translate_engine: BaseTranslator,
         translation_config: TranslationConfig,
+        tokenizer=None,
     ):
         self.translate_engine = translate_engine
         self.translation_config = translation_config
         self.font_mapper = FontMapper(translation_config)
 
+        if tokenizer is None:
+            self.tokenizer = tiktoken.encoding_for_model("gpt-4o")
+        else:
+            self.tokenizer = tokenizer
+
+    def calc_token_count(self, text: str) -> int:
+        return len(self.tokenizer.encode(text))
+
     def translate(self, docs: Document):
         tracker = DocumentTranslateTracker()
         # count total paragraph
@@ -173,6 +183,7 @@ class ILTranslator:
                 tracker.new_paragraph(),
                 page_font_map,
                 page_xobj_font_map,
+                priority=1048576 - self.calc_token_count(paragraph.unicode),
             )
 
     class TranslateInput:

{babeldoc-0.2.20 → babeldoc-0.2.22}/babeldoc/document_il/midend/il_translator_llm_only.py

@@ -1,4 +1,3 @@
-import concurrent.futures
 import json
 import logging
 import re
@@ -17,10 +16,11 @@ 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.document_il.utils.priority_thread_pool_executor import (
+    PriorityThreadPoolExecutor,
+)
 from babeldoc.translation_config import TranslationConfig
 
-enc = tiktoken.encoding_for_model("gpt-4o")
-
 logger = logging.getLogger(__name__)
 
 
@@ -39,6 +39,7 @@ class ILTranslatorLLMOnly:
         self,
         translate_engine: BaseTranslator,
         translation_config: TranslationConfig,
+        tokenizer=None,
     ):
         self.translate_engine = translate_engine
         self.translation_config = translation_config
@@ -47,9 +48,15 @@ class ILTranslatorLLMOnly:
             translation_config.shared_context_cross_split_part
         )
 
+        if tokenizer is None:
+            self.tokenizer = tiktoken.encoding_for_model("gpt-4o")
+        else:
+            self.tokenizer = tokenizer
+
         self.il_translator = ILTranslator(
             translate_engine=translate_engine,
             translation_config=translation_config,
+            tokenizer=self.tokenizer,
         )
 
         try:
@@ -57,6 +64,9 @@ class ILTranslatorLLMOnly:
         except NotImplementedError as e:
             raise ValueError("LLM translator not supported") from e
 
+    def calc_token_count(self, text: str) -> int:
+        return len(self.tokenizer.encode(text))
+
     def find_title_paragraph(self, docs: Document) -> PdfParagraph | None:
         """Find the first paragraph with layout_label 'title' in the document.
 
@@ -103,10 +113,10 @@ class ILTranslatorLLMOnly:
             self.stage_name,
             total,
         ) as pbar:
-            with concurrent.futures.ThreadPoolExecutor(
+            with PriorityThreadPoolExecutor(
                 max_workers=self.translation_config.qps,
             ) as executor2:
-                with concurrent.futures.ThreadPoolExecutor(
+                with PriorityThreadPoolExecutor(
                     max_workers=self.translation_config.qps,
                 ) as executor:
                     for page in docs.page:
@@ -128,10 +138,10 @@ class ILTranslatorLLMOnly:
     def process_page(
         self,
         page: Page,
-        executor: concurrent.futures.ThreadPoolExecutor,
+        executor: PriorityThreadPoolExecutor,
         pbar: tqdm | None = None,
         tracker: PageTranslateTracker = None,
-        executor2: concurrent.futures.ThreadPoolExecutor | None = None,
+        executor2: PriorityThreadPoolExecutor | None = None,
     ):
         self.translation_config.raise_if_cancelled()
         page_font_map = {}
@@ -145,17 +155,17 @@ class ILTranslatorLLMOnly:
 
         paragraphs = []
 
-        total_unicode_counts = 0
+        total_token_count = 0
         for paragraph in page.pdf_paragraph:
             if paragraph.debug_id is None or paragraph.unicode is None:
                 continue
             # self.translate_paragraph(paragraph, pbar,tracker.new_paragraph(), page_font_map, page_xobj_font_map)
-            total_unicode_counts += len(paragraph.unicode)
+            total_token_count += self.calc_token_count(paragraph.unicode)
             paragraphs.append(paragraph)
             if paragraph.layout_label == "title":
                 self.shared_context_cross_split_part.recent_title_paragraph = paragraph
 
-            if total_unicode_counts > 1200 or len(paragraphs) > 4:
+            if total_token_count > 800 or len(paragraphs) > 5:
                 executor.submit(
                     self.translate_paragraph,
                     BatchParagraph(paragraphs, tracker),
@@ -165,9 +175,10 @@ class ILTranslatorLLMOnly:
                     self.translation_config.shared_context_cross_split_part.first_paragraph,
                     self.translation_config.shared_context_cross_split_part.recent_title_paragraph,
                     executor2,
+                    priority=1048576 - total_token_count,
                 )
                 paragraphs = []
-                total_unicode_counts = 0
+                total_token_count = 0
 
         if paragraphs:
             executor.submit(
@@ -179,6 +190,7 @@ class ILTranslatorLLMOnly:
                 self.translation_config.shared_context_cross_split_part.first_paragraph,
                 self.translation_config.shared_context_cross_split_part.recent_title_paragraph,
                 executor2,
+                priority=1048576 - total_token_count,
             )
 
     def translate_paragraph(
@@ -189,10 +201,11 @@ class ILTranslatorLLMOnly:
         xobj_font_map: dict[int, dict[str, PdfFont]] = None,
         title_paragraph: PdfParagraph | None = None,
         local_title_paragraph: PdfParagraph | None = None,
-        executor=None,
+        executor: PriorityThreadPoolExecutor | None = None,
     ):
         """Translate a paragraph using pre and post processing functions."""
         self.translation_config.raise_if_cancelled()
+        should_translate_paragraph = []
         try:
             inputs = []
 
@@ -205,6 +218,7 @@ class ILTranslatorLLMOnly:
                 if text is None:
                     pbar.advance(1)
                     continue
+                should_translate_paragraph.append(i)
                 inputs.append((text, translate_input, paragraph, tracker))
             if not inputs:
                 return
@@ -313,8 +327,8 @@ class ILTranslatorLLMOnly:
                     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))
+                    input_token_count = self.calc_token_count(input_unicode)
+                    output_token_count = self.calc_token_count(output_unicode)
 
                     if not (0.3 < output_token_count / input_token_count < 3):
                         logger.warning(
@@ -354,11 +368,16 @@ class ILTranslatorLLMOnly:
                             inputs[id_][3],
                             page_font_map,
                             xobj_font_map,
+                            priority=1048576 - self.calc_token_count(paragraph.unicode),
                         )
 
         except Exception as e:
             logger.warning(f"Error {e} during translation. try fallback")
-            for i in range(len(batch_paragraph.paragraphs)):
+            if not should_translate_paragraph:
+                should_translate_paragraph = list(
+                    range(len(batch_paragraph.paragraphs))
+                )
+            for i in should_translate_paragraph:
                 paragraph = batch_paragraph.paragraphs[i]
                 tracker = batch_paragraph.trackers[i]
                 if paragraph.debug_id is None:
@@ -370,6 +389,7 @@ class ILTranslatorLLMOnly:
                     tracker,
                     page_font_map,
                     xobj_font_map,
+                    priority=1048576 - self.calc_token_count(paragraph.unicode),
                 )
 
     def _clean_json_output(self, llm_output: str) -> str:

{babeldoc-0.2.20 → babeldoc-0.2.22}/babeldoc/document_il/translator/translator.py

@@ -13,6 +13,7 @@ from tenacity import stop_after_attempt
 from tenacity import wait_exponential
 
 from babeldoc.document_il.translator.cache import TranslationCache
+from babeldoc.document_il.utils.atomic_integer import AtomicInteger
 
 logger = logging.getLogger(__name__)
 
@@ -62,31 +63,6 @@ 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)

babeldoc-0.2.22/babeldoc/document_il/utils/atomic_integer.py

@@ -0,0 +1,26 @@
+import threading
+
+
+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

babeldoc-0.2.22/babeldoc/document_il/utils/priority_thread_pool_executor.py

@@ -0,0 +1,259 @@
+# thanks to:
+# https://github.com/oleglpts/PriorityThreadPoolExecutor/blob/master/PriorityThreadPoolExecutor/__init__.py
+# https://github.com/oleglpts/PriorityThreadPoolExecutor/issues/4
+
+import atexit
+import itertools
+import logging
+import queue
+import random
+import sys
+import threading
+import weakref
+from concurrent.futures import _base
+from concurrent.futures.thread import BrokenThreadPool
+from concurrent.futures.thread import ThreadPoolExecutor
+from concurrent.futures.thread import _python_exit
+from concurrent.futures.thread import _threads_queues
+from concurrent.futures.thread import _WorkItem
+from heapq import heappop
+from heapq import heappush
+
+logger = logging.getLogger(__name__)
+
+########################################################################################################################
+#                                                Global variables                                                      #
+########################################################################################################################
+
+NULL_ENTRY = (sys.maxsize, _WorkItem(None, None, (), {}))
+_shutdown = False
+
+########################################################################################################################
+#                                           Before system exit procedure                                               #
+########################################################################################################################
+
+
+def python_exit():
+    """
+
+    Cleanup before system exit
+
+    """
+    global _shutdown
+    _shutdown = True
+    items = list(_threads_queues.items())
+    for _t, q in items:
+        q.put(NULL_ENTRY)
+    for t, _q in items:
+        t.join()
+
+
+# change default cleanup
+
+
+atexit.unregister(_python_exit)
+atexit.register(python_exit)
+
+
+class PriorityQueue(queue.Queue):
+    """Variant of Queue that retrieves open entries in priority order (lowest first).
+
+    Entries are typically tuples of the form:  (priority number, data).
+    """
+
+    REMOVED = "<removed-task>"
+    DEFAULT_PRIORITY = 100
+
+    def _init(self, maxsize):
+        self.queue = []
+        self.entry_finder = {}
+        self.counter = itertools.count()
+
+    def _qsize(self):
+        return len(self.queue)
+
+    def _put(self, item):
+        # heappush(self.queue, item)
+        try:
+            if item[1] in self.entry_finder:
+                self.remove(item[1])
+            count = next(self.counter)
+            entry = [item[0], count, item[1]]
+            self.entry_finder[item[1]] = entry
+            heappush(self.queue, entry)
+        except TypeError:  # handle item==None
+            self._put((self.DEFAULT_PRIORITY, None))
+
+    def remove(self, task):
+        """
+        This simply replaces the data with the REMOVED value,
+        which will get cleared out once _get reaches it.
+        """
+        entry = self.entry_finder.pop(task)
+        entry[-1] = self.REMOVED
+
+    def _get(self):
+        while self.queue:
+            entry = heappop(self.queue)
+            if entry[2] is not self.REMOVED:
+                del self.entry_finder[entry[2]]
+                return entry
+        return None
+
+
+def _worker(executor_reference, work_queue, initializer, initargs):
+    if initializer is not None:
+        try:
+            initializer(*initargs)
+        except BaseException:
+            _base.LOGGER.critical("Exception in initializer:", exc_info=True)
+            executor = executor_reference()
+            if executor is not None:
+                executor._initializer_failed()
+            return
+    try:
+        while True:
+            work_item = work_queue.get(block=True)
+            try:
+                if work_item[2] is not None:
+                    work_item[2].run()
+                    # Delete references to object. See issue16284
+                    del work_item
+
+                    # attempt to increment idle count
+                    executor = executor_reference()
+                    if executor is not None:
+                        executor._idle_semaphore.release()
+                    del executor
+                    continue
+
+                executor = executor_reference()
+                # Exit if:
+                #   - The interpreter is shutting down OR
+                #   - The executor that owns the worker has been collected OR
+                #   - The executor that owns the worker has been shutdown.
+                if _shutdown or executor is None or executor._shutdown:
+                    # Flag the executor as shutting down as early as possible if it
+                    # is not gc-ed yet.
+                    if executor is not None:
+                        executor._shutdown = True
+                    # Notice other workers
+                    work_queue.put(None)
+                    return
+                del executor
+            finally:
+                work_queue.task_done()
+    except BaseException:
+        _base.LOGGER.critical("Exception in worker", exc_info=True)
+
+
+class PriorityThreadPoolExecutor(ThreadPoolExecutor):
+    """
+    Thread pool executor with priority queue (priorities must be different, lowest first)
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # change work queue type to queue.PriorityQueue
+        self._work_queue: PriorityQueue = PriorityQueue()
+
+    def submit(self, fn, *args, **kwargs):
+        """
+
+        Sending the function to the execution queue
+
+        :param fn: function being executed
+        :type fn: callable
+        :param args: function's positional arguments
+        :param kwargs: function's keywords arguments
+        :return: future instance
+        :rtype: _base.Future
+
+        Added keyword:
+
+        - priority (integer later sys.maxsize)
+
+        """
+        with self._shutdown_lock:
+            if self._broken:
+                raise BrokenThreadPool(self._broken)
+
+            if self._shutdown:
+                raise RuntimeError("cannot schedule new futures after shutdown")
+            if _shutdown:
+                raise RuntimeError(
+                    "cannot schedule new futures after interpreter shutdown"
+                )
+
+            priority = kwargs.get("priority", random.randint(0, sys.maxsize - 1))  # noqa: S311
+            if "priority" in kwargs:
+                del kwargs["priority"]
+
+            f = _base.Future()
+            w = _WorkItem(f, fn, args, kwargs)
+
+            self._work_queue.put((priority, w))
+            self._adjust_thread_count()
+            return f
+
+    def _adjust_thread_count(self):
+        # if idle threads are available, don't spin new threads
+        if self._idle_semaphore.acquire(timeout=0):
+            return
+
+        # When the executor gets lost, the weakref callback will wake up
+        # the worker threads.
+        def weakref_cb(_, q=self._work_queue):
+            q.put(None)
+
+        num_threads = len(self._threads)
+        if num_threads < self._max_workers:
+            thread_name = f"{self._thread_name_prefix or self}_{num_threads:d}"
+            t = threading.Thread(
+                name=thread_name,
+                target=_worker,
+                args=(
+                    weakref.ref(self, weakref_cb),
+                    self._work_queue,
+                    self._initializer,
+                    self._initargs,
+                ),
+            )
+            t.start()
+            self._threads.add(t)
+            _threads_queues[t] = self._work_queue
+
+    def shutdown(self, wait=True, *, cancel_futures=False):
+        logger.info("Shutting down executor %s", self._thread_name_prefix or self)
+        if wait:
+            logger.info(
+                "Waiting for all tasks done %s", self._thread_name_prefix or self
+            )
+            self._work_queue.join()
+            logger.info("All tasks done %s", self._thread_name_prefix or self)
+
+        with self._shutdown_lock:
+            self._shutdown = True
+            if cancel_futures:
+                # Drain all work items from the queue, and then cancel their
+                # associated futures.
+                while True:
+                    try:
+                        work_item = self._work_queue.get_nowait()
+                    except queue.Empty:
+                        break
+                    if work_item is not None:
+                        work_item.future.cancel()
+
+            # Send a wake-up to prevent threads calling
+            # _work_queue.get(block=True) from permanently blocking.
+            self._work_queue.put(None)
+        if wait:
+            logger.info(
+                "Waiting for all thread done %s", self._thread_name_prefix or self
+            )
+            for t in self._threads:
+                self._work_queue.put(None)
+                t.join()
+        logger.info("shutdown finish %s", self._thread_name_prefix or self)

{babeldoc-0.2.20 → babeldoc-0.2.22}/babeldoc/main.py

@@ -23,7 +23,7 @@ from babeldoc.translation_config import TranslationConfig
 from babeldoc.translation_config import WatermarkOutputMode
 
 logger = logging.getLogger(__name__)
-__version__ = "0.2.20"
+__version__ = "0.2.22"
 
 
 def create_parser():

{babeldoc-0.2.20 → babeldoc-0.2.22}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "BabelDOC"
-version = "0.2.20"
+version = "0.2.22"
 description = "Yet Another Document Translator"
 license = "AGPL-3.0"
 readme = "README.md"
@@ -148,7 +148,7 @@ pythonpath = [".", "src"]
 testpaths = ["tests"]
 
 [bumpver]
-current_version = "0.2.20"
+current_version = "0.2.22"
 version_pattern = "MAJOR.MINOR.PATCH[.PYTAGNUM]"
 
 [bumpver.file_patterns]

{babeldoc-0.2.20 → babeldoc-0.2.22}/uv.lock

@@ -48,7 +48,7 @@ wheels = [
 
 [[package]]
 name = "babeldoc"
-version = "0.2.20"
+version = "0.2.22"
 source = { editable = "." }
 dependencies = [
     { name = "bitstring" },

babeldoc-0.2.20/babeldoc/__init__.py

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