Summer-0.5B-S0 / example_vllm.py
tf-bao's picture
Upload Summer-0.5B-S0
16cf0c6 verified
Raw History Blame Contribute Delete
1.62 kB
"""用 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 = [
"机器翻译的基本任务是",
"中文分词的目标是",
]
def encode(p):
return [tok.bos_token_id] + tok.encode(p, add_special_tokens=False) # 训练时每行都以 <bos> 开头,不加是分布外输入
outputs = llm.generate(
[TokensPrompt(prompt_token_ids=encode(p)) for p in prompts],
SamplingParams(temperature=0.0, max_tokens=64,
stop_token_ids=tok.stop_token_ids),
)
for prompt, out in zip(prompts, outputs):
print(f"\n{prompt}", end="")
print(tok.decode(list(out.outputs[0].token_ids)))