Qwen3-1.7B-Base-ReTok / example_load.py
tf-bao's picture
Update inference code and model card
0c9fe61 verified
Raw History Blame Contribute Delete
1.03 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.encode(prompt, add_special_tokens=False)
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 == tok.eos_token_id:
break
out.append(nxt)
x = torch.cat([x, torch.tensor([[nxt]], device=x.device)], dim=1)
print(prompt + tok.decode(out))