BabelDOC 0.1.6__py3-none-any.whl

babeldoc/document_il/translator/cache.py

@@ -0,0 +1,145 @@
+import json
+from pathlib import Path
+
+from peewee import SQL
+from peewee import AutoField
+from peewee import CharField
+from peewee import Model
+from peewee import SqliteDatabase
+from peewee import TextField
+
+# we don't init the database here
+db = SqliteDatabase(None)
+
+
+class _TranslationCache(Model):
+    id = AutoField()
+    translate_engine = CharField(max_length=20)
+    translate_engine_params = TextField()
+    original_text = TextField()
+    translation = TextField()
+
+    class Meta:
+        database = db
+        constraints = [
+            SQL(
+                """
+            UNIQUE (
+                translate_engine,
+                translate_engine_params,
+                original_text
+                )
+            ON CONFLICT REPLACE
+            """,
+            ),
+        ]
+
+
+class TranslationCache:
+    @staticmethod
+    def _sort_dict_recursively(obj):
+        if isinstance(obj, dict):
+            return {
+                k: TranslationCache._sort_dict_recursively(v)
+                for k in sorted(obj.keys())
+                for v in [obj[k]]
+            }
+        elif isinstance(obj, list):
+            return [TranslationCache._sort_dict_recursively(item) for item in obj]
+        return obj
+
+    def __init__(self, translate_engine: str, translate_engine_params: dict = None):
+        self.translate_engine = translate_engine
+        self.replace_params(translate_engine_params)
+
+    # The program typically starts multi-threaded translation
+    # only after cache parameters are fully configured,
+    # so thread safety doesn't need to be considered here.
+    def replace_params(self, params: dict = None):
+        if params is None:
+            params = {}
+        self.params = params
+        params = self._sort_dict_recursively(params)
+        self.translate_engine_params = json.dumps(params)
+
+    def update_params(self, params: dict = None):
+        if params is None:
+            params = {}
+        self.params.update(params)
+        self.replace_params(self.params)
+
+    def add_params(self, k: str, v):
+        self.params[k] = v
+        self.replace_params(self.params)
+
+    # Since peewee and the underlying sqlite are thread-safe,
+    # get and set operations don't need locks.
+    def get(self, original_text: str) -> str | None:
+        result = _TranslationCache.get_or_none(
+            translate_engine=self.translate_engine,
+            translate_engine_params=self.translate_engine_params,
+            original_text=original_text,
+        )
+        return result.translation if result else None
+
+    def set(self, original_text: str, translation: str):
+        _TranslationCache.create(
+            translate_engine=self.translate_engine,
+            translate_engine_params=self.translate_engine_params,
+            original_text=original_text,
+            translation=translation,
+        )
+
+
+def init_db(remove_exists=False):
+    cache_folder = Path.home() / ".cache" / "pdf2zh"
+    cache_folder.mkdir(parents=True, exist_ok=True)
+    # The current version does not support database migration, so add the version number to the file name.
+    cache_db_path = cache_folder / "cache.v1.db"
+    if remove_exists and cache_db_path.exists():
+        cache_db_path.unlink()
+    db.init(
+        cache_db_path,
+        pragmas={
+            "journal_mode": "wal",
+            "busy_timeout": 1000,
+        },
+    )
+    db.create_tables([_TranslationCache], safe=True)
+
+
+def init_test_db():
+    import tempfile
+
+    temp_file = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
+    cache_db_path = temp_file.name
+    temp_file.close()
+
+    test_db = SqliteDatabase(
+        cache_db_path,
+        pragmas={
+            "journal_mode": "wal",
+            "busy_timeout": 1000,
+        },
+    )
+    test_db.bind([_TranslationCache], bind_refs=False, bind_backrefs=False)
+    test_db.connect()
+    test_db.create_tables([_TranslationCache], safe=True)
+    return test_db
+
+
+def clean_test_db(test_db):
+    test_db.drop_tables([_TranslationCache])
+    test_db.close()
+    db_path = Path(test_db.database)
+    if db_path.exists():
+        db_path.unlink()
+    wal_path = Path(str(db_path) + "-wal")
+    if wal_path.exists():
+        wal_path.unlink()
+    shm_path = Path(str(db_path) + "-shm")
+    if shm_path.exists():
+        shm_path.unlink()
+
+
+init_db()

babeldoc/document_il/translator/translator.py

@@ -0,0 +1,267 @@
+import html
+import logging
+import re
+import threading
+import time
+import unicodedata
+from abc import ABC
+
+import openai
+import requests
+
+from babeldoc.document_il.translator.cache import TranslationCache
+
+logger = logging.getLogger(__name__)
+
+
+def remove_control_characters(s):
+    return "".join(ch for ch in s if unicodedata.category(ch)[0] != "C")
+
+
+class RateLimiter:
+    def __init__(self, max_qps: int):
+        self.max_qps = max_qps
+        self.min_interval = 1.0 / max_qps
+        self.last_requests = []  # Track last N requests
+        self.window_size = max_qps  # Track requests in a sliding window
+        self.lock = threading.Lock()
+
+    def wait(self):
+        with self.lock:
+            now = time.time()
+
+            # Clean up old requests outside the 1-second window
+            while self.last_requests and now - self.last_requests[0] > 1.0:
+                self.last_requests.pop(0)
+
+            # If we have less than max_qps requests in the last second, allow immediately
+            if len(self.last_requests) < self.max_qps:
+                self.last_requests.append(now)
+                return
+
+            # Otherwise, wait until we can make the next request
+            next_time = self.last_requests[0] + 1.0
+            if next_time > now:
+                time.sleep(next_time - now)
+            self.last_requests.pop(0)
+            self.last_requests.append(next_time)
+
+    def set_max_qps(self, max_qps):
+        self.max_qps = max_qps
+        self.min_interval = 1.0 / max_qps
+        self.window_size = max_qps
+
+
+_translate_rate_limiter = RateLimiter(5)
+
+
+def set_translate_rate_limiter(max_qps):
+    _translate_rate_limiter.set_max_qps(max_qps)
+
+
+class BaseTranslator(ABC):
+    # Due to cache limitations, name should be within 20 characters.
+    # cache.py: translate_engine = CharField(max_length=20)
+    name = "base"
+    lang_map = {}
+
+    def __init__(self, lang_in, lang_out, ignore_cache):
+        self.ignore_cache = ignore_cache
+        lang_in = self.lang_map.get(lang_in.lower(), lang_in)
+        lang_out = self.lang_map.get(lang_out.lower(), lang_out)
+        self.lang_in = lang_in
+        self.lang_out = lang_out
+
+        self.cache = TranslationCache(
+            self.name,
+            {
+                "lang_in": lang_in,
+                "lang_out": lang_out,
+            },
+        )
+
+        self.translate_call_count = 0
+        self.translate_cache_call_count = 0
+
+    def __del__(self):
+        logger.info(f"{self.name} translate call count: {self.translate_call_count}")
+        logger.info(
+            f"{self.name} translate cache call count: {self.translate_cache_call_count}",
+        )
+
+    def add_cache_impact_parameters(self, k: str, v):
+        """
+        Add parameters that affect the translation quality to distinguish the translation effects under different parameters.
+        :param k: key
+        :param v: value
+        """
+        self.cache.add_params(k, v)
+
+    def translate(self, text, ignore_cache=False):
+        """
+        Translate the text, and the other part should call this method.
+        :param text: text to translate
+        :return: translated text
+        """
+        self.translate_call_count += 1
+        if not (self.ignore_cache or ignore_cache):
+            cache = self.cache.get(text)
+            if cache is not None:
+                self.translate_cache_call_count += 1
+                return cache
+        _translate_rate_limiter.wait()
+        translation = self.do_translate(text)
+        if not (self.ignore_cache or ignore_cache):
+            self.cache.set(text, translation)
+        return translation
+
+    def do_translate(self, text):
+        """
+        Actual translate text, override this method
+        :param text: text to translate
+        :return: translated text
+        """
+        logger.critical(
+            f"Do not call BaseTranslator.do_translate. "
+            f"Translator: {self}. "
+            f"Text: {text}. ",
+        )
+        raise NotImplementedError
+
+    def __str__(self):
+        return f"{self.name} {self.lang_in} {self.lang_out} {self.model}"
+
+    def get_rich_text_left_placeholder(self, placeholder_id: int):
+        return f"<b{placeholder_id}>"
+
+    def get_rich_text_right_placeholder(self, placeholder_id: int):
+        return f"</b{placeholder_id}>"
+
+    def get_formular_placeholder(self, placeholder_id: int):
+        return self.get_rich_text_left_placeholder(placeholder_id)
+
+
+class GoogleTranslator(BaseTranslator):
+    name = "google"
+    lang_map = {"zh": "zh-CN"}
+
+    def __init__(self, lang_in, lang_out, ignore_cache=False):
+        super().__init__(lang_in, lang_out, ignore_cache)
+        self.session = requests.Session()
+        self.endpoint = "http://translate.google.com/m"
+        self.headers = {
+            "User-Agent": "Mozilla/4.0 (compatible;MSIE 6.0;Windows NT 5.1;SV1;.NET CLR 1.1.4322;.NET CLR 2.0.50727;.NET CLR 3.0.04506.30)",
+        }
+
+    def do_translate(self, text):
+        text = text[:5000]  # google translate max length
+        response = self.session.get(
+            self.endpoint,
+            params={"tl": self.lang_out, "sl": self.lang_in, "q": text},
+            headers=self.headers,
+        )
+        re_result = re.findall(
+            r'(?s)class="(?:t0|result-container)">(.*?)<',
+            response.text,
+        )
+        if response.status_code == 400:
+            result = "IRREPARABLE TRANSLATION ERROR"
+        else:
+            response.raise_for_status()
+            result = html.unescape(re_result[0])
+        return remove_control_characters(result)
+
+
+class BingTranslator(BaseTranslator):
+    # https://github.com/immersive-translate/old-immersive-translate/blob/6df13da22664bea2f51efe5db64c63aca59c4e79/src/background/translationService.js
+    name = "bing"
+    lang_map = {"zh": "zh-Hans"}
+
+    def __init__(self, lang_in, lang_out, ignore_cache=False):
+        super().__init__(lang_in, lang_out, ignore_cache)
+        self.session = requests.Session()
+        self.endpoint = "https://www.bing.com/translator"
+        self.headers = {
+            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36 Edg/131.0.0.0",
+        }
+
+    def find_sid(self):
+        response = self.session.get(self.endpoint)
+        response.raise_for_status()
+        url = response.url[:-10]
+        ig = re.findall(r"\"ig\":\"(.*?)\"", response.text)[0]
+        iid = re.findall(r"data-iid=\"(.*?)\"", response.text)[-1]
+        key, token = re.findall(
+            r"params_AbusePreventionHelper\s=\s\[(.*?),\"(.*?)\",",
+            response.text,
+        )[0]
+        return url, ig, iid, key, token
+
+    def do_translate(self, text):
+        text = text[:1000]  # bing translate max length
+        url, ig, iid, key, token = self.find_sid()
+        response = self.session.post(
+            f"{url}ttranslatev3?IG={ig}&IID={iid}",
+            data={
+                "fromLang": self.lang_in,
+                "to": self.lang_out,
+                "text": text,
+                "token": token,
+                "key": key,
+            },
+            headers=self.headers,
+        )
+        response.raise_for_status()
+        return response.json()[0]["translations"][0]["text"]
+
+
+class OpenAITranslator(BaseTranslator):
+    # https://github.com/openai/openai-python
+    name = "openai"
+
+    def __init__(
+        self,
+        lang_in,
+        lang_out,
+        model,
+        base_url=None,
+        api_key=None,
+        ignore_cache=False,
+    ):
+        super().__init__(lang_in, lang_out, ignore_cache)
+        self.options = {"temperature": 0}  # 随机采样可能会打断公式标记
+        self.client = openai.OpenAI(base_url=base_url, api_key=api_key)
+        self.add_cache_impact_parameters("temperature", self.options["temperature"])
+        self.model = model
+        self.add_cache_impact_parameters("model", self.model)
+        self.add_cache_impact_parameters("prompt", self.prompt(""))
+
+    def do_translate(self, text) -> str:
+        response = self.client.chat.completions.create(
+            model=self.model,
+            **self.options,
+            messages=self.prompt(text),
+        )
+        return response.choices[0].message.content.strip()
+
+    def prompt(self, text):
+        return [
+            {
+                "role": "system",
+                "content": "You are a professional,authentic machine translation engine.",
+            },
+            {
+                "role": "user",
+                "content": f";; Treat next line as plain text input and translate it into {self.lang_out}, output translation ONLY. If translation is unnecessary (e.g. proper nouns, codes, {'{{1}}, etc. '}), return the original text. NO explanations. NO notes. Input:\n\n{text}",
+            },
+        ]
+
+    def get_formular_placeholder(self, placeholder_id: int):
+        return "{{v" + str(placeholder_id) + "}}"
+        return "{{" + str(placeholder_id) + "}}"
+
+    def get_rich_text_left_placeholder(self, placeholder_id: int):
+        return self.get_formular_placeholder(placeholder_id)
+
+    def get_rich_text_right_placeholder(self, placeholder_id: int):
+        return self.get_formular_placeholder(placeholder_id + 1)

babeldoc/document_il/utils/fontmap.py

@@ -0,0 +1,227 @@
+import functools
+import logging
+import re
+from pathlib import Path
+
+import pymupdf
+from babeldoc.const import get_cache_file_path
+from babeldoc.document_il import PdfFont
+from babeldoc.document_il import il_version_1
+from babeldoc.translation_config import TranslationConfig
+
+logger = logging.getLogger(__name__)
+
+
+class FontMapper:
+    stage_name = "添加字体"
+
+    def __init__(self, translation_config: TranslationConfig):
+        self.font_names = [
+            "source-han-serif-cn.ttf",
+            "SourceHanSansSC-Regular.ttf",
+            "source-han-serif-cn-bold.ttf",
+            "SourceHanSansSC-Bold.ttf",
+        ]
+        self.fonts = {
+            Path(file_name).name.split(".")[0].replace("-", "").lower(): pymupdf.Font(
+                fontfile=str(get_cache_file_path(file_name)),
+            )
+            for file_name in self.font_names
+        }
+        for k, v in self.fonts.items():
+            v.font_id = k
+        self.translation_config = translation_config
+        self.base_font_path = translation_config.font
+        self.fallback_font_path = get_cache_file_path("noto.ttf")
+        self.base_font = pymupdf.Font(fontfile=str(self.base_font_path))
+        self.fallback_font = pymupdf.Font(fontfile=str(self.fallback_font_path))
+
+        self.kai_font_path = get_cache_file_path("LXGWWenKai-Regular.ttf")
+        self.kai_font = pymupdf.Font(fontfile=str(self.kai_font_path))
+
+        self.base_font.font_id = "base"
+        self.fallback_font.font_id = "fallback"
+        self.kai_font.font_id = "kai"
+
+        # Set ascent and descent for base font
+        self.base_font.ascent_fontmap = 1151
+        self.base_font.descent_fontmap = -286
+
+        # Set ascent and descent for fallback font
+        self.fallback_font.ascent_fontmap = 1069
+        self.fallback_font.descent_fontmap = -293
+
+        # Set ascent and descent for kai font
+        self.kai_font.ascent_fontmap = 928
+        self.kai_font.descent_fontmap = -256
+
+        self.fontid2font = {f.font_id: f for f in self.fonts.values()}
+        self.fontid2font["base"] = self.base_font
+        self.fontid2font["fallback"] = self.fallback_font
+        self.fontid2font["kai"] = self.kai_font
+
+        # Set ascent and descent for other fonts
+        font_metrics = {
+            "sourcehanserifcn": (1151, 0),
+            "sourcehansansscregular": (1160, -288),
+            "sourcehanserifcnbold": (1151, -286),
+            "sourcehansansscbold": (1160, -288),
+        }
+
+        for font_id, (ascent, descent) in font_metrics.items():
+            if font_id in self.fontid2font:
+                self.fontid2font[font_id].ascent_fontmap = ascent
+                self.fontid2font[font_id].descent_fontmap = descent
+
+        for font in self.fontid2font.values():
+            font.char_lengths = functools.lru_cache(maxsize=10240, typed=True)(
+                font.char_lengths,
+            )
+
+    def has_char(self, char_unicode: str):
+        if len(char_unicode) != 1:
+            return False
+        current_char = ord(char_unicode)
+        for font in self.fonts.values():
+            if font.has_glyph(current_char):
+                return True
+        if self.base_font.has_glyph(current_char):
+            return True
+        if self.fallback_font.has_glyph(current_char):
+            return True
+        return False
+
+    def map(self, original_font: PdfFont, char_unicode: str):
+        current_char = ord(char_unicode)
+        if isinstance(original_font, pymupdf.Font):
+            bold = original_font.is_bold
+            italic = original_font.is_italic
+            monospaced = original_font.is_monospaced
+            serif = original_font.is_serif
+        elif isinstance(original_font, PdfFont):
+            bold = original_font.bold
+            italic = original_font.italic
+            monospaced = original_font.monospace
+            serif = original_font.serif
+        else:
+            logger.error(
+                f"Unknown font type: {type(original_font)}. "
+                f"Original font: {original_font}. "
+                f"Char unicode: {char_unicode}. ",
+            )
+            return None
+        if italic and self.kai_font.has_glyph(current_char):
+            return self.kai_font
+        for _k, font in self.fonts.items():
+            if not font.has_glyph(current_char):
+                continue
+            if bold != font.is_bold:
+                continue
+            # 不知道什么原因，思源黑体的 serif 属性为1，先workaround
+            if serif == 1 and "serif" not in font.font_id:
+                continue
+            if serif == 0 and "serif" in font.font_id:
+                continue
+            return font
+        if self.base_font.has_glyph(current_char):
+            return self.base_font
+
+        if self.fallback_font.has_glyph(current_char):
+            return self.fallback_font
+
+        logger.error(
+            f"Can't find font for {char_unicode}({current_char}). "
+            f"Original font: {original_font}. "
+            f"Char unicode: {char_unicode}. ",
+        )
+        return None
+
+    def add_font(self, doc_zh: pymupdf.Document, il: il_version_1.Document):
+        font_list = [
+            ("base", self.base_font_path),
+            ("fallback", self.fallback_font_path),
+            ("kai", self.kai_font_path),
+        ]
+        font_list.extend(
+            [
+                (
+                    Path(file_name).name.split(".")[0].replace("-", "").lower(),
+                    get_cache_file_path(file_name),
+                )
+                for file_name in self.font_names
+            ],
+        )
+        font_id = {}
+        xreflen = doc_zh.xref_length()
+        with self.translation_config.progress_monitor.stage_start(
+            self.stage_name,
+            xreflen - 1 + len(font_list) + len(il.page) + len(font_list),
+        ) as pbar:
+            for font in font_list:
+                font_id[font[0]] = doc_zh[0].insert_font(font[0], font[1])
+                pbar.advance(1)
+            for xref in range(1, xreflen):
+                pbar.advance(1)
+                for label in ["Resources/", ""]:  # 可能是基于 xobj 的 res
+                    try:  # xref 读写可能出错
+                        font_res = doc_zh.xref_get_key(xref, f"{label}Font")
+                        if font_res is None:
+                            continue
+                        target_key_prefix = f"{label}Font/"
+                        if font_res[0] == "xref":
+                            resource_xref_id = re.search(
+                                "(\\d+) 0 R",
+                                font_res[1],
+                            ).group(1)
+                            xref = int(resource_xref_id)
+                            font_res = ("dict", doc_zh.xref_object(xref))
+                            target_key_prefix = ""
+                        if font_res[0] == "dict":
+                            for font in font_list:
+                                target_key = f"{target_key_prefix}{font[0]}"
+                                font_exist = doc_zh.xref_get_key(xref, target_key)
+                                if font_exist[0] == "null":
+                                    doc_zh.xref_set_key(
+                                        xref,
+                                        target_key,
+                                        f"{font_id[font[0]]} 0 R",
+                                    )
+                    except Exception:
+                        pass
+
+            # Create PdfFont for each font
+            # 预先创建所有字体对象
+            pdf_fonts = []
+            for font_name, font_path in font_list:
+                font = pymupdf.Font(fontfile=str(font_path))
+                # Get descent_fontmap from fontid2font
+                descent_fontmap = None
+                if font_name in self.fontid2font:
+                    mupdf_font = self.fontid2font[font_name]
+                    if hasattr(mupdf_font, "descent_fontmap"):
+                        descent_fontmap = mupdf_font.descent_fontmap
+                    if hasattr(mupdf_font, "ascent_fontmap"):
+                        ascent_fontmap = mupdf_font.ascent_fontmap
+
+                pdf_fonts.append(
+                    il_version_1.PdfFont(
+                        name=font_name,
+                        xref_id=font_id[font_name],
+                        font_id=font_name,
+                        encoding_length=2,
+                        bold=font.is_bold,
+                        italic=font.is_italic,
+                        monospace=font.is_monospaced,
+                        serif=font.is_serif,
+                        descent=descent_fontmap,
+                        ascent=ascent_fontmap,
+                    ),
+                )
+                pbar.advance(1)
+
+            # 批量添加字体到页面和XObject
+            for page in il.page:
+                page.pdf_font.extend(pdf_fonts)
+                for xobj in page.pdf_xobject:
+                    xobj.pdf_font.extend(pdf_fonts)
+                pbar.advance(1)