Summer-0.5B-S0 / example_load.py
tf-bao's picture
Upload Summer-0.5B-S0
16cf0c6 verified
Raw History Blame Contribute Delete
1.17 kB
"""跑通这个发布包 —— 只需要 torch 和 piece_tokenizer。
pip install torch
pip install git+https://github.com/Ismantic/PieceTokenizer
python example_load.py
"""
import torch
from model import Qwen3ForCausalLM
from tokenizer import PieceTokenizerWrapper
HERE = "."
tok = PieceTokenizerWrapper(HERE)
model = Qwen3ForCausalLM.from_pretrained(
HERE, device="cuda" if torch.cuda.is_available() else "cpu",
dtype=torch.bfloat16)
prompt = "机器翻译的基本任务是"
ids = [tok.bos_token_id] + tok.encode(prompt, add_special_tokens=False) # 这个模型训练时每行都以 <bos> 开头,不加是分布外输入——会生成退化的重复内容
print(f"{prompt!r} -> {len(ids)} tokens")
# 贪心续写 40 步。没有 KV cache —— 每步重算前缀,短续写够用。
x = torch.tensor([ids], device=next(model.parameters()).device)
out = []
with torch.no_grad():
for _ in range(40):
nxt = int(model(x)[0, -1].argmax())
if nxt in tok.stop_token_ids:
break
out.append(nxt)
x = torch.cat([x, torch.tensor([[nxt]], device=x.device)], dim=1)
print(prompt + tok.decode(out))