scripts/translate.py

"""
交互式翻译推理脚本

使用方式:
    # 命令行交互翻译
    python scripts/translate.py --checkpoint checkpoints/best_model.pt

    # 翻译文件
    python scripts/translate.py --checkpoint checkpoints/best_model.pt --input input.txt --output output.txt

    # 启动 Gradio Web UI
    python scripts/translate.py --checkpoint checkpoints/best_model.pt --web
"""

import argparse
import sys
from pathlib import Path

import torch
import yaml

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))

from easytranslate.evaluation.evaluator import Evaluator


def parse_args():
    parser = argparse.ArgumentParser(description="EasyTranslate Inference")
    parser.add_argument("--config", type=str, default="configs/default_config.yaml")
    parser.add_argument("--checkpoint", type=str, required=True)
    parser.add_argument("--input", type=str, default=None, help="输入文件路径")
    parser.add_argument("--output", type=str, default=None, help="输出文件路径")
    parser.add_argument("--web", action="store_true", help="启动 Gradio Web UI")
    return parser.parse_args()


def interactive_translate(evaluator):
    """
    命令行交互翻译。

    TODO [Person D]:
    1. 循环读取用户输入
    2. 调用 evaluator.translate_single()
    3. 打印翻译结果
    4. 输入 'quit' 退出
    """
    print("\nInteractive Translation Mode (type 'quit' to exit)")
    print("-" * 40)

    while True:
        try:
            text = input("\n[EN] > ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nBye!")
            break

        if text.lower() in ("quit", "exit", "q"):
            print("Bye!")
            break

        if not text:
            continue

        translation = evaluator.translate_single(text)
        print(f"[ZH] > {translation}")


def translate_file(evaluator, input_path: str, output_path: str):
    """
    文件翻译。

    TODO [Person D]:
    1. 读取输入文件 (一行一句)
    2. 批量翻译
    3. 将结果写入输出文件
    """
    input_file = Path(input_path)
    if not input_file.exists():
        print(f"Error: input file not found: {input_path}")
        return

    with open(input_file, "r", encoding="utf-8") as f:
        lines = [line.strip() for line in f if line.strip()]

    print(f"Translating {len(lines)} sentences...")

    # 分批翻译
    batch_size = 32
    translations = []
    for i in range(0, len(lines), batch_size):
        batch = lines[i : i + batch_size]
        batch_translations = evaluator.translate(batch)
        translations.extend(batch_translations)
        print(f"  Translated {min(i + batch_size, len(lines))}/{len(lines)}")

    # 写入输出文件
    out_file = Path(output_path)
    out_file.parent.mkdir(parents=True, exist_ok=True)
    with open(out_file, "w", encoding="utf-8") as f:
        for t in translations:
            f.write(t + "\n")

    print(f"Results saved to: {output_path}")


def launch_web_ui(evaluator):
    """
    启动 Gradio Web UI。

    TODO [Person D]:
    1. 创建 Gradio Interface
    2. 输入: 英文文本框
    3. 输出: 中文翻译结果
    4. 调用 evaluator.translate_single()
    """
    try:
        import gradio as gr
    except ImportError:
        print("Error: gradio is not installed. Install with: pip install gradio")
        return

    def translate_fn(text):
        if not text.strip():
            return ""
        return evaluator.translate_single(text)

    interface = gr.Interface(
        fn=translate_fn,
        inputs=gr.Textbox(label="English", placeholder="Enter English text..."),
        outputs=gr.Textbox(label="Chinese Translation"),
        title="EasyTranslate - English to Chinese",
        description="Transformer-based English to Chinese translation system.",
    )
    interface.launch()


def main():
    args = parse_args()

    print("=" * 60)
    print("  EasyTranslate - Translation")
    print("=" * 60)

    # 加载配置
    with open(args.config, "r", encoding="utf-8") as f:
        config = yaml.safe_load(f)

    # 加载模型
    checkpoint = torch.load(args.checkpoint, map_location="cpu")

    from easytranslate.model.transformer import Transformer
    from easytranslate.data.tokenizer import build_tokenizer

    model_config = config.get("model", {})
    model = Transformer(model_config)
    model.load_state_dict(checkpoint["model_state_dict"])

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    model.eval()

    tokenizer = build_tokenizer(config.get("data", {}).get("tokenizer", {}))

    # 构建 evaluator
    evaluator = Evaluator(model=model, tokenizer=tokenizer, config=config)

    # 根据参数选择模式
    if args.web:
        launch_web_ui(evaluator)
    elif args.input:
        output_path = args.output or args.input.replace(".txt", "_translated.txt")
        translate_file(evaluator, args.input, output_path)
    else:
        interactive_translate(evaluator)


if __name__ == "__main__":
    main()