"""用 vLLM 跑这个模型。

    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,自己编码好 token id 用 TokensPrompt
喂进去,拿回来的 id 自己解码。

推论:`vllm serve` 那种 OpenAI 兼容服务**开箱用不了** —— 它要把文本转成
token,而它转不了。调用方得自己传 token id。
"""
import torch
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)

prompts = [
    "机器翻译的基本任务是",
    "中文分词的目标是",
]

outputs = llm.generate(
    [TokensPrompt(prompt_token_ids=tok.encode(p, add_special_tokens=False))
     for p in prompts],
    SamplingParams(temperature=0.0, max_tokens=64,
                   stop_token_ids=[tok.eos_token_id]),
)

for prompt, out in zip(prompts, outputs):
    print(f"\n{prompt}", end="")
    print(tok.decode(list(out.outputs[0].token_ids)))