"""批量翻译 —— vLLM 后端。

    pip install vllm
    python example_vllm.py

**`vllm serve` 开箱用不了。** piece 不是 HF 分词器,vLLM 起不来它的
tokenizer,所以必须:

  1. `LLM(..., skip_tokenizer_init=True)` —— 不让 vLLM 自己加载分词器
  2. 提示词自己编码,用 `TokensPrompt(prompt_token_ids=...)` 喂进去
  3. 输出拿 `token_ids` 再用本目录的 `tokenizer.py` 解码

这也是训练时评测走的路径,数字可比。
"""
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from tokenizer import PieceTokenizerWrapper  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))

PROMPT = {
    "zh-en": "Translate the following text from Chinese to English.\n"
             "Chinese: {src}\nEnglish:",
    "en-zh": "Translate the following text from English to Chinese.\n"
             "English: {src}\nChinese:",
}


def main():
    from vllm import LLM, SamplingParams
    from vllm.inputs import TokensPrompt

    tok = PieceTokenizerWrapper(HERE)
    llm = LLM(model=HERE, skip_tokenizer_init=True, dtype="bfloat16",
              max_model_len=1024, trust_remote_code=True)

    batch = [("人工智能正在深刻改变我们的生活方式。", "zh-en"),
             ("机器翻译的基本任务是把一种语言转换成另一种语言。", "zh-en"),
             ("The quick brown fox jumps over the lazy dog.", "en-zh"),
             ("Machine translation has improved dramatically in recent years.",
              "en-zh")]

    prompts = [TokensPrompt(prompt_token_ids=tok.apply_chat_template(
        [{"role": "user", "content": PROMPT[d].format(src=s)}],
        tokenize=True, add_generation_prompt=True)) for s, d in batch]

    # 贪心;停在 </s>(id 2)。温度不为 0 会掉分,训练和评测用的都是贪心。
    params = SamplingParams(max_tokens=256, temperature=0.0,
                            stop_token_ids=[tok.eos_token_id])

    for (src, d), out in zip(batch, llm.generate(prompts, params)):
        ids = list(out.outputs[0].token_ids)
        if ids and ids[-1] == tok.eos_token_id:
            ids = ids[:-1]
        print(f"[{d}] {src}\n      → {tok.decode(ids, skip_special_tokens=True).strip()}\n")


if __name__ == "__main__":
    main()