Merge pull request #229 from funstory-ai/pdf2zh-v2-rc

feat(highlevel): track and report peak memory usage during translation

Author: awwaawwa

babeldoc/__init__.py

@@ -1 +1 @@
-__version__ = "0.3.4"
+__version__ = "0.3.5"

babeldoc/const.py

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

babeldoc/high_level.py

@@ -381,7 +381,7 @@ def do_translate(
         original_pdf_path = translation_config.input_file
         logger.info(f"start to translate: {original_pdf_path}")
         start_time = time.time()
-
+        peak_memory_usage = 0
         with MemoryMonitor() as memory_monitor:
             # Check if split translation is enabled
             if not translation_config.split_strategy:
@@ -504,6 +504,7 @@ def do_translate(
                         logger.info("start merge results")
                         result = merger.merge_results(results)
                         logger.info("finish merge results")
+            peak_memory_usage = memory_monitor.peak_memory_usage
 
         finish_time = time.time()
         result.total_seconds = finish_time - start_time
@@ -512,6 +513,7 @@ def do_translate(
             f"finish translate: {original_pdf_path}, cost: {finish_time - start_time} s",
         )
         result.original_pdf_path = translation_config.input_file
+        result.peak_memory_usage = peak_memory_usage
         pm.translate_done(result)
         return result
 

babeldoc/main.py

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

babeldoc/translation_config.py

@@ -256,6 +256,7 @@ class TranslateResult:
     dual_pdf_path: str | None
     no_watermark_mono_pdf_path: str | None
     no_watermark_dual_pdf_path: str | None
+    peak_memory_usage: int | None
 
     def __init__(self, mono_pdf_path: str | None, dual_pdf_path: str | None):
         self.mono_pdf_path = mono_pdf_path
@@ -299,6 +300,9 @@ def __str__(self):
                 f"\tNo-watermark Dual-language PDF: {self.no_watermark_dual_pdf_path}"
             )
 
+        if hasattr(self, "peak_memory_usage") and self.peak_memory_usage:
+            result.append(f"\tPeak memory usage: {self.peak_memory_usage} MB")
+
         if result:
             result.insert(0, "Translation results:")
 

pyproject.toml

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