Summer-0.5B-Chat / example_vllm.py
tf-bao's picture
Upload Summer-0.5B-Chat
d05f7e5 verified
Raw History Blame Contribute Delete
2.13 kB
"""用 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": "……下一个问题……"})