File size: 3,181 Bytes
a4251f4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | """Interpreter-Qwen3-1.7B-ReTok —— 最小可运行示例。**只要 torch。**
huggingface-cli download Ismantic/Interpreter-Qwen3-1.7B-ReTok --local-dir Interpreter-ReTok
pip install git+https://github.com/Ismantic/PieceTokenizer
cd Interpreter-ReTok && python example_load.py
本目录自带 `model.py` / `checkpoint.py` / `tokenizer.py` —— Qwen3 的 forward、
safetensors 读取、piece 分词都在里面,**不需要 transformers,也不需要
safetensors 库**。用 transformers 加载也可以(权重是标准 HF 命名),
但那时分词器仍然要用本目录的 `tokenizer.py`:piece 不是 HF 分词器,
`AutoTokenizer.from_pretrained` 加载不了。
批量推理见 `example_vllm.py`。
"""
import os
import sys
import torch
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from model import Qwen3ForCausalLM # noqa: E402 随包自带,只依赖 torch
from tokenizer import PieceTokenizerWrapper # noqa: E402
HERE = os.path.dirname(os.path.abspath(__file__))
# 特殊 token(权威来源是本目录的 token_mapping.json,这里抄一份便于对照):
#
# <unk> 0 <pad> 81899
# <s> 1 <user> 81900
# </s> 2 <assistant> 81901
# <system> 81902
#
# 对话格式是 `<s> <user> {prompt} <assistant> {response} </s>`,**不是 ChatML**。
PROMPT = {
"zh-en": "Translate the following text from Chinese to English.\n"
"Chinese: {src}\nEnglish:",
"en-zh": "Translate the following text from English to Chinese.\n"
"English: {src}\nChinese:",
}
@torch.no_grad()
def translate(model, tok, text, direction="zh-en", max_new_tokens=256):
"""贪心解码。训练和评测用的都是贪心,采样会掉分。"""
ids = tok.apply_chat_template(
[{"role": "user", "content": PROMPT[direction].format(src=text)}],
tokenize=True, add_generation_prompt=True,
)
cur = torch.tensor([ids], device=next(model.parameters()).device)
out = []
for _ in range(max_new_tokens):
nxt = int(model(cur)[0, -1].argmax())
if nxt == tok.eos_token_id:
break
out.append(nxt)
cur = torch.cat([cur, torch.tensor([[nxt]], device=cur.device)], dim=1)
return tok.decode(out, skip_special_tokens=True).strip()
def main():
dev = "cuda" if torch.cuda.is_available() else "cpu"
tok = PieceTokenizerWrapper(HERE)
model = Qwen3ForCausalLM.from_pretrained(HERE, device=dev,
dtype=torch.bfloat16)
print(f"vocab={tok.vocab_size} bos={tok.bos_token_id} "
f"eos={tok.eos_token_id} pad={tok.pad_token_id} device={dev}\n")
for text, d in [
("人工智能正在深刻改变我们的生活方式。", "zh-en"),
("机器翻译的基本任务是把一种语言转换成另一种语言。", "zh-en"),
("The quick brown fox jumps over the lazy dog.", "en-zh"),
("Machine translation has improved dramatically in recent years.", "en-zh"),
]:
print(f"[{d}] {text}")
print(f" → {translate(model, tok, text, d)}\n")
if __name__ == "__main__":
main()
|