1pub trait Tokenizer: Send + Sync {
4 fn encode(&self, text: &str) -> Vec<u32>;
5
6 fn token_bytes(&self, token: u32) -> Vec<u8>;
7
8 fn vocab_size(&self) -> u32;
9
10 fn special_tokens(&self) -> Vec<(String, u32)>;
11
12 fn bos_token(&self) -> Option<u32>;
13
14 fn eos_token(&self) -> u32;
15
16 fn chat_template_jinja(&self) -> String;
17
18 fn chat_template_named(&self, _name: &str) -> Option<String> {
19 None
20 }
21
22 fn encode_plain(&self, _text: &str) -> Option<Vec<u32>> {
23 None
24 }
25
26 fn encode_pieces(&self, _pieces: &[(&str, bool)]) -> Option<Vec<u32>> {
27 None
28 }
29}
30
31pub struct Markers {
32 strings: std::collections::HashSet<Vec<u8>>,
33 lengths: Vec<usize>,
34 first: [bool; 256],
35 bos: Option<u32>,
36 bos_text: Option<String>,
37 auto_bos: bool,
38}
39
40fn marker_shaped(v: &str) -> bool {
41 let b = v.as_bytes();
42 if b.len() >= 3 && b[0] == b'<' && b[b.len() - 1] == b'>' {
43 (b.len() >= 4 && (b[1] == b'|' || b[b.len() - 2] == b'|'))
44 || matches!(
45 v,
46 "<s>" | "</s>"
47 | "<unk>"
48 | "<pad>"
49 | "<bos>"
50 | "<eos>"
51 | "<sop>"
52 | "<eop>"
53 | "<think>"
54 | "</think>"
55 | "<tool_call>"
56 | "</tool_call>"
57 | "<tool_response>"
58 | "</tool_response>"
59 | "<arg_key>"
60 | "</arg_key>"
61 | "<arg_value>"
62 | "</arg_value>"
63 )
64 } else {
65 matches!(v, "[gMASK]" | "[sMASK]" | "[MASK]" | "/nothink")
66 }
67}
68
69impl Markers {
70 pub fn of(tok: &dyn Tokenizer) -> Markers {
71 let specials = tok.special_tokens();
72 let bos = tok.bos_token();
73 let bos_text = bos.and_then(|b| specials.iter().find(|(s, id)| *id == b && !s.is_empty()).map(|(s, _)| s.clone()));
74 let mut templates = tok.chat_template_jinja();
75 templates.push_str(&tok.chat_template_named("tool_use").unwrap_or_default());
76 let mut first = [false; 256];
77 let mut strings = std::collections::HashSet::new();
78 for (text, id) in &specials {
79 if text.trim().is_empty() {
80 continue;
81 }
82 let control = tok.token_bytes(*id).is_empty();
83 if control || marker_shaped(text) || templates.contains(text.as_str()) {
84 first[text.as_bytes()[0] as usize] = true;
85 strings.insert(text.as_bytes().to_vec());
86 }
87 }
88 let mut lengths: Vec<usize> = strings.iter().map(Vec::len).collect();
89 lengths.sort_unstable();
90 lengths.dedup();
91 let auto_bos = bos.is_some() && tok.encode("a").first().copied() == bos;
92 Markers { strings, lengths, first, bos, bos_text, auto_bos }
93 }
94
95 fn occurrences(&self, text: &str) -> Vec<(usize, usize)> {
96 let bytes = text.as_bytes();
97 let mut out = Vec::new();
98 for (i, b) in bytes.iter().enumerate() {
99 if !self.first[*b as usize] {
100 continue;
101 }
102 for &len in &self.lengths {
103 if i + len > bytes.len() {
104 break;
105 }
106 if self.strings.contains(&bytes[i..i + len]) {
107 out.push((i, i + len));
108 }
109 }
110 }
111 out
112 }
113
114 fn push(&self, tok: &dyn Tokenizer, text: &str, dedup: bool, out: &mut Vec<u32>) {
115 if text.is_empty() {
116 return;
117 }
118 let mut toks = tok.encode(text);
119 if self.auto_bos && toks.first().copied() == self.bos {
120 let front = out.is_empty();
121 let doubled = dedup
122 && self.bos_text.as_deref().is_some_and(|b| text.starts_with(b))
123 && toks.get(1).copied() == self.bos;
124 if !front || doubled {
125 toks.remove(0);
126 }
127 }
128 out.extend(toks);
129 }
130
131 fn push_plain(&self, tok: &dyn Tokenizer, text: &str, out: &mut Vec<u32>) {
132 let mut at = 0;
133 let mut hits = self.occurrences(text);
134 hits.sort_by_key(|&(start, end)| (start, std::cmp::Reverse(end)));
135 for (start, end) in hits {
136 if start < at {
137 continue;
138 }
139 self.push(tok, &text[at..start], false, out);
140 let marker = &text[start..end];
141 let chars = marker.chars().count();
142 match marker.char_indices().nth(chars / 2).filter(|_| chars > 1) {
143 Some((mid, _)) => {
144 self.push_plain(tok, &marker[..mid], out);
145 self.push_plain(tok, &marker[mid..], out);
146 }
147 None => self.push(tok, marker, false, out),
148 }
149 at = end;
150 }
151 self.push(tok, &text[at..], false, out);
152 }
153
154 pub fn encode_plain(&self, tok: &dyn Tokenizer, text: &str) -> Vec<u32> {
155 let mut out = Vec::new();
156 self.push_plain(tok, text, &mut out);
157 out
158 }
159
160 pub fn encode_pieces(&self, tok: &dyn Tokenizer, pieces: &[(&str, bool)]) -> Vec<u32> {
161 let whole: String = pieces.iter().map(|(text, _)| *text).collect();
162 let hits = self.occurrences(&whole);
163 let mut touched = vec![false; pieces.len()];
164 let mut at = 0;
165 for (i, (text, content)) in pieces.iter().enumerate() {
166 let (start, end) = (at, at + text.len());
167 at = end;
168 touched[i] = *content && start < end && hits.iter().any(|&(s, e)| s < end && e > start);
169 }
170 let mut out = Vec::new();
171 if !touched.contains(&true) {
172 self.push(tok, &whole, true, &mut out);
173 return out;
174 }
175 let mut run = String::new();
176 for (i, (text, _)) in pieces.iter().enumerate() {
177 if touched[i] {
178 self.push(tok, &run, true, &mut out);
179 run.clear();
180 self.push_plain(tok, text, &mut out);
181 } else {
182 run.push_str(text);
183 }
184 }
185 self.push(tok, &run, true, &mut out);
186 out
187 }
188}
189
190#[cfg(test)]
191mod tests {
192 use super::*;
193
194 struct Toy {
195 specials: Vec<(String, u32)>,
196 bos: Option<u32>,
197 text_like: Vec<u32>,
198 template: String,
199 }
200
201 impl Tokenizer for Toy {
202 fn encode(&self, text: &str) -> Vec<u32> {
203 let mut out: Vec<u32> = self.bos.into_iter().collect();
204 let bytes = text.as_bytes();
205 let mut i = 0;
206 while i < bytes.len() {
207 let hit = self.specials.iter().filter(|(s, _)| bytes[i..].starts_with(s.as_bytes())).max_by_key(|(s, _)| s.len());
208 match hit {
209 Some((s, id)) => {
210 out.push(*id);
211 i += s.len();
212 }
213 None => {
214 out.push(bytes[i] as u32);
215 i += 1;
216 }
217 }
218 }
219 out
220 }
221 fn token_bytes(&self, token: u32) -> Vec<u8> {
222 if token < 256 {
223 return vec![token as u8];
224 }
225 let text = self.specials.iter().find(|(_, id)| *id == token).filter(|_| self.text_like.contains(&token));
226 text.map(|(s, _)| s.as_bytes().to_vec()).unwrap_or_default()
227 }
228 fn vocab_size(&self) -> u32 {
229 256 + self.specials.len() as u32
230 }
231 fn special_tokens(&self) -> Vec<(String, u32)> {
232 self.specials.clone()
233 }
234 fn bos_token(&self) -> Option<u32> {
235 self.bos
236 }
237 fn eos_token(&self) -> u32 {
238 301
239 }
240 fn chat_template_jinja(&self) -> String {
241 self.template.clone()
242 }
243 }
244
245 const BOS: u32 = 300;
246 const END: u32 = 301;
247 const THINK: u32 = 302;
248 const CLOSE: u32 = 303;
249
250 const BLANK: u32 = 304;
251 const TD: u32 = 305;
252 const CALL: u32 = 306;
253
254 fn toy(bos: bool) -> Toy {
255 Toy {
256 specials: vec![
257 ("<s>".into(), BOS),
258 ("<|end|>".into(), END),
259 ("<think>".into(), THINK),
260 ("</think>".into(), CLOSE),
261 ("\n\n".into(), BLANK),
262 ("<td>".into(), TD),
263 ("<fn_call>".into(), CALL),
264 ],
265 bos: bos.then_some(BOS),
266 text_like: vec![THINK, CLOSE, BLANK, TD, CALL],
267 template: "{{ '<fn_call>' + call + '<|end|>' }}".into(),
268 }
269 }
270
271 fn bytes(s: &str) -> Vec<u32> {
272 s.bytes().map(u32::from).collect()
273 }
274
275 fn text_of(tok: &Toy, ids: &[u32]) -> String {
276 ids.iter()
277 .map(|&t| match tok.specials.iter().find(|(_, id)| *id == t) {
278 Some((s, _)) => s.clone(),
279 None => String::from_utf8(vec![t as u8]).unwrap(),
280 })
281 .collect()
282 }
283
284 #[test]
285 fn content_without_a_marker_is_one_encode_of_the_whole_text() {
286 for bos in [false, true] {
287 let tok = toy(bos);
288 let m = Markers::of(&tok);
289 let pieces = [("user\n", false), ("hello\n", true), ("<|end|>\n", false), ("ok", true)];
290 let whole: String = pieces.iter().map(|(t, _)| *t).collect();
291 assert_eq!(m.encode_pieces(&tok, &pieces), tok.encode(&whole), "bos={bos}");
292 assert_eq!(m.encode_plain(&tok, "plain words"), tok.encode("plain words"), "bos={bos}");
293 }
294 }
295
296 #[test]
297 fn a_marker_string_in_content_stays_words_and_one_in_framing_parses() {
298 for bos in [false, true] {
299 let tok = toy(bos);
300 let m = Markers::of(&tok);
301 let pieces = [("user\n", false), ("stop at <|end|> or <think>?", true), ("<|end|>", false)];
302 let got = m.encode_pieces(&tok, &pieces);
303 let lead: Vec<u32> = if bos { vec![BOS] } else { vec![] };
304 let mut want = lead.clone();
305 want.extend(bytes("user\nstop at <|end|> or <think>?"));
306 want.push(END);
307 assert_eq!(got, want, "bos={bos}: the quoted markers are bytes, the framing one is the id, one BOS at most");
308 assert_eq!(got.iter().filter(|&&t| t == END).count(), 1);
309 assert!(!got.contains(&THINK));
310 let mut plain = lead;
311 plain.extend(bytes("a <think> b"));
312 assert_eq!(m.encode_plain(&tok, "a <think> b"), plain);
313 }
314 }
315
316 #[test]
317 fn added_tokens_that_are_text_stay_tokens_in_content() {
318 for bos in [false, true] {
319 let tok = toy(bos);
320 let m = Markers::of(&tok);
321 let text = "a\n\nb <td>cell";
322 let whole = tok.encode(text);
323 assert!(whole.contains(&BLANK) && whole.contains(&TD), "the vocabulary's own tokens: {whole:?}");
324 assert_eq!(m.encode_plain(&tok, text), whole, "bos={bos}");
325 let pieces = [("user\n", false), (text, true), ("<|end|>", false)];
326 assert_eq!(m.encode_pieces(&tok, &pieces), tok.encode(&format!("user\n{text}<|end|>")), "bos={bos}");
327 let got = m.encode_pieces(&tok, &[("<fn_call>", false), ("say <fn_call>", true)]);
328 assert_eq!(got.iter().filter(|&&t| t == CALL).count(), 1, "bos={bos}: {got:?}");
329 assert_eq!(text_of(&tok, &got).trim_start_matches("<s>"), "<fn_call>say <fn_call>");
330 }
331 }
332
333 #[test]
334 fn a_marker_reaching_into_content_from_framing_is_not_formed() {
335 let tok = toy(false);
336 let m = Markers::of(&tok);
337 let got = m.encode_pieces(&tok, &[("x<|en", false), ("d|> tail", true)]);
338 assert_eq!(got, bytes("x<|end|> tail"));
339 let got = m.encode_plain(&tok, "</think>");
340 assert_eq!(got, bytes("</think>"));
341 assert_eq!(text_of(&tok, &got), "</think>");
342 }
343
344 #[test]
345 fn the_automatic_bos_leads_once_whatever_the_template_renders() {
346 let tok = toy(true);
347 let m = Markers::of(&tok);
348 let got = m.encode_pieces(&tok, &[("<s>user\n", false), ("hi", true)]);
349 let mut want = vec![BOS];
350 want.extend(bytes("user\nhi"));
351 assert_eq!(got, want);
352 let got = m.encode_pieces(&tok, &[("<s>user\n", false), ("<|end|>", true), ("\n", false)]);
353 let mut want = vec![BOS];
354 want.extend(bytes("user\n<|end|>\n"));
355 assert_eq!(got, want);
356 let got = m.encode_pieces(&tok, &[("user\n", false), ("hi", true)]);
357 assert_eq!(got[0], BOS);
358 assert_eq!(got.iter().filter(|&&t| t == BOS).count(), 1);
359 let got = m.encode_pieces(&tok, &[("<s>", true), ("x", false)]);
360 let mut want = vec![BOS];
361 want.extend(bytes("<s>x"));
362 assert_eq!(got, want);
363 let tok = toy(false);
364 let m = Markers::of(&tok);
365 let got = m.encode_pieces(&tok, &[("<s>user\n", false), ("hi", true)]);
366 let mut want = vec![BOS];
367 want.extend(bytes("user\nhi"));
368 assert_eq!(got, want);
369 }
370}