"""Interpreter-Qwen3-1.7B-ReTok —— 最小可运行示例。**只要 torch。**

    huggingface-cli download Ismantic/Interpreter-Qwen3-1.7B-ReTok --local-dir Interpreter-ReTok
    pip install git+https://github.com/Ismantic/PieceTokenizer
    cd Interpreter-ReTok && python example_load.py

本目录自带 `model.py` / `checkpoint.py` / `tokenizer.py` —— Qwen3 的 forward、
safetensors 读取、piece 分词都在里面,**不需要 transformers,也不需要
safetensors 库**。用 transformers 加载也可以(权重是标准 HF 命名),
但那时分词器仍然要用本目录的 `tokenizer.py`:piece 不是 HF 分词器,
`AutoTokenizer.from_pretrained` 加载不了。

批量推理见 `example_vllm.py`。
"""
import os
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from model import Qwen3ForCausalLM        # noqa: E402  随包自带,只依赖 torch
from tokenizer import PieceTokenizerWrapper  # noqa: E402

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

# 特殊 token(权威来源是本目录的 token_mapping.json,这里抄一份便于对照):
#
#     <unk>  0      <pad>        81899
#     <s>    1      <user>       81900
#     </s>   2      <assistant>  81901
#                   <system>     81902
#
# 对话格式是 `<s> <user> {prompt} <assistant> {response} </s>`,**不是 ChatML**。
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:",
}


@torch.no_grad()
def translate(model, tok, text, direction="zh-en", max_new_tokens=256):
    """贪心解码。训练和评测用的都是贪心,采样会掉分。"""
    ids = tok.apply_chat_template(
        [{"role": "user", "content": PROMPT[direction].format(src=text)}],
        tokenize=True, add_generation_prompt=True,
    )
    cur = torch.tensor([ids], device=next(model.parameters()).device)
    out = []
    for _ in range(max_new_tokens):
        nxt = int(model(cur)[0, -1].argmax())
        if nxt == tok.eos_token_id:
            break
        out.append(nxt)
        cur = torch.cat([cur, torch.tensor([[nxt]], device=cur.device)], dim=1)
    return tok.decode(out, skip_special_tokens=True).strip()


def main():
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    tok = PieceTokenizerWrapper(HERE)
    model = Qwen3ForCausalLM.from_pretrained(HERE, device=dev,
                                             dtype=torch.bfloat16)
    print(f"vocab={tok.vocab_size}  bos={tok.bos_token_id} "
          f"eos={tok.eos_token_id} pad={tok.pad_token_id}  device={dev}\n")

    for text, d in [
        ("人工智能正在深刻改变我们的生活方式。", "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"),
    ]:
        print(f"[{d}] {text}")
        print(f"      → {translate(model, tok, text, d)}\n")


if __name__ == "__main__":
    main()