Upload hf_backend/translator_runner.py with huggingface_hub

hf_backend/translator_runner.py

@@ -0,0 +1,303 @@
+from __future__ import annotations
+
+import os
+import re
+from queue import Empty, Queue
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+from threading import Thread
+
+from hf_backend.config import AppConfig
+from hf_backend.ebook_convert import normalize_epub_for_translation
+
+
+SUPPORTED_EXTENSIONS = {".epub", ".txt", ".srt"}
+TQDM_PROGRESS_RE = re.compile(
+    r"(?P<percent>\d{1,3})%\|.*?\|\s*(?P<current>\d+)/(?P<total>\d+)\s*\[",
+)
+BACKMATTER_MARKER_RE = re.compile(r"BBM_BACKMATTER_SKIPPED::(?P<title>.+)")
+CHECKPOINT_MARKER_RE = re.compile(r"BBM_CHECKPOINT_SAVED::(?P<count>\d+)")
+
+
+def _key_args(config: AppConfig, model: str) -> list[str]:
+    if model == "claude" and config.claude_key:
+        return ["--claude_key", config.claude_key]
+    if model == "deepl" and config.deepl_key:
+        return ["--deepl_key", config.deepl_key]
+    if model == "caiyun" and config.caiyun_key:
+        return ["--caiyun_key", config.caiyun_key]
+    if model == "gemini" and config.gemini_key:
+        return ["--gemini_key", config.gemini_key]
+    if model == "customapi" and config.custom_api:
+        return ["--custom_api", config.custom_api]
+    if model in {"chatgptapi", "gpt4"} and config.openai_key:
+        return ["--openai_key", config.openai_key]
+    return []
+
+
+def extract_progress_update(line: str) -> dict[str, int] | None:
+    cleaned = re.sub(r"\x1b\[[0-9;]*[A-Za-z]", "", line).strip()
+    matched = TQDM_PROGRESS_RE.search(cleaned)
+    if not matched:
+        return None
+
+    percent = int(matched.group("percent"))
+    current = int(matched.group("current"))
+    total = int(matched.group("total"))
+    return {
+        "percent": percent,
+        "current": current,
+        "total": total,
+    }
+
+
+def extract_backmatter_event(line: str) -> dict[str, str | bool] | None:
+    cleaned = re.sub(r"\x1b\[[0-9;]*[A-Za-z]", "", line).strip()
+    matched = BACKMATTER_MARKER_RE.search(cleaned)
+    if not matched:
+        return None
+
+    title = matched.group("title").strip()
+    return {
+        "stage": "translating",
+        "message": f"检测到尾部附录章节，后续内容保留原文：{title}",
+        "backmatter_skipped": True,
+        "backmatter_title": title,
+    }
+
+
+def extract_checkpoint_event(line: str) -> dict[str, str | bool | int] | None:
+    cleaned = re.sub(r"\x1b\[[0-9;]*[A-Za-z]", "", line).strip()
+    matched = CHECKPOINT_MARKER_RE.search(cleaned)
+    if not matched:
+        return None
+
+    current = int(matched.group("count"))
+    return {
+        "stage": "translating",
+        "message": f"已保存断点：{current}",
+        "checkpoint_saved": True,
+        "checkpoint_progress_current": current,
+    }
+
+
+def _read_stream_chunks(stream, stream_name: str, queue: Queue) -> None:
+    if stream is None:
+        return
+
+    if hasattr(stream, "read"):
+        buffer = ""
+        while True:
+            chunk = stream.read(1)
+            if not chunk:
+                if buffer:
+                    queue.put((stream_name, buffer))
+                return
+            buffer += chunk
+            if chunk in {"\n", "\r"}:
+                queue.put((stream_name, buffer))
+                buffer = ""
+    else:
+        while True:
+            chunk = stream.readline()
+            if not chunk:
+                return
+            queue.put((stream_name, chunk))
+
+
+def _run_streaming_translation(
+    command: list[str],
+    *,
+    config: AppConfig,
+    progress_callback=None,
+) -> str:
+    env = os.environ.copy()
+    env["PYTHONIOENCODING"] = "utf-8"
+    env["PYTHONUNBUFFERED"] = "1"
+    process = subprocess.Popen(
+        command,
+        cwd=config.repo_root,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+        encoding="utf-8",
+        errors="replace",
+        bufsize=1,
+        env=env,
+    )
+
+    output_queue: Queue = Queue()
+    readers = [
+        Thread(
+            target=_read_stream_chunks,
+            args=(process.stdout, "stdout", output_queue),
+            daemon=True,
+        ),
+        Thread(
+            target=_read_stream_chunks,
+            args=(process.stderr, "stderr", output_queue),
+            daemon=True,
+        ),
+    ]
+    for reader in readers:
+        reader.start()
+
+    combined_parts: list[str] = []
+    last_progress_signature: tuple[int, int, int] | None = None
+
+    while True:
+        if process.poll() is not None and all(not reader.is_alive() for reader in readers) and output_queue.empty():
+            break
+        try:
+            _stream_name, chunk = output_queue.get(timeout=0.1)
+        except Empty:
+            continue
+
+        text = chunk.strip()
+        if text:
+            combined_parts.append(text)
+            backmatter_event = extract_backmatter_event(text)
+            if backmatter_event and progress_callback:
+                progress_callback(backmatter_event)
+            checkpoint_event = extract_checkpoint_event(text)
+            if checkpoint_event and progress_callback:
+                progress_callback(checkpoint_event)
+            progress = extract_progress_update(text)
+            if progress and progress_callback:
+                signature = (
+                    progress["percent"],
+                    progress["current"],
+                    progress["total"],
+                )
+                if signature != last_progress_signature:
+                    progress_callback(
+                        {
+                            "stage": "translating",
+                            "message": f"正在翻译：{progress['current']} / {progress['total']}",
+                            **progress,
+                        }
+                    )
+                    last_progress_signature = signature
+
+    for reader in readers:
+        reader.join(timeout=1)
+
+    if process.wait() != 0:
+        combined_output = "\n".join(combined_parts)
+        raise RuntimeError(
+            f"translation command failed with code {process.returncode}\n{combined_output[-3000:]}"
+        )
+
+    return "\n".join(combined_parts)
+
+
+def run_translation_job(
+    config: AppConfig,
+    job: dict,
+    source_path: Path,
+    *,
+    progress_callback=None,
+    checkpoint_source_path: Path | None = None,
+) -> tuple[Path, str]:
+    suffix = source_path.suffix.lower()
+    if suffix not in SUPPORTED_EXTENSIONS:
+        raise RuntimeError(f"unsupported file type: {suffix}")
+
+    work_dir = config.work_dir / job["job_id"]
+    work_dir.mkdir(parents=True, exist_ok=True)
+
+    local_source = work_dir / job["filename"]
+    shutil.copy2(source_path, local_source)
+
+    translation_source = local_source
+    if suffix == ".epub":
+        translation_source = normalize_epub_for_translation(config, local_source, work_dir)
+        if checkpoint_source_path is not None:
+            seeded_checkpoint = translation_source.parent / f".{translation_source.stem}.temp.bin"
+            shutil.copy2(checkpoint_source_path, seeded_checkpoint)
+
+    model = config.resolve_model(job.get("model"))
+    language = config.resolve_language(job.get("language"))
+    command = [
+        sys.executable,
+        str(config.make_book_script),
+        "--book_name",
+        str(translation_source),
+        "--model",
+        model,
+        "--language",
+        language,
+    ]
+
+    command.extend(_key_args(config, model))
+    if config.api_base:
+        command.extend(["--api_base", config.api_base])
+    if suffix == ".epub" and getattr(config, "epub_accumulated_num", 1) > 1:
+        command.extend(["--accumulated_num", str(config.epub_accumulated_num)])
+    if suffix == ".epub":
+        command.extend(
+            [
+                "--checkpoint_interval",
+                str(getattr(config, "epub_checkpoint_interval", 50)),
+            ]
+        )
+        if checkpoint_source_path is not None:
+            command.append("--resume")
+        if getattr(config, "epub_backmatter_skip_after_percent", 0) > 0:
+            command.extend(
+                [
+                    "--skip_backmatter_after_percent",
+                    str(config.epub_backmatter_skip_after_percent),
+                ]
+            )
+        if getattr(config, "epub_backmatter_titles", "").strip():
+            command.extend(
+                [
+                    "--skip_backmatter_titles",
+                    config.epub_backmatter_titles,
+                ]
+            )
+    if config.batch_size > 0 and suffix == ".txt":
+        command.extend(["--batch_size", str(config.batch_size)])
+
+    if progress_callback:
+        combined_output = _run_streaming_translation(
+            command,
+            config=config,
+            progress_callback=progress_callback,
+        )
+    else:
+        completed = subprocess.run(
+            command,
+            cwd=config.repo_root,
+            check=False,
+            capture_output=True,
+            text=True,
+            encoding="utf-8",
+            errors="replace",
+        )
+        combined_output = "\n".join(
+            part for part in [completed.stdout.strip(), completed.stderr.strip()] if part
+        )
+
+        if completed.returncode != 0:
+            raise RuntimeError(
+                f"translation command failed with code {completed.returncode}\n{combined_output[-3000:]}"
+            )
+
+    result_path = translation_source.with_name(
+        f"{translation_source.stem}_bilingual{translation_source.suffix}"
+    )
+    if not result_path.exists():
+        raise RuntimeError(f"expected translated file not found: {result_path.name}")
+
+    canonical_result_path = local_source.with_name(
+        f"{local_source.stem}_bilingual{local_source.suffix}"
+    )
+    if canonical_result_path != result_path:
+        shutil.copy2(result_path, canonical_result_path)
+        result_path = canonical_result_path
+
+    return result_path, combined_output