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

    # 启动 Streamlit Web UI（无需模型即可预览界面）
    python scripts/translate.py --web
    # 或者带模型启动翻译
    python scripts/translate.py --checkpoint checkpoints/best_model.pt --web
"""

import argparse
import subprocess
import sys
from pathlib import Path

try:
    from omegaconf import OmegaConf
except ImportError:  # pragma: no cover
    OmegaConf = None
    import yaml

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

import torch

from easytranslate.data.collator import TranslationCollator
from easytranslate.data.dataset import (
    TranslationDataset,
    load_custom_dataset,
    load_opus_dataset,
    load_wmt_dataset,
)
from easytranslate.data.tokenizer import TokenizerWrapper, build_tokenizer
from easytranslate.evaluation.evaluator import Evaluator
from easytranslate.model import TransformerTranslationModel
from easytranslate.model.finetune import load_pretrained_model


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, default=None)
    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="启动 Streamlit Web UI")
    parser.add_argument("--streamlit-app", action="store_true", help=argparse.SUPPRESS)
    args, unknown = parser.parse_known_args()
    return args, unknown


def _get_config(config, *keys, default=None):
    value = config
    for key in keys:
        if isinstance(value, dict):
            value = value.get(key, default)
        else:
            value = getattr(value, key, default)
        if value is default:
            break
    return value


def _load_config(path, cli_overrides=None):
    if OmegaConf is not None:
        config = OmegaConf.load(path)
        if cli_overrides:
            config = OmegaConf.merge(config, OmegaConf.from_cli(cli_overrides))
        return config

    with open(path, "r", encoding="utf-8") as fin:
        config = yaml.safe_load(fin)
    if cli_overrides:
        print("Warning: OmegaConf is not installed; CLI overrides are ignored.")
    return config


def interactive_translate(evaluator):
    print("输入英⽂句⼦，按回车翻译；输入 'quit' 退出。")
    while True:
        try:
            text = input("> ").strip()
        except EOFError:
            break
        if not text:
            continue
        if text.lower() in {"quit", "exit"}:
            break
        translation = evaluator.translate_single(text)
        print(translation)


def translate_file(evaluator, input_path: str, output_path: str):
    source_lines = []
    with Path(input_path).open("r", encoding="utf-8") as fin:
        for line in fin:
            line = line.strip()
            if line:
                source_lines.append(line)

    translations = evaluator.translate(source_lines)

    output_file = Path(output_path)
    output_file.parent.mkdir(parents=True, exist_ok=True)
    with output_file.open("w", encoding="utf-8") as fout:
        for line in translations:
            fout.write(f"{line}\n")

    print(f"Translation complete: {len(translations)} lines written to {output_path}")


def launch_streamlit_app(args):
    import streamlit as st

    @st.cache_resource
    def load_evaluator():
        config = _load_config(args.config)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model, tokenizer = build_model_and_tokenizer(config, device)
        model = load_checkpoint(model, args.checkpoint, device)
        return Evaluator(model, tokenizer, config)

    st.set_page_config(page_title="EasyTranslate", layout="wide")
    st.title("EasyTranslate")
    st.write("English to Chinese translation powered by EasyTranslate.")

    if args.checkpoint is None:
        st.warning(
            "当前未提供模型 checkpoint，页面仅用于预览界面效果。"
            " 如需翻译，请传入 --checkpoint 或先训练生成模型。"
        )

    input_text = st.text_area("English Input", value="", height=200)
    if st.button("Translate"):
        if not input_text.strip():
            st.warning("请输入要翻译的英文文本。")
        elif args.checkpoint is None:
            st.error("未提供 checkpoint，无法执行翻译。请使用 --checkpoint 参数启动。")
        else:
            with st.spinner("Translating..."):
                try:
                    evaluator = load_evaluator()
                    translation = evaluator.translate_single(input_text.strip())
                    st.text_area("Chinese Translation", value=translation, height=200)
                except Exception as exc:
                    st.error(f"模型加载或翻译失败：{exc}")

    st.markdown("---")
    st.caption("此页面用于展示前端界面；在未提供模型时，翻译功能会被禁用。")


def launch_streamlit_process(args):
    script_path = Path(__file__).resolve()
    cmd = [sys.executable, "-m", "streamlit", "run", str(script_path), "--", "--streamlit-app", "--config", args.config]
    if args.checkpoint:
        cmd.extend(["--checkpoint", args.checkpoint])
    if args.input:
        cmd.extend(["--input", args.input])
    if args.output:
        cmd.extend(["--output", args.output])
    subprocess.run(cmd, check=True)


def _get_tokenizer_train_texts(config, allow_auto: bool = False) -> list[str] | None:
    dataset_name = _get_config(config, "data", "dataset_name")
    if dataset_name == "custom":
        custom = _get_config(config, "data", "custom") or {}
        data = load_custom_dataset(
            train_src=custom.get("train_src"),
            train_tgt=custom.get("train_tgt"),
            val_src=custom.get("val_src"),
            val_tgt=custom.get("val_tgt"),
            test_src=custom.get("test_src"),
            test_tgt=custom.get("test_tgt"),
            preprocessing_config=_get_config(config, "data", "preprocessing"),
        )
        return list(data["train"]["src"]) + list(data["train"]["tgt"])

    if not allow_auto:
        return None

    if dataset_name == "wmt":
        try:
            dataset = load_wmt_dataset(
                year=_get_config(config, "data", "wmt", "year"),
                language_pair=_get_config(config, "data", "wmt", "language_pair"),
                split="train",
            )
        except Exception:
            dataset = load_wmt_dataset(
                year=_get_config(config, "data", "wmt", "year"),
                language_pair=_get_config(config, "data", "wmt", "language_pair"),
                split="validation",
            )
        return list(dataset["src"]) + list(dataset["tgt"])

    if dataset_name == "opus":
        try:
            dataset = load_opus_dataset(
                subset=_get_config(config, "data", "opus", "subset"),
                split="train",
            )
        except Exception:
            dataset = load_opus_dataset(
                subset=_get_config(config, "data", "opus", "subset"),
                split="validation",
            )
        return list(dataset["src"]) + list(dataset["tgt"])

    return None


def build_model_and_tokenizer(config, device):
    model_type = _get_config(config, "model", "type")
    if model_type == "transformer_scratch":
        tokenizer_config = _get_config(config, "tokenizer") or {}
        tokenizer_path = tokenizer_config.get("path") or tokenizer_config.get("tokenizer_path")
        tokenizer_type = tokenizer_config.get("type", "bpe")
        auto_train = bool(tokenizer_config.get("auto_train", False))
        if tokenizer_type in {"bpe", "sentencepiece"} and not tokenizer_path:
            train_texts = _get_tokenizer_train_texts(config, allow_auto=auto_train)
            if train_texts is None:
                raise ValueError(
                    "BPE tokenizer requires tokenizer.path or a local custom dataset with train texts. "
                    "Automatic download from WMT/OPUS is disabled by default. "
                    "Set tokenizer.auto_train=true to enable it, or provide tokenizer.path/pretrained tokenizer."
                )
            tokenizer = build_tokenizer(tokenizer_config, train_texts=train_texts)
        else:
            tokenizer = build_tokenizer(tokenizer_config)

        model = TransformerTranslationModel(
            src_vocab_size=tokenizer.vocab_size,
            tgt_vocab_size=tokenizer.vocab_size,
            d_model=_get_config(config, "model", "transformer", "d_model"),
            nhead=_get_config(config, "model", "transformer", "nhead"),
            num_encoder_layers=_get_config(config, "model", "transformer", "num_encoder_layers"),
            num_decoder_layers=_get_config(config, "model", "transformer", "num_decoder_layers"),
            dim_feedforward=_get_config(config, "model", "transformer", "dim_feedforward"),
            dropout=_get_config(config, "model", "transformer", "dropout"),
            activation=_get_config(config, "model", "transformer", "activation"),
            max_seq_len=_get_config(config, "model", "transformer", "max_seq_len"),
            use_flash_attention=_get_config(config, "model", "transformer", "use_flash_attention"),
            use_rotary_embedding=_get_config(config, "model", "transformer", "use_rotary_embedding"),
            pre_norm=_get_config(config, "model", "transformer", "pre_norm"),
            pad_id=tokenizer.pad_token_id,
        )
        return model.to(device), tokenizer

    model, hf_tokenizer = load_pretrained_model(
        config.model.pretrained.model_name,
        config.model.pretrained.src_lang,
        config.model.pretrained.tgt_lang,
        device=str(device),
    )
    tokenizer = TokenizerWrapper(
        hf_tokenizer,
        pad_token=getattr(hf_tokenizer, "pad_token", "<pad>"),
        unk_token=getattr(hf_tokenizer, "unk_token", "<unk>"),
        bos_token=getattr(hf_tokenizer, "bos_token", "<s>"),
        eos_token=getattr(hf_tokenizer, "eos_token", "</s>"),
    )
    return model, tokenizer


def load_checkpoint(model, checkpoint_path, device):
    checkpoint = torch.load(checkpoint_path, map_location=device)
    if isinstance(checkpoint, dict):
        if "model_state_dict" in checkpoint:
            model.load_state_dict(checkpoint["model_state_dict"])
        elif "state_dict" in checkpoint:
            model.load_state_dict(checkpoint["state_dict"])
        else:
            try:
                model.load_state_dict(checkpoint)
            except Exception as exc:
                raise ValueError("Checkpoint does not contain a valid model state dict") from exc
    else:
        raise ValueError("Unsupported checkpoint format")
    return model


def main():
    args, cli_overrides = parse_args()

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

    if args.web and not args.streamlit_app:
        launch_streamlit_process(args)
        return

    if args.streamlit_app:
        launch_streamlit_app(args)
        return

    config = _load_config(args.config, cli_overrides)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model, tokenizer = build_model_and_tokenizer(config, device)
    model = load_checkpoint(model, args.checkpoint, device)
    evaluator = Evaluator(model, tokenizer, config)

    if args.input and args.output:
        translate_file(evaluator, args.input, args.output)
    else:
        interactive_translate(evaluator)


if __name__ == "__main__":
    main()