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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
//! Tokenizer specification returned by [`crate::ModelSource::tokenizer`].
use std::collections::HashMap;
use std::path::PathBuf;
/// Where to find the tokenizer and which special tokens it defines.
#[derive(Debug, Clone)]
pub struct TokenizerSpec {
/// Path to the HuggingFace `tokenizer.json`.
pub tokenizer_json: PathBuf,
/// Added special tokens parsed from `tokenizer_config.json`
/// (`added_tokens_decoder`): token id → token string (e.g. `<|im_end|>`).
pub added_tokens: HashMap<u32, String>,
/// Raw chat template string (Jinja) if present in `tokenizer_config.json`.
/// Phase 1 uses a built-in ChatML wrap instead of evaluating Jinja.
pub chat_template: Option<String>,
/// Whether prompts should be prefixed with BOS (HF `add_bos_token` /
/// GGUF `tokenizer.ggml.add_bos_token`). `None` = unspecified, which the
/// engine treats as "prepend when the model declares a BOS id". Qwen2
/// declares a BOS id but sets this to `false`; ignoring it prepends
/// `<|endoftext|>` to every prompt.
pub add_bos: Option<bool>,
}
impl TokenizerSpec {
/// Dummy spec used by weight-only safetensors sources that have no
/// tokenizer. `tokenizer()` on those sources errors before this is used.
pub(crate) fn placeholder() -> Self {
Self {
tokenizer_json: PathBuf::new(),
added_tokens: HashMap::new(),
chat_template: None,
add_bos: None,
}
}
/// Looks up the id of an added special token by its string, e.g.
/// `spec.special_token_id("<|im_end|>")`.
pub fn special_token_id(&self, token: &str) -> Option<u32> {
self.added_tokens
.iter()
.find(|(_, s)| s.as_str() == token)
.map(|(id, _)| *id)
}
/// Wraps a user prompt in the ChatML template used by SmolLM2-style
/// instruction models. Returns `None` if the tokenizer does not define
/// `<|im_start|>` / `<|im_end|>` tokens.
pub fn chatml_wrap(&self, user_prompt: &str) -> Option<String> {
let start = self.special_token_id("<|im_start|>")?;
self.special_token_id("<|im_end|>")?;
let _ = start;
Some(format!(
"<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
))
}
/// Which chat template this tokenizer's special tokens imply.
pub fn chat_template_kind(&self) -> ChatTemplate {
if self.special_token_id("<start_of_turn>").is_some() {
ChatTemplate::Gemma
} else {
ChatTemplate::Chatml
}
}
/// Wraps (role, content) message pairs into the model's chat format,
/// ending with the assistant turn left open for generation. Unknown
/// roles are coerced to `user` (same convention as serve).
pub fn wrap_messages(&self, messages: &[(String, String)]) -> String {
match self.chat_template_kind() {
ChatTemplate::Chatml => {
let mut out = String::new();
for (role, content) in messages {
let role = match role.as_str() {
"system" | "user" | "assistant" => role.as_str(),
_ => "user",
};
out.push_str(&format!("<|im_start|>{role}\n{content}<|im_end|>\n"));
}
out.push_str("<|im_start|>assistant\n");
out
}
ChatTemplate::Gemma => {
// Gemma-3 template: <bos> prefix, turns are
// `<start_of_turn>{user|model}\n{content}<end_of_turn>\n`,
// and system content is folded into the first user turn
// (the HF template prepends it with a blank line).
let mut out = String::from("<bos>");
let mut system_prefix = String::new();
let mut first_user_seen = false;
for (role, content) in messages {
match role.as_str() {
"system" => {
system_prefix.push_str(content);
system_prefix.push_str("\n\n");
}
_ => {
let turn_role = match role.as_str() {
"assistant" | "model" => "model",
_ => "user",
};
let mut body = String::new();
if turn_role == "user" && !first_user_seen {
body.push_str(&system_prefix);
first_user_seen = true;
}
body.push_str(content);
out.push_str(&format!(
"<start_of_turn>{turn_role}\n{body}<end_of_turn>\n"
));
}
}
}
out.push_str("<start_of_turn>model\n");
out
}
}
}
}
/// Chat template flavor detected from the tokenizer's special tokens.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ChatTemplate {
/// `<|im_start|>role\n…<|im_end|>` (SmolLM2 / Qwen style).
Chatml,
/// `<start_of_turn>user\n…<end_of_turn>` (Gemma style).
Gemma,
}