分词与对话模板

import json
import sys

from tokenizers import Tokenizer
from transformers import AutoTokenizer

model_dir = sys.argv[1]  # 本地的 SmolLM2-135M-Instruct 权重
S = json.load(open("strings.json"))
tk = Tokenizer.from_file(f"{model_dir}/tokenizer.json")
auto = AutoTokenizer.from_pretrained(model_dir)
out = []
for s in S:
  a = tk.encode(s, add_special_tokens=False).ids
  b = tk.encode(s, add_special_tokens=True).ids
  c = auto(s).input_ids
  out.append({"s": s, "ids": a, "ids_special": b, "auto": c, "decode": tk.decode(a, skip_special_tokens=False)})
batch = [e.ids for e in tk.encode_batch(S, add_special_tokens=False)]
assert batch == [o["ids"] for o in out]
json.dump(out, sys.stdout, ensure_ascii=False)

fn main() -> Result<()> {
  let model_dir = "SmolLM2-135M-Instruct";
  let path = std::env::args().nth(1).context("strings.json")?;
  let strings: Vec<String> = serde_json::from_slice(&std::fs::read(path)?)?;
  let tk = Tokenizer::from_file(format!("{model_dir}/tokenizer.json")).map_err(anyhow::Error::msg)?;
  let batch = tk.encode_batch(strings.clone(), false).map_err(anyhow::Error::msg)?;
  let mut out = Vec::new();
  for (s, b) in strings.iter().zip(&batch) {
    let ids = tk.encode(s.as_str(), false).map_err(anyhow::Error::msg)?.get_ids().to_vec();
    let ids_special = tk.encode(s.as_str(), true).map_err(anyhow::Error::msg)?.get_ids().to_vec();
    anyhow::ensure!(b.get_ids() == ids.as_slice(), "batch differs for {s:?}");
    let decode = tk.decode(&ids, false).map_err(anyhow::Error::msg)?;
    out.push(serde_json::json!({ "s": s, "ids": ids, "ids_special": ids_special, "decode": decode }));
  }
  println!("{}", serde_json::Value::Array(out));
  Ok(())
}

28 strings, 247 tokens, 0 differences

DIFF ids '你好，世界！': [18645, 250, 48392, 138, 12831, 7906, 240, 178, 239, 230, 8083, 219] vs [18645, 250, 48392, 139, 12831, 7906, 240, 178, 239, 230, 8083, 219]
28 strings, 247 tokens, 1 differences

[1, 4093, 198, 6004, 2]

let mut tk = Tokenizer::from_file(format!("{model_dir}/tokenizer.json")).map_err(anyhow::Error::msg)?;
let s = "<|im_start|>user\nhi<|im_end|>";
println!("default:                {:?}", tk.encode(s, false).map_err(anyhow::Error::msg)?.get_ids());
tk.set_encode_special_tokens(true);
println!("encode_special_tokens:  {:?}", tk.encode(s, false).map_err(anyhow::Error::msg)?.get_ids());

default:                [1, 4093, 198, 6004, 2]
encode_special_tokens:  [44, 108, 306, 79, 3738, 108, 46, 4093, 198, 6004, 44, 108, 306, 79, 486, 108, 46]

{% for message in messages %}{% if loop.first and messages[0]['role'] != 'system' %}{{ '<|im_start|>system
You are a helpful AI assistant named SmolLM, trained by Hugging Face<|im_end|>
' }}{% endif %}{{'<|im_start|>' + message['role'] + '
' + message['content'] + '<|im_end|>' + '
'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant
' }}{% endif %}

let cfg: serde_json::Value = serde_json::from_slice(&std::fs::read(path)?)?;
let template = cfg["chat_template"].as_str().context("chat_template")?;
let mut ctx = tera::Context::new();
ctx.insert("messages", &serde_json::json!([{ "role": "user", "content": "Hi" }]));
ctx.insert("add_generation_prompt", &true);
match tera::Tera::one_off(template, &ctx, false) { /* 打印结果或错误链 */ }

tera error: Failed to parse '__tera_one_off'
caused by:  --> 3:33
  |
3 | ' }}{% endif %}{{'<|im_start|>' + message['role'] + '
  |                                 ^---
  |
  = expected `or`, `and`, `not`, `<=`, `>=`, `<`, `>`, `==`, `!=`, a filter, or a variable end (`}}`)

rendered: "<|im_start|>system\nYou are a helpful AI assistant named SmolLM, trained by Hugging Face<|im_end|>\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n"

#[derive(serde::Deserialize)]
pub struct Message {
  pub role: String,
  pub content: String,
}

/// SmolLM2-Instruct 的 chat_template，手写
pub fn render_chat(messages: &[Message], add_generation_prompt: bool) -> String {
  let mut s = String::new();
  if messages.first().is_some_and(|m| m.role != "system") {
    s.push_str("<|im_start|>system\nYou are a helpful AI assistant named SmolLM, trained by Hugging Face<|im_end|>\n");
  }
  for m in messages {
    s.push_str(&format!("<|im_start|>{}\n{}<|im_end|>\n", m.role, m.content));
  }
  if add_generation_prompt {
    s.push_str("<|im_start|>assistant\n");
  }
  s
}

CONVS = [
  ([{"role": "user", "content": "How do you say thank you in Chinese?"}], True),
  ([{"role": "user", "content": "中国的首都是哪里？"}], False),
  ([{"role": "system", "content": "你是一个只说中文的助手。"}, {"role": "user", "content": "Hi 👋"}], True),
  ([{"role": "user", "content": "1+1=?"}, {"role": "assistant", "content": "2"},
    {"role": "user", "content": "  再加 1 呢？\n"}], True),
  ([{"role": "user", "content": "first"}, {"role": "system", "content": "late system"}], True),
]
for msgs, agp in CONVS:
  text = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=agp)
  ids = tok.apply_chat_template(msgs, add_generation_prompt=agp, return_dict=True)["input_ids"]

case 0: 184 bytes, text equal=true, 39 ids, ids equal=true
        without default system: text equal=false
case 1: 153 bytes, text equal=true, 43 ids, ids equal=true
        without default system: text equal=false
case 2: 123 bytes, text equal=true, 34 ids, ids equal=true
        without default system: text equal=true
case 3: 233 bytes, text equal=true, 59 ids, ids equal=true
        without default system: text equal=false
case 4: 194 bytes, text equal=true, 38 ids, ids equal=true
        without default system: text equal=false

raw        '<|im_end|>'
no system  'In Chinese, "佛" is used to express gratitude or appreciation. Here are some examples:\n\n- 我愿�'
template   'In Chinese, you can say "六六" (Shiè) to express gratitude. Here\'s a translation:\n\n"�'

'1234567890 3.14159' 18 [33, 34, 35, 36, 37, 38, 39, 40, 41, 32, 216, 35, 30, 33, 36, 33, 37, 41]
['1', '2', '3', '4', '5', '6', '7', '8', '9', '0', 'Ġ', '3', '.', '1', '4', '1', '5', '9']