main.py

import asyncio
import logging
from pathlib import Path

import configargparse
import tqdm
from rich.progress import BarColumn
from rich.progress import MofNCompleteColumn
from rich.progress import Progress
from rich.progress import TextColumn
from rich.progress import TimeElapsedColumn
from rich.progress import TimeRemainingColumn

import babeldoc.high_level
from babeldoc.const import get_cache_file_path
from babeldoc.document_il.translator.translator import BingTranslator
from babeldoc.document_il.translator.translator import GoogleTranslator
from babeldoc.document_il.translator.translator import OpenAITranslator
from babeldoc.document_il.translator.translator import set_translate_rate_limiter
from babeldoc.translation_config import TranslationConfig

logger = logging.getLogger(__name__)
__version__ = "0.1.11"


def create_parser():
    parser = configargparse.ArgParser(
        config_file_parser_class=configargparse.TomlConfigParser(["babeldoc"]),
    )
    parser.add_argument(
        "-c",
        "--my-config",
        required=False,
        is_config_file=True,
        help="config file path",
    )
    parser.add_argument(
        "--version",
        action="version",
        version=f"%(prog)s {__version__}",
    )
    parser.add_argument(
        "--files",
        type=str,
        # nargs="*",
        action="append",
        help="One or more paths to PDF files.",
    )
    parser.add_argument(
        "--debug",
        "-d",
        default=False,
        action="store_true",
        help="Use debug logging level.",
    )
    parser.add_argument(
        "--warmup",
        default=False,
        action="store_true",
        help="Only download and verify required assets then exit.",
    )
    parser.add_argument(
        "--rpc-doclayout",
        default=None,
        help="RPC service host address for document layout analysis",
        type=str,
    )
    translation_params = parser.add_argument_group(
        "Translation",
        description="Used during translation",
    )
    translation_params.add_argument(
        "--pages",
        "-p",
        type=str,
        help="Pages to translate. If not set, translate all pages. like: 1,2,1-,-3,3-5",
    )
    translation_params.add_argument(
        "--min-text-length",
        type=int,
        default=5,
        help="Minimum text length to translate (default: 5)",
    )
    translation_params.add_argument(
        "--lang-in",
        "-li",
        type=str,
        default="en",
        help="The code of source language.",
    )
    translation_params.add_argument(
        "--lang-out",
        "-lo",
        type=str,
        default="zh",
        help="The code of target language.",
    )
    translation_params.add_argument(
        "--output",
        "-o",
        type=str,
        default=None,
        help="Output directory for files. if not set, use same as input.",
    )
    translation_params.add_argument(
        "--qps",
        "-q",
        type=int,
        default=4,
        help="QPS limit of translation service",
    )
    translation_params.add_argument(
        "--ignore-cache",
        "-ic",
        default=False,
        action="store_true",
        help="Ignore translation cache.",
    )
    translation_params.add_argument(
        "--no-dual",
        default=False,
        action="store_true",
        help="Do not output bilingual PDF files",
    )
    translation_params.add_argument(
        "--no-mono",
        default=False,
        action="store_true",
        help="Do not output monolingual PDF files",
    )
    translation_params.add_argument(
        "--formular-font-pattern",
        type=str,
        default=None,
        help="Font pattern to identify formula text",
    )
    translation_params.add_argument(
        "--formular-char-pattern",
        type=str,
        default=None,
        help="Character pattern to identify formula text",
    )
    translation_params.add_argument(
        "--split-short-lines",
        default=False,
        action="store_true",
        help="Force split short lines into different paragraphs (may cause poor typesetting & bugs)",
    )
    translation_params.add_argument(
        "--short-line-split-factor",
        type=float,
        default=0.8,
        help="Split threshold factor. The actual threshold is the median length of all lines on the current page * this factor",
    )
    translation_params.add_argument(
        "--skip-clean",
        default=False,
        action="store_true",
        help="Skip PDF cleaning step",
    )
    translation_params.add_argument(
        "--dual-translate-first",
        default=False,
        action="store_true",
        help="Put translated pages first in dual PDF mode",
    )
    translation_params.add_argument(
        "--disable-rich-text-translate",
        default=False,
        action="store_true",
        help="Disable rich text translation (may help improve compatibility with some PDFs)",
    )
    translation_params.add_argument(
        "--enhance-compatibility",
        default=False,
        action="store_true",
        help="Enable all compatibility enhancement options (equivalent to --skip-clean --dual-translate-first --disable-rich-text-translate)",
    )
    translation_params.add_argument(
        "--report-interval",
        type=float,
        default=0.1,
        help="Progress report interval in seconds (default: 0.1)",
    )
    service_params = translation_params.add_mutually_exclusive_group()
    service_params.add_argument(
        "--openai",
        default=False,
        action="store_true",
        help="Use OpenAI translator.",
    )
    service_params.add_argument(
        "--google",
        default=False,
        action="store_true",
        help="Use Google translator.",
    )
    service_params.add_argument(
        "--bing",
        default=False,
        action="store_true",
        help="Use Bing translator.",
    )
    openai_params = parser.add_argument_group(
        "Translation - OpenAI Options",
        description="OpenAI specific options",
    )
    openai_params.add_argument(
        "--openai-model",
        "-m",
        type=str,
        default="gpt-4o-mini",
        help="The OpenAI model to use for translation.",
    )
    openai_params.add_argument(
        "--openai-base-url",
        "-b",
        type=str,
        default=None,
        help="The base URL for the OpenAI API.",
    )
    openai_params.add_argument(
        "--openai-api-key",
        "-k",
        type=str,
        default=None,
        help="The API key for the OpenAI API.",
    )

    return parser


def create_progress_handler(translation_config: TranslationConfig):
    """Create a progress handler function based on the configuration.

    Args:
        translation_config: The translation configuration.

    Returns:
        A tuple of (progress_context, progress_handler), where progress_context is a context
        manager that should be used to wrap the translation process, and progress_handler
        is a function that will be called with progress events.
    """
    if translation_config.use_rich_pbar:
        progress = Progress(
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            MofNCompleteColumn(),
            TimeElapsedColumn(),
            TimeRemainingColumn(),
        )
        translate_task_id = progress.add_task("translate", total=100)
        stage_tasks = {}

        def progress_handler(event):
            if event["type"] == "progress_start":
                stage_tasks[event["stage"]] = progress.add_task(
                    f"{event['stage']}",
                    total=event.get("stage_total", 100),
                )
            elif event["type"] == "progress_update":
                stage = event["stage"]
                if stage in stage_tasks:
                    progress.update(
                        stage_tasks[stage],
                        completed=event["stage_current"],
                        total=event["stage_total"],
                        description=f"{event['stage']} ({event['stage_current']}/{event['stage_total']})",
                        refresh=True,
                    )
                progress.update(
                    translate_task_id,
                    completed=event["overall_progress"],
                    refresh=True,
                )
            elif event["type"] == "progress_end":
                stage = event["stage"]
                if stage in stage_tasks:
                    progress.update(
                        stage_tasks[stage],
                        completed=event["stage_total"],
                        total=event["stage_total"],
                        description=f"{event['stage']} (Complete)",
                        refresh=True,
                    )
                    progress.update(
                        translate_task_id,
                        completed=event["overall_progress"],
                        refresh=True,
                    )
                progress.refresh()

        return progress, progress_handler
    else:
        pbar = tqdm.tqdm(total=100, desc="translate")

        def progress_handler(event):
            if event["type"] == "progress_update":
                pbar.update(event["overall_progress"] - pbar.n)
                pbar.set_description(
                    f"{event['stage']} ({event['stage_current']}/{event['stage_total']})",
                )
            elif event["type"] == "progress_end":
                pbar.set_description(f"{event['stage']} (Complete)")
                pbar.refresh()

        return pbar, progress_handler


async def main():
    parser = create_parser()
    args = parser.parse_args()

    if args.debug:
        logging.getLogger().setLevel(logging.DEBUG)

    if args.warmup:
        logger.info("Warmup completed, exiting...")
        return

    # 验证翻译服务选择
    if not (args.openai or args.google or args.bing):
        parser.error("必须选择一个翻译服务：--openai、--google 或 --bing")

    # 验证 OpenAI 参数
    if args.openai and not args.openai_api_key:
        parser.error("使用 OpenAI 服务时必须提供 API key")

    # 实例化翻译器
    if args.openai:
        translator = OpenAITranslator(
            lang_in=args.lang_in,
            lang_out=args.lang_out,
            model=args.openai_model,
            base_url=args.openai_base_url,
            api_key=args.openai_api_key,
            ignore_cache=args.ignore_cache,
        )
    elif args.bing:
        translator = BingTranslator(
            lang_in=args.lang_in,
            lang_out=args.lang_out,
            ignore_cache=args.ignore_cache,
        )
    else:
        translator = GoogleTranslator(
            lang_in=args.lang_in,
            lang_out=args.lang_out,
            ignore_cache=args.ignore_cache,
        )

    # 设置翻译速率限制
    set_translate_rate_limiter(args.qps)

    # 初始化文档布局模型
    if args.rpc_doclayout:
        from babeldoc.docvision.rpc_doclayout import RpcDocLayoutModel

        doc_layout_model = RpcDocLayoutModel(host=args.rpc_doclayout)
    else:
        from babeldoc.docvision.doclayout import DocLayoutModel

        doc_layout_model = DocLayoutModel.load_onnx()

    pending_files = []
    for file in args.files:
        # 清理文件路径，去除两端的引号
        if file.startswith("--files="):
            file = file[len("--files=") :]
        file = file.lstrip("-").strip("\"'")
        if not Path(file).exists():
            logger.error(f"文件不存在：{file}")
            exit(1)
        if not file.endswith(".pdf"):
            logger.error(f"文件不是 PDF 文件：{file}")
            exit(1)
        pending_files.append(file)

    font_path = get_cache_file_path("source-han-serif-cn.ttf")

    # 验证字体
    if font_path:
        if not Path(font_path).exists():
            logger.error(f"字体文件不存在：{font_path}")
            exit(1)
        if not str(font_path).endswith(".ttf"):
            logger.error(f"字体文件不是 TTF 文件：{font_path}")
            exit(1)

    if args.output:
        if not Path(args.output).exists():
            logger.info(f"输出目录不存在，创建：{args.output}")
            try:
                Path(args.output).mkdir(parents=True, exist_ok=True)
            except OSError:
                logger.critical(
                    f"Failed to create output folder at {args.output}",
                    exc_info=True,
                )
                exit(1)
    else:
        args.output = None

    for file in pending_files:
        # 清理文件路径，去除两端的引号
        file = file.strip("\"'")
        # 创建配置对象
        config = TranslationConfig(
            input_file=file,
            font=font_path,
            pages=args.pages,
            output_dir=args.output,
            translator=translator,
            debug=args.debug,
            lang_in=args.lang_in,
            lang_out=args.lang_out,
            no_dual=args.no_dual,
            no_mono=args.no_mono,
            qps=args.qps,
            formular_font_pattern=args.formular_font_pattern,
            formular_char_pattern=args.formular_char_pattern,
            split_short_lines=args.split_short_lines,
            short_line_split_factor=args.short_line_split_factor,
            doc_layout_model=doc_layout_model,
            skip_clean=args.skip_clean,
            dual_translate_first=args.dual_translate_first,
            disable_rich_text_translate=args.disable_rich_text_translate,
            enhance_compatibility=args.enhance_compatibility,
            report_interval=args.report_interval,
            min_text_length=args.min_text_length,
        )

        # Create progress handler
        progress_context, progress_handler = create_progress_handler(config)

        # 开始翻译
        with progress_context:
            async for event in babeldoc.high_level.async_translate(config):
                progress_handler(event)
                if config.debug:
                    logger.debug(event)
                if event["type"] == "finish":
                    result = event["translate_result"]
                    logger.info("Translation Result:")
                    logger.info(f"  Original PDF: {result.original_pdf_path}")
                    logger.info(f"  Time Cost: {result.total_seconds:.2f}s")
                    logger.info(f"  Mono PDF: {result.mono_pdf_path or 'None'}")
                    logger.info(f"  Dual PDF: {result.dual_pdf_path or 'None'}")
                    break


# for backward compatibility
def create_cache_folder():
    return babeldoc.high_level.create_cache_folder()


# for backward compatibility
def download_font_assets():
    return babeldoc.high_level.download_font_assets()


def cli():
    """Command line interface entry point."""
    from rich.logging import RichHandler

    logging.basicConfig(level=logging.INFO, handlers=[RichHandler()])

    logging.getLogger("httpx").setLevel("CRITICAL")
    logging.getLogger("httpx").propagate = False
    logging.getLogger("openai").setLevel("CRITICAL")
    logging.getLogger("openai").propagate = False
    logging.getLogger("httpcore").setLevel("CRITICAL")
    logging.getLogger("httpcore").propagate = False
    logging.getLogger("http11").setLevel("CRITICAL")
    logging.getLogger("http11").propagate = False
    for v in logging.Logger.manager.loggerDict.values():
        if getattr(v, "name", None) is None:
            continue
        if (
            v.name.startswith("pdfminer")
            or v.name.startswith("peewee")
            or v.name.startswith("httpx")
            or "http11" in v.name
            or "openai" in v.name
        ):
            v.disabled = True
            v.propagate = False

    babeldoc.high_level.init()
    asyncio.run(main())


if __name__ == "__main__":
    cli()