"""用 vLLM 跑这个 chat 模型,支持多轮。

    pip install vllm
    pip install git+https://github.com/Ismantic/PieceTokenizer
    python example_vllm.py

## 为什么要 skip_tokenizer_init

vLLM 能加载**权重** —— config.json 里 architectures 是 Qwen3ForCausalLM,
它用自己那份 Qwen3 实现,按 state_dict 的 key 名灌进去。

但它**认不了这个词表**:81903 的 piece 词表不是 HF 标准格式,AutoTokenizer
走不通。所以要 skip_tokenizer_init=True,自己用 apply_chat_template 编码好
token id 用 TokensPrompt 喂进去,拿回来的 id 自己解码。

推论:`vllm serve` 那种 OpenAI 兼容服务**开箱用不了** —— 它要把文本转成
token,而它转不了。调用方得自己传 token id。

## repetition_penalty=1.15

**不是可选的美化项。** 纯贪心(temperature=0)容易陷入复读循环,这是这个
项目自己反复测出来的结论,不是随便调的参数——见仓库 docs/POSTTRAIN.md。
想要跑分可比的贪心结果,把 temperature 设成 0、repetition_penalty 设成 1.0。
"""
from vllm import LLM, SamplingParams
from vllm.inputs import TokensPrompt

from tokenizer import PieceTokenizerWrapper

HERE = "."

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

conversation = [
    {"role": "user", "content": "机器翻译的基本任务是什么？"},
]
ids = tok.apply_chat_template(conversation, tokenize=True,
                              add_generation_prompt=True)

outputs = llm.generate(
    [TokensPrompt(prompt_token_ids=ids)],
    SamplingParams(temperature=0.6, repetition_penalty=1.15, max_tokens=600,
                   stop_token_ids=tok.stop_token_ids),
)

reply = tok.decode(list(outputs[0].outputs[0].token_ids), skip_special_tokens=True)
print(reply)

# 多轮:把回复接回 conversation,再走一次 apply_chat_template。
# conversation.append({"role": "assistant", "content": reply})
# conversation.append({"role": "user", "content": "……下一个问题……"})