1use serde::Deserialize;
13use std::collections::{HashMap, HashSet};
14use std::path::Path;
15use std::sync::OnceLock;
16use unicode_normalization::UnicodeNormalization;
17
18const DEFAULT_SPLIT: &str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+";
21
22pub const TOOL_MARKUP_TOKENS: &[&str] = &[
34 "<tool_call>",
35 "</tool_call>",
36 "<function",
37 "</function>",
38 "<param",
39 "</param>",
40];
41
42pub struct Tokenizer {
44 vocab: HashMap<String, u32>,
46 id_to_token: Vec<String>,
48 ranks: HashMap<(String, String), u32>,
50 sp_ids: OnceLock<Option<SpMerges>>,
54 added: Vec<(String, u32)>,
56 added_ids: HashSet<u32>,
58 special_ids: HashSet<u32>,
60 split_res: Vec<fancy_regex::Regex>,
64 split_drop_gap: Vec<bool>,
70 eow_suffix: Option<String>,
75 lowercase: bool,
78 collapse_ws: bool,
81 sp_prepend: bool,
84 sp_prepend_first: bool,
90 metaspace: bool,
93 nfc: bool,
96 byte_to_char: [char; 256],
98 char_to_byte: HashMap<char, u8>,
100 pub bos_token_id: Option<u32>,
102 pub eos_token_id: Option<u32>,
103 pub pad_token_id: Option<u32>,
104 pub im_start_id: Option<u32>,
106 pub im_end_id: Option<u32>,
107 pub chat_template: Option<String>,
110 pub extra_eos: HashSet<u32>,
112 pub add_bos: bool,
114}
115
116impl std::fmt::Debug for Tokenizer {
117 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
118 f.debug_struct("Tokenizer")
119 .field("vocab", &self.vocab.len())
120 .field("merges", &self.ranks.len())
121 .field("added", &self.added.len())
122 .finish()
123 }
124}
125
126fn bytes_to_unicode() -> ([char; 256], HashMap<char, u8>) {
129 let mut b2c = ['\0'; 256];
130 let mut c2b = HashMap::with_capacity(256);
131 let mut n = 0u32;
132 for b in 0..=255u16 {
133 let printable =
134 (0x21..=0x7E).contains(&b) || (0xA1..=0xAC).contains(&b) || (0xAE..=0xFF).contains(&b);
135 let c = if printable {
136 char::from_u32(b as u32).unwrap()
137 } else {
138 let c = char::from_u32(256 + n).unwrap();
139 n += 1;
140 c
141 };
142 b2c[b as usize] = c;
143 c2b.insert(c, b as u8);
144 }
145 (b2c, c2b)
146}
147
148#[derive(Deserialize)]
150struct HfTokenizerJson {
151 model: HfModel,
152 #[serde(default)]
153 added_tokens: Vec<HfAddedToken>,
154 #[serde(default)]
155 pre_tokenizer: Option<serde_json::Value>,
156 #[serde(default)]
157 normalizer: Option<serde_json::Value>,
158 #[serde(default)]
159 post_processor: Option<serde_json::Value>,
160}
161
162#[derive(Deserialize)]
163struct HfModel {
164 vocab: HashMap<String, u32>,
165 #[serde(default)]
166 merges: Vec<HfMerge>,
167 #[serde(default)]
168 byte_fallback: bool,
169 #[serde(default)]
171 end_of_word_suffix: Option<String>,
172}
173
174#[derive(Deserialize)]
177#[serde(untagged)]
178enum HfMerge {
179 Pair([String; 2]),
180 Text(String),
181}
182
183#[derive(Deserialize)]
184struct HfAddedToken {
185 id: u32,
186 content: String,
187 special: bool,
188}
189
190fn collect_split_patterns(pt: &serde_json::Value, out: &mut Vec<String>) {
199 if pt.get("type").and_then(|t| t.as_str()) == Some("Split") {
200 if let Some(r) = pt
201 .get("pattern")
202 .and_then(|p| p.get("Regex"))
203 .and_then(|r| r.as_str())
204 {
205 out.push(r.to_string());
206 }
207 return;
208 }
209 if let Some(list) = pt.get("pretokenizers").and_then(|l| l.as_array()) {
210 for p in list {
211 collect_split_patterns(p, out);
212 }
213 }
214}
215
216fn collect_split_drop_gaps(pt: &serde_json::Value, out: &mut Vec<bool>) {
220 if pt.get("type").and_then(|t| t.as_str()) == Some("Split") {
221 let behavior = pt.get("behavior").and_then(|b| b.as_str()).unwrap_or("");
222 let invert = pt.get("invert").and_then(|i| i.as_bool()).unwrap_or(false);
223 out.push(behavior == "Removed" && invert);
224 return;
225 }
226 if let Some(list) = pt.get("pretokenizers").and_then(|l| l.as_array()) {
227 for p in list {
228 collect_split_drop_gaps(p, out);
229 }
230 }
231}
232
233fn find_prepend_scheme(pt: &serde_json::Value) -> Option<String> {
236 if pt.get("type").and_then(|t| t.as_str()) == Some("Metaspace") {
237 return pt
238 .get("prepend_scheme")
239 .and_then(|p| p.as_str())
240 .map(String::from);
241 }
242 if let Some(list) = pt.get("pretokenizers").and_then(|l| l.as_array()) {
243 return list.iter().find_map(find_prepend_scheme);
244 }
245 None
246}
247
248pub(crate) fn strip_generation_tags(tpl: &str) -> std::borrow::Cow<'_, str> {
264 if !tpl.contains("generation") {
265 return std::borrow::Cow::Borrowed(tpl);
266 }
267 let mut out = String::with_capacity(tpl.len());
268 let mut rest = tpl;
269 let mut touched = false;
270 while let Some(open) = rest.find("{%") {
271 let Some(close_rel) = rest[open..].find("%}") else {
272 break;
273 };
274 let close = open + close_rel + 2;
275 let tag = &rest[open..close];
276 let inner = tag[2..tag.len() - 2].trim();
277 let lead = inner.starts_with('-');
278 let trail = inner.ends_with('-');
279 let name = inner.trim_matches('-').trim();
280 out.push_str(&rest[..open]);
281 if name == "generation" || name == "endgeneration" {
282 out.push_str(if lead { "{%-" } else { "{%" });
283 out.push_str(" set _generation_span = true ");
284 out.push_str(if trail { "-%}" } else { "%}" });
285 touched = true;
286 } else {
287 out.push_str(tag);
288 }
289 rest = &rest[close..];
290 }
291 if !touched {
292 return std::borrow::Cow::Borrowed(tpl);
293 }
294 out.push_str(rest);
295 std::borrow::Cow::Owned(out)
296}
297
298fn close_direct_think(rendered: String, enable_thinking: Option<bool>) -> String {
307 if enable_thinking != Some(false) || rendered.contains("</think>") {
308 return rendered;
309 }
310 if let Some(pos) = rendered.rfind("<|assistant|>") {
311 let marker_end = pos + "<|assistant|>".len();
312 let after = &rendered[marker_end..];
313 let mut out = String::with_capacity(rendered.len() + 24);
314 out.push_str(&rendered[..marker_end]);
315 if let Some(rest) = after.strip_prefix("<think>") {
316 out.push_str("<think></think>\n\n");
317 out.push_str(rest);
318 } else {
319 out.push_str(after);
320 if !out.ends_with('\n') {
321 out.push('\n');
322 }
323 out.push_str("<think>\n\n</think>\n\n");
324 }
325 return out;
326 }
327 if let Some(pos) = rendered.rfind("\nassistant") {
331 let mut insert_at = pos + 1 + "assistant".len();
332 if let Some(idx) = rendered[insert_at..].find('\n') {
333 insert_at += idx + 1;
334 }
335 let mut out = String::with_capacity(rendered.len() + 24);
336 out.push_str(&rendered[..insert_at]);
337 if !out.ends_with('\n') {
338 out.push('\n');
339 }
340 out.push_str("<think>\n\n</think>\n\n");
341 out.push_str(&rendered[insert_at..]);
342 return out;
343 }
344 rendered
345}
346
347impl Tokenizer {
348 pub fn from_file(path: impl AsRef<Path>) -> Result<Self, TokenizerError> {
350 let data = std::fs::read_to_string(path.as_ref())
351 .map_err(|e| TokenizerError::Io(e.to_string()))?;
352 Self::from_json(&data)
353 }
354
355 pub fn from_bytes(bytes: &[u8]) -> Result<Self, TokenizerError> {
357 let s = std::str::from_utf8(bytes)
358 .map_err(|e| TokenizerError::Parse(format!("vocab is not UTF-8: {e}")))?;
359 Self::from_json(s)
360 }
361
362 pub fn from_json(json: &str) -> Result<Self, TokenizerError> {
364 let hf: HfTokenizerJson =
365 serde_json::from_str(json).map_err(|e| TokenizerError::Parse(e.to_string()))?;
366
367 let mut vocab = hf.model.vocab;
368 let mut ranks = HashMap::new();
369 for (rank, m) in hf.model.merges.into_iter().enumerate() {
370 let (a, b) = match m {
371 HfMerge::Pair([a, b]) => (a, b),
372 HfMerge::Text(s) => {
373 let mut it = s.splitn(2, ' ');
374 match (it.next(), it.next()) {
375 (Some(a), Some(b)) => (a.to_string(), b.to_string()),
376 _ => continue,
377 }
378 }
379 };
380 ranks.insert((a, b), rank as u32);
381 }
382
383 let mut saw_gemma_bos = false;
388 let add_bos_detected = hf
389 .post_processor
390 .as_ref()
391 .map(|p| {
392 let pp = p.to_string();
393 pp.contains("\"<s>\"") || pp.contains("\"<bos>\"")
394 })
395 .unwrap_or(false);
396 let nfc = hf
397 .normalizer
398 .as_ref()
399 .map(|n| n.to_string().contains("NFC"))
400 .unwrap_or(false);
401 let metaspace = hf.model.byte_fallback
402 || hf
403 .normalizer
404 .as_ref()
405 .map(|n| n.to_string().contains("\u{2581}") || n.to_string().contains("▁"))
406 .unwrap_or(false);
407 let sp_prepend = hf
408 .normalizer
409 .as_ref()
410 .map(|n| n.to_string().contains("Prepend"))
411 .unwrap_or(false);
412 let (sp_prepend, sp_prepend_first) = if sp_prepend {
416 (true, false)
417 } else {
418 match hf.pre_tokenizer.as_ref().and_then(find_prepend_scheme) {
419 Some(s) if s == "always" => (true, false),
420 Some(s) if s == "first" => (false, true),
421 _ => (false, false),
422 }
423 };
424 let (split_res, split_drop_gap) = if metaspace {
425 (Vec::new(), Vec::new())
426 } else {
427 let mut pats = Vec::new();
428 if let Some(pt) = hf.pre_tokenizer.as_ref() {
429 collect_split_patterns(pt, &mut pats);
430 }
431 if pats.is_empty() {
432 pats.push(DEFAULT_SPLIT.to_string());
433 }
434 let mut split_drop_gap = Vec::new();
435 if let Some(pt) = hf.pre_tokenizer.as_ref() {
436 collect_split_drop_gaps(pt, &mut split_drop_gap);
437 }
438 while split_drop_gap.len() < pats.len() {
439 split_drop_gap.push(false);
440 }
441 split_drop_gap.truncate(pats.len());
442 (
443 pats.iter()
444 .map(|p| {
445 fancy_regex::Regex::new(p)
446 .map_err(|e| TokenizerError::Parse(format!("pre-tokenizer regex: {e}")))
447 })
448 .collect::<Result<Vec<_>, _>>()?,
449 split_drop_gap,
450 )
451 };
452 let norm_s = hf
454 .normalizer
455 .as_ref()
456 .map(|n| n.to_string())
457 .unwrap_or_default();
458 let collapse_ws = norm_s.contains("\"Replace\"") && norm_s.contains("\\\\s+");
459 let lowercase = norm_s.contains("\"Lowercase\"");
460
461 let mut bos_token_id = None;
463 let mut eos_token_id = None;
464 let mut pad_token_id = None;
465 let mut im_start_id = None;
466 let mut im_end_id = None;
467 let mut special_ids = HashSet::new();
468 let mut added_ids = HashSet::new();
469 let mut added = Vec::new();
470
471 for at in &hf.added_tokens {
472 vocab.insert(at.content.clone(), at.id);
473 added.push((at.content.clone(), at.id));
474 added_ids.insert(at.id);
475 if at.special && !TOOL_MARKUP_TOKENS.contains(&at.content.as_str()) {
476 special_ids.insert(at.id);
477 }
478 match at.content.as_str() {
479 "<|endoftext|>" | "</s>" | "[EOS]" => eos_token_id = Some(at.id),
480 "<|im_start|>" => im_start_id = Some(at.id),
481 "<|im_end|>" => im_end_id = Some(at.id),
482 "<s>" | "[BOS]" | "<|startoftext|>" => bos_token_id = Some(at.id),
483 "<bos>" => {
487 bos_token_id = Some(at.id);
488 saw_gemma_bos = true;
489 }
490 "<pad>" => pad_token_id = Some(at.id),
491 _ => {}
492 }
493 }
494
495 const DSV41_SPECIALS: &[&str] = &[
500 "<|begin▁of▁sentence|>",
501 "<|end▁of▁sentence|>",
502 "<|User|>",
503 "<|Assistant|>",
504 "<|System|>",
505 "<|latest_reminder|>",
506 "<|deepseek_image|>",
507 "<|action|>",
508 "<|query|>",
509 "<|authority|>",
510 "<|domain|>",
511 "<|title|>",
512 "<|read_url|>",
513 "<think>",
514 "</think>",
515 "|DSML|",
516 ];
517 for token in DSV41_SPECIALS {
518 if let Some(&id) = vocab.get(*token) {
519 if !added.iter().any(|(content, _)| content.as_str() == *token) {
520 added.push(((*token).to_string(), id));
521 }
522 added_ids.insert(id);
523 match *token {
524 "<|begin▁of▁sentence|>" => bos_token_id = Some(id),
525 "<|end▁of▁sentence|>" => eos_token_id = Some(id),
526 _ => {}
527 }
528 }
529 }
530 added.sort_by_key(|(c, _)| std::cmp::Reverse(c.len()));
531
532 let gemma_family = saw_gemma_bos
539 || vocab.contains_key("<start_of_turn>")
540 || added.iter().any(|(c, _)| c == "<start_of_turn>");
541
542 if let Some(pp) = hf.post_processor.as_ref() {
546 let pp = pp.to_string();
547 for name in ["<bos>", "<s>"] {
548 if pp.contains(&format!("\"{name}\"")) {
549 if let Some(&id) = vocab.get(name) {
550 bos_token_id = Some(id);
551 }
552 break;
553 }
554 }
555 }
556
557 let max_id = vocab.values().copied().max().unwrap_or(0) as usize;
559 let mut id_to_token = vec![String::new(); max_id + 1];
560 for (token, &id) in &vocab {
561 if (id as usize) < id_to_token.len() {
562 id_to_token[id as usize] = token.clone();
563 }
564 }
565
566 let (byte_to_char, char_to_byte) = bytes_to_unicode();
567
568 tracing::info!(
569 "Tokenizer loaded: {} vocab, {} merges, {} added, eos={:?}",
570 vocab.len(),
571 ranks.len(),
572 added.len(),
573 eos_token_id
574 );
575
576 Ok(Self {
577 vocab,
578 id_to_token,
579 ranks,
580 sp_ids: OnceLock::new(),
581 added,
582 added_ids,
583 special_ids,
584 split_res,
585 split_drop_gap,
586 eow_suffix: hf.model.end_of_word_suffix,
587 lowercase,
588 collapse_ws,
589 metaspace,
590 sp_prepend,
591 sp_prepend_first,
592 nfc,
593 byte_to_char,
594 char_to_byte,
595 bos_token_id,
596 eos_token_id,
597 pad_token_id,
598 im_start_id,
599 im_end_id,
600 chat_template: None,
601 extra_eos: HashSet::new(),
602 add_bos: add_bos_detected || gemma_family,
603 })
604 }
605
606 pub fn byte_level() -> Self {
608 let mut vocab = HashMap::new();
609 let mut id_to_token = Vec::with_capacity(256);
610 for i in 0..256u32 {
611 let tok = format!("<0x{:02X}>", i);
612 vocab.insert(tok.clone(), i);
613 id_to_token.push(tok);
614 }
615 let (byte_to_char, char_to_byte) = bytes_to_unicode();
616 Self {
617 vocab,
618 id_to_token,
619 ranks: HashMap::new(),
620 sp_ids: OnceLock::new(),
621 added: Vec::new(),
622 added_ids: HashSet::new(),
623 special_ids: HashSet::new(),
624 split_res: Vec::new(),
625 split_drop_gap: Vec::new(),
626 eow_suffix: None,
627 lowercase: false,
628 collapse_ws: false,
629 metaspace: false,
630 sp_prepend: false,
631 sp_prepend_first: false,
632 nfc: false,
633 byte_to_char,
634 char_to_byte,
635 bos_token_id: None,
636 eos_token_id: None,
637 pad_token_id: None,
638 im_start_id: None,
639 im_end_id: None,
640 chat_template: None,
641 extra_eos: HashSet::new(),
642 add_bos: false,
643 }
644 }
645
646 pub fn encode(&self, text: &str) -> Vec<u32> {
648 let mut ids = Vec::new();
649 let mut rest = text;
651 let mut head = true;
656 'outer: while !rest.is_empty() {
657 let mut best: Option<(usize, usize, u32)> = None; for (content, id) in &self.added {
659 if let Some(pos) = rest.find(content.as_str()) {
660 let better = match best {
661 None => true,
662 Some((bp, bl, _)) => pos < bp || (pos == bp && content.len() > bl),
663 };
664 if better {
665 best = Some((pos, content.len(), *id));
666 }
667 if pos == 0 {
668 break; }
670 }
671 }
672 match best {
673 Some((pos, len, id)) => {
674 self.encode_segment_at(&rest[..pos], head, &mut ids);
675 ids.push(id);
676 rest = &rest[pos + len..];
677 head = false;
678 }
679 None => {
680 self.encode_segment_at(rest, head, &mut ids);
681 break 'outer;
682 }
683 }
684 }
685 ids
686 }
687
688 pub fn encode_plain(&self, text: &str) -> Vec<u32> {
694 let mut ids = Vec::new();
695 self.encode_segment_at(text, true, &mut ids);
696 ids
697 }
698
699 fn encode_segment_at(&self, segment: &str, head: bool, out: &mut Vec<u32>) {
704 if segment.is_empty() {
705 return;
706 }
707 let mut norm: String = if self.nfc {
708 segment.nfc().collect()
709 } else {
710 segment.to_string()
711 };
712 if self.collapse_ws {
713 let mut collapsed = String::with_capacity(norm.len());
715 let mut in_ws = false;
716 for ch in norm.chars() {
717 if ch.is_whitespace() {
718 if !in_ws {
719 collapsed.push(' ');
720 in_ws = true;
721 }
722 } else {
723 collapsed.push(ch);
724 in_ws = false;
725 }
726 }
727 norm = collapsed;
728 }
729 if self.lowercase {
730 norm = norm.to_lowercase();
731 }
732 if self.metaspace {
733 let sp = if self.sp_prepend {
737 format!("\u{2581}{}", norm).replace(' ', "\u{2581}")
740 } else {
741 let replaced = norm.replace(' ', "\u{2581}");
745 if self.sp_prepend_first && head && !replaced.starts_with('\u{2581}') {
746 format!("\u{2581}{replaced}")
747 } else {
748 replaced
749 }
750 };
751 self.bpe_piece_sp(&sp, out);
752 return;
753 }
754 if !self.split_res.is_empty() {
755 let mut pieces: Vec<(usize, usize)> = vec![(0, norm.len())];
760 for (stage, re) in self.split_res.iter().enumerate() {
761 let drop = self.split_drop_gap.get(stage).copied().unwrap_or(false);
762 let mut next: Vec<(usize, usize)> = Vec::with_capacity(pieces.len() * 2);
763 for (ps, pe) in pieces {
764 let seg = &norm[ps..pe];
765 let mut last = 0usize;
766 for m in re.find_iter(seg) {
767 let m = match m {
768 Ok(m) => m,
769 Err(e) => {
770 tracing::error!("pre-tokenizer regex failed: {e}");
771 break;
772 }
773 };
774 if m.start() > last && !drop {
775 next.push((ps + last, ps + m.start()));
776 }
777 if m.end() > m.start() {
778 next.push((ps + m.start(), ps + m.end()));
779 }
780 last = m.end();
781 }
782 if last < seg.len() && !drop {
783 next.push((ps + last, pe));
784 }
785 }
786 pieces = next;
787 }
788 for (ps, pe) in pieces {
789 self.bpe_piece(&norm[ps..pe], out);
790 }
791 } else {
792 {
793 for b in norm.bytes() {
795 let tok = format!("<0x{:02X}>", b);
796 if let Some(&id) = self.vocab.get(&tok) {
797 out.push(id);
798 }
799 }
800 }
801 }
802 }
803
804 fn sp_merges(&self) -> Option<&SpMerges> {
806 self.sp_ids
807 .get_or_init(|| SpMerges::build(&self.vocab, &self.ranks))
808 .as_ref()
809 }
810
811 fn bpe_piece_sp(&self, piece: &str, out: &mut Vec<u32>) {
820 if piece.is_empty() {
821 return;
822 }
823 if !SP_STRINGS.with(|f| f.get()) {
824 if let Some(m) = self.sp_merges() {
825 m.encode(&self.vocab, piece, out);
826 return;
827 }
828 }
829 self.bpe_piece_sp_strings(piece, out);
830 }
831
832 #[doc(hidden)]
835 pub fn encode_reference(&self, text: &str) -> Vec<u32> {
836 SP_STRINGS.with(|f| f.set(true));
837 let ids = self.encode(text);
838 SP_STRINGS.with(|f| f.set(false));
839 ids
840 }
841
842 fn bpe_piece_sp_strings(&self, piece: &str, out: &mut Vec<u32>) {
843 let mut sym: Vec<String> = piece.chars().map(|c| c.to_string()).collect();
844 loop {
845 let mut best: Option<(u32, usize)> = None;
846 for i in 0..sym.len().saturating_sub(1) {
847 if let Some(&r) = self.ranks.get(&(sym[i].clone(), sym[i + 1].clone())) {
848 if best.map(|(br, _)| r < br).unwrap_or(true) {
849 best = Some((r, i));
850 }
851 }
852 }
853 let Some((_, i)) = best else { break };
854 let merged = format!("{}{}", sym[i], sym[i + 1]);
855 let (left, right) = (sym[i].clone(), sym[i + 1].clone());
856 let mut j = 0;
857 while j + 1 < sym.len() {
858 if sym[j] == left && sym[j + 1] == right {
859 sym[j] = merged.clone();
860 sym.remove(j + 1);
861 }
862 j += 1;
863 }
864 }
865 for t in &sym {
866 if let Some(&id) = self.vocab.get(t) {
867 out.push(id);
868 } else {
869 let mut ok = true;
870 for byte in t.bytes() {
871 let tok = format!("<0x{:02X}>", byte);
872 match self.vocab.get(&tok) {
873 Some(&id) => out.push(id),
874 None => {
875 ok = false;
876 break;
877 }
878 }
879 }
880 if !ok {
881 tracing::error!("tokenizer: no id for SP symbol {t:?} — dropped");
882 }
883 }
884 }
885 }
886
887 fn bpe_piece(&self, piece: &str, out: &mut Vec<u32>) {
889 if piece.is_empty() {
890 return;
891 }
892 let mapped: Vec<String> = piece
893 .bytes()
894 .map(|b| self.byte_to_char[b as usize].to_string())
895 .collect();
896 let mut sym = mapped;
897 if let Some(sfx) = &self.eow_suffix {
901 if let Some(last) = sym.last_mut() {
902 last.push_str(sfx);
903 }
904 }
905
906 loop {
908 let mut best: Option<(u32, usize)> = None;
909 for i in 0..sym.len().saturating_sub(1) {
910 if let Some(&r) = self.ranks.get(&(sym[i].clone(), sym[i + 1].clone())) {
911 if best.map(|(br, _)| r < br).unwrap_or(true) {
912 best = Some((r, i));
913 }
914 }
915 }
916 let Some((_, i)) = best else { break };
917 let merged = format!("{}{}", sym[i], sym[i + 1]);
918 let (left, right) = (sym[i].clone(), sym[i + 1].clone());
920 let mut j = 0;
921 while j + 1 < sym.len() {
922 if sym[j] == left && sym[j + 1] == right {
923 sym[j] = merged.clone();
924 sym.remove(j + 1);
925 }
926 j += 1;
927 }
928 }
929
930 for s in &sym {
931 if let Some(&id) = self.vocab.get(s) {
932 out.push(id);
933 } else {
934 let mut ok = true;
936 for ch in s.chars() {
937 let Some(&b) = self.char_to_byte.get(&ch) else {
938 ok = false;
939 break;
940 };
941 let tok = format!("<0x{:02X}>", b);
942 if let Some(&id) = self.vocab.get(&tok) {
943 out.push(id);
944 } else {
945 ok = false;
946 break;
947 }
948 }
949 if !ok {
950 tracing::error!("tokenizer: no id for symbol {s:?} — dropped");
951 }
952 }
953 }
954 }
955
956 pub fn decode(&self, ids: &[u32]) -> String {
959 let mut bytes: Vec<u8> = Vec::new();
960 for &id in ids {
961 if self.special_ids.contains(&id) {
962 continue;
963 }
964 let idx = id as usize;
965 if idx >= self.id_to_token.len() {
966 continue;
967 }
968 let tok = &self.id_to_token[idx];
969 if self.added_ids.contains(&id) {
970 if self.metaspace && tok.contains('\u{2581}') {
973 bytes.extend_from_slice(tok.replace('\u{2581}', " ").as_bytes());
974 } else {
975 bytes.extend_from_slice(tok.as_bytes());
976 }
977 continue;
978 }
979 if tok.starts_with("<0x") && tok.ends_with('>') && tok.len() == 6 {
981 if let Ok(b) = u8::from_str_radix(&tok[3..5], 16) {
982 bytes.push(b);
983 continue;
984 }
985 }
986 if self.metaspace {
987 for ch in tok.chars() {
989 if ch == '\u{2581}' {
990 bytes.push(b' ');
991 } else {
992 let mut buf = [0u8; 4];
993 bytes.extend_from_slice(ch.encode_utf8(&mut buf).as_bytes());
994 }
995 }
996 continue;
997 }
998 for ch in tok.chars() {
999 match self.char_to_byte.get(&ch) {
1000 Some(&b) => bytes.push(b),
1001 None => {
1004 let mut buf = [0u8; 4];
1005 bytes.extend_from_slice(ch.encode_utf8(&mut buf).as_bytes());
1006 }
1007 }
1008 }
1009 }
1010 let text = String::from_utf8_lossy(&bytes).into_owned();
1011 if self.metaspace && (self.sp_prepend || self.sp_prepend_first) {
1012 if let Some(stripped) = text.strip_prefix(' ') {
1014 return stripped.to_string();
1015 }
1016 }
1017 text
1018 }
1019
1020 pub fn decode_token(&self, id: u32) -> String {
1023 if self.special_ids.contains(&id) {
1024 return String::new();
1025 }
1026 let idx = id as usize;
1027 if idx >= self.id_to_token.len() {
1028 return String::new();
1029 }
1030 let tok = &self.id_to_token[idx];
1031 if self.added_ids.contains(&id) {
1032 if self.metaspace && tok.contains('\u{2581}') {
1033 return tok.replace('\u{2581}', " ");
1034 }
1035 return tok.clone();
1036 }
1037 if tok.starts_with("<0x") && tok.ends_with('>') && tok.len() == 6 {
1038 if let Ok(b) = u8::from_str_radix(&tok[3..5], 16) {
1039 return String::from_utf8_lossy(&[b]).into_owned();
1040 }
1041 }
1042 if self.metaspace {
1043 return tok.replace('\u{2581}', " ");
1044 }
1045 let mut bytes = Vec::new();
1046 for ch in tok.chars() {
1047 match self.char_to_byte.get(&ch) {
1048 Some(&b) => bytes.push(b),
1049 None => {
1050 let mut buf = [0u8; 4];
1051 bytes.extend_from_slice(ch.encode_utf8(&mut buf).as_bytes());
1052 }
1053 }
1054 }
1055 String::from_utf8_lossy(&bytes).into_owned()
1056 }
1057
1058 pub fn decode_token_for_hash(&self, id: u32) -> String {
1061 let idx = id as usize;
1062 if idx >= self.id_to_token.len() {
1063 return String::new();
1064 }
1065 if self.special_ids.contains(&id) {
1066 return self.id_to_token[idx].clone();
1067 }
1068 self.decode_token(id)
1069 }
1070
1071 pub fn decode_for_protocol(&self, ids: &[u32]) -> String {
1074 let mut out = String::new();
1075 for &id in ids {
1076 let idx = id as usize;
1077 if self.special_ids.contains(&id) {
1078 if let Some(token) = self.id_to_token.get(idx) {
1079 out.push_str(token);
1080 }
1081 } else {
1082 out.push_str(&self.decode_token(id));
1083 }
1084 }
1085 out
1086 }
1087
1088 pub fn raw_token_for_hash(&self, id: u32) -> String {
1090 self.id_to_token
1091 .get(id as usize)
1092 .cloned()
1093 .unwrap_or_default()
1094 }
1095
1096 pub fn apply_chat_template(&self, messages: &[(String, String)]) -> Vec<u32> {
1100 self.apply_chat_template_opts(messages, None)
1101 }
1102
1103 pub fn apply_chat_template_json(
1120 &self,
1121 messages: &[serde_json::Value],
1122 tools: Option<&[serde_json::Value]>,
1123 enable_thinking: Option<bool>,
1124 ) -> Vec<u32> {
1125 self.apply_chat_template_json_with_reasoning_effort(messages, tools, enable_thinking, None)
1126 }
1127
1128 pub fn apply_chat_template_json_with_reasoning_effort(
1131 &self,
1132 messages: &[serde_json::Value],
1133 tools: Option<&[serde_json::Value]>,
1134 enable_thinking: Option<bool>,
1135 reasoning_effort: Option<&str>,
1136 ) -> Vec<u32> {
1137 match self.try_apply_chat_template_json_with_reasoning_effort(
1138 messages,
1139 tools,
1140 enable_thinking,
1141 reasoning_effort,
1142 ) {
1143 Ok(ids) => ids,
1144 Err(e) => {
1145 tracing::error!("chat template render failed ({e}); ChatML fallback");
1146 self.chatml_json_fallback(messages, enable_thinking)
1147 }
1148 }
1149 }
1150
1151 pub fn try_apply_chat_template_json(
1162 &self,
1163 messages: &[serde_json::Value],
1164 tools: Option<&[serde_json::Value]>,
1165 enable_thinking: Option<bool>,
1166 ) -> Result<Vec<u32>, String> {
1167 self.try_apply_chat_template_json_with_reasoning_effort(
1168 messages,
1169 tools,
1170 enable_thinking,
1171 None,
1172 )
1173 }
1174
1175 pub fn try_apply_chat_template_json_with_reasoning_effort(
1179 &self,
1180 messages: &[serde_json::Value],
1181 tools: Option<&[serde_json::Value]>,
1182 enable_thinking: Option<bool>,
1183 reasoning_effort: Option<&str>,
1184 ) -> Result<Vec<u32>, String> {
1185 if let Some(tpl) = &self.chat_template {
1186 return self
1187 .render_template_json(tpl, messages, tools, enable_thinking, reasoning_effort)
1188 .map(|text| self.with_bos(self.encode(&text)))
1189 .map_err(|e| format!("{e:#}"));
1190 }
1191 Ok(self.chatml_json_fallback(messages, enable_thinking))
1192 }
1193
1194 fn chatml_json_fallback(
1195 &self,
1196 messages: &[serde_json::Value],
1197 enable_thinking: Option<bool>,
1198 ) -> Vec<u32> {
1199 let pairs: Vec<(String, String)> = messages
1200 .iter()
1201 .map(|m| {
1202 (
1203 m.get("role")
1204 .and_then(|v| v.as_str())
1205 .unwrap_or("user")
1206 .to_string(),
1207 m.get("content")
1208 .and_then(|v| v.as_str())
1209 .unwrap_or("")
1210 .to_string(),
1211 )
1212 })
1213 .collect();
1214 self.with_bos(self.chatml_fallback_opts(&pairs, enable_thinking))
1215 }
1216
1217 pub fn render_chat_json(
1219 &self,
1220 messages: &[serde_json::Value],
1221 tools: Option<&[serde_json::Value]>,
1222 enable_thinking: Option<bool>,
1223 ) -> Option<String> {
1224 self.render_chat_json_with_reasoning_effort(messages, tools, enable_thinking, None)
1225 }
1226
1227 pub fn render_chat_json_with_reasoning_effort(
1230 &self,
1231 messages: &[serde_json::Value],
1232 tools: Option<&[serde_json::Value]>,
1233 enable_thinking: Option<bool>,
1234 reasoning_effort: Option<&str>,
1235 ) -> Option<String> {
1236 let tpl = self.chat_template.as_ref()?;
1237 match self.render_template_json(tpl, messages, tools, enable_thinking, reasoning_effort) {
1238 Ok(t) => Some(t),
1239 Err(e) => {
1240 tracing::error!("chat template render (json): {e:#}");
1241 eprintln!("chat template render (json): {e:#}");
1242 None
1243 }
1244 }
1245 }
1246
1247 fn render_template_json(
1248 &self,
1249 tpl: &str,
1250 messages: &[serde_json::Value],
1251 tools: Option<&[serde_json::Value]>,
1252 enable_thinking: Option<bool>,
1253 reasoning_effort: Option<&str>,
1254 ) -> Result<String, minijinja::Error> {
1255 let mut env = crate::chat_template::environment();
1256 let tpl_src = strip_generation_tags(tpl);
1257 env.add_template("chat", &tpl_src)?;
1258 let msgs: Vec<minijinja::Value> = messages
1259 .iter()
1260 .map(minijinja::Value::from_serialize)
1261 .collect();
1262 let tools_v: Option<Vec<minijinja::Value>> =
1263 tools.map(|ts| ts.iter().map(minijinja::Value::from_serialize).collect());
1264 let tpl = env.get_template("chat")?;
1265 let rendered = match (tools_v, enable_thinking, reasoning_effort) {
1272 (Some(ts), Some(v), Some(effort)) => tpl.render(minijinja::context! {
1273 messages => msgs, tools => ts, add_generation_prompt => true, enable_thinking => v,
1274 reasoning_effort => effort,
1275 tool_call_format => "json",
1276 })?,
1277 (Some(ts), Some(v), None) => tpl.render(minijinja::context! {
1278 messages => msgs, tools => ts, add_generation_prompt => true, enable_thinking => v,
1279 reasoning_effort => if !v { Some("low") } else { None::<&str> },
1284 tool_call_format => "json",
1285 })?,
1286 (Some(ts), None, Some(effort)) => tpl.render(minijinja::context! {
1287 messages => msgs, tools => ts, add_generation_prompt => true,
1288 reasoning_effort => effort,
1289 tool_call_format => "json",
1290 })?,
1291 (Some(ts), None, None) => tpl.render(minijinja::context! {
1292 messages => msgs, tools => ts, add_generation_prompt => true,
1293 tool_call_format => "json",
1294 })?,
1295 (None, Some(v), Some(effort)) => tpl.render(minijinja::context! {
1296 messages => msgs, add_generation_prompt => true, enable_thinking => v,
1297 reasoning_effort => effort,
1298 tool_call_format => "json",
1299 })?,
1300 (None, Some(v), None) => tpl.render(minijinja::context! {
1301 messages => msgs, add_generation_prompt => true, enable_thinking => v,
1302 reasoning_effort => if !v { Some("low") } else { None::<&str> },
1304 tool_call_format => "json",
1305 })?,
1306 (None, None, Some(effort)) => tpl.render(minijinja::context! {
1307 messages => msgs, add_generation_prompt => true,
1308 reasoning_effort => effort,
1309 tool_call_format => "json",
1310 })?,
1311 (None, None, None) => tpl.render(minijinja::context! {
1312 messages => msgs, add_generation_prompt => true,
1313 tool_call_format => "json",
1314 })?,
1315 };
1316 Ok(close_direct_think(rendered, enable_thinking))
1317 }
1318
1319 pub fn apply_chat_template_opts(
1320 &self,
1321 messages: &[(String, String)],
1322 enable_thinking: Option<bool>,
1323 ) -> Vec<u32> {
1324 if let Some(tpl) = &self.chat_template {
1325 match self.render_template(tpl, messages, enable_thinking) {
1326 Ok(text) => return self.with_bos(self.encode(&text)),
1327 Err(e) => {
1328 tracing::error!("chat template render failed ({e}); ChatML fallback");
1329 }
1330 }
1331 }
1332 self.with_bos(self.chatml_fallback_opts(messages, enable_thinking))
1333 }
1334
1335 pub fn with_bos(&self, mut ids: Vec<u32>) -> Vec<u32> {
1337 if self.add_bos {
1338 if let Some(b) = self.bos_token_id {
1339 if ids.first() != Some(&b) {
1340 ids.insert(0, b);
1341 }
1342 }
1343 }
1344 ids
1345 }
1346
1347 pub fn render_chat(&self, messages: &[(String, String)]) -> Option<String> {
1349 self.render_chat_opts(messages, None)
1350 }
1351
1352 pub fn render_chat_opts(
1354 &self,
1355 messages: &[(String, String)],
1356 enable_thinking: Option<bool>,
1357 ) -> Option<String> {
1358 let tpl = self.chat_template.as_ref()?;
1359 match self.render_template(tpl, messages, enable_thinking) {
1360 Ok(t) => Some(t),
1361 Err(e) => {
1362 tracing::error!("chat template render: {e:#}");
1363 None
1364 }
1365 }
1366 }
1367
1368 fn render_template(
1369 &self,
1370 tpl: &str,
1371 messages: &[(String, String)],
1372 enable_thinking: Option<bool>,
1373 ) -> Result<String, minijinja::Error> {
1374 let mut env = crate::chat_template::environment();
1375 let tpl_src = strip_generation_tags(tpl);
1376 env.add_template("chat", &tpl_src)?;
1377 let msgs: Vec<minijinja::Value> = messages
1378 .iter()
1379 .map(|(role, content)| {
1380 minijinja::context! { role => role, content => content }
1381 })
1382 .collect();
1383 let rendered = match enable_thinking {
1386 Some(v) => env.get_template("chat")?.render(minijinja::context! {
1387 messages => msgs,
1388 add_generation_prompt => true,
1389 enable_thinking => v,
1390 reasoning_effort => if !v { Some("low") } else { None::<&str> },
1393 })?,
1394 None => env.get_template("chat")?.render(minijinja::context! {
1395 messages => msgs,
1396 add_generation_prompt => true,
1397 })?,
1398 };
1399 Ok(close_direct_think(rendered, enable_thinking))
1400 }
1401
1402 fn chatml_fallback(&self, messages: &[(String, String)]) -> Vec<u32> {
1404 self.chatml_fallback_opts(messages, None)
1405 }
1406
1407 fn chatml_fallback_opts(
1409 &self,
1410 messages: &[(String, String)],
1411 enable_thinking: Option<bool>,
1412 ) -> Vec<u32> {
1413 let mut tokens = Vec::new();
1414
1415 for (role, content) in messages {
1416 if let Some(start_id) = self.im_start_id {
1418 tokens.push(start_id);
1419 }
1420 tokens.extend(self.encode(&format!("{}\n{}", role, content)));
1421 if let Some(end_id) = self.im_end_id {
1422 tokens.push(end_id);
1423 }
1424 tokens.extend(self.encode("\n"));
1425 }
1426
1427 if let Some(start_id) = self.im_start_id {
1429 tokens.push(start_id);
1430 }
1431 tokens.extend(self.encode("assistant\n"));
1432 if enable_thinking == Some(false) {
1433 tokens.extend(self.encode("<think>\n\n</think>\n\n"));
1434 }
1435
1436 tokens
1437 }
1438
1439 pub fn vocab_size(&self) -> usize {
1441 self.id_to_token.len()
1442 }
1443
1444 pub fn token_to_id(&self, token: &str) -> Option<u32> {
1448 self.vocab.get(token).copied()
1449 }
1450
1451 pub fn convert_tokens_to_ids(&self, token: &str) -> Option<u32> {
1454 self.token_to_id(token)
1455 }
1456
1457 pub fn is_eos(&self, id: u32) -> bool {
1459 self.eos_token_id == Some(id) || self.im_end_id == Some(id) || self.extra_eos.contains(&id)
1460 }
1461}
1462
1463#[derive(Debug, thiserror::Error)]
1464pub enum TokenizerError {
1465 #[error("IO error: {0}")]
1466 Io(String),
1467 #[error("Parse error: {0}")]
1468 Parse(String),
1469}
1470
1471struct SpMerges {
1473 pairs: HashMap<u64, (u32, u32)>,
1474 bytes: [u32; 256],
1476}
1477
1478const SP_NONE: u32 = u32::MAX;
1479
1480thread_local! {
1481 static SP_STRINGS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
1483}
1484
1485impl SpMerges {
1486 fn build(vocab: &HashMap<String, u32>, ranks: &HashMap<(String, String), u32>) -> Option<Self> {
1487 if ranks.is_empty() {
1488 return None;
1489 }
1490 let mut pairs = HashMap::with_capacity(ranks.len());
1491 let mut cat = String::new();
1492 for ((a, b), &r) in ranks {
1493 let (&ia, &ib) = (vocab.get(a)?, vocab.get(b)?);
1494 cat.clear();
1495 cat.push_str(a);
1496 cat.push_str(b);
1497 let &m = vocab.get(&cat)?;
1498 pairs.insert(((ia as u64) << 32) | ib as u64, (r, m));
1499 }
1500 let mut bytes = [SP_NONE; 256];
1501 for (b, slot) in bytes.iter_mut().enumerate() {
1502 if let Some(&id) = vocab.get(&format!("<0x{b:02X}>")) {
1503 *slot = id;
1504 }
1505 }
1506 Some(SpMerges { pairs, bytes })
1507 }
1508
1509 #[inline]
1510 fn rank(&self, a: u32, b: u32) -> Option<(u32, u32)> {
1511 if a == SP_NONE || b == SP_NONE {
1512 return None;
1513 }
1514 self.pairs.get(&(((a as u64) << 32) | b as u64)).copied()
1515 }
1516
1517 fn encode(&self, vocab: &HashMap<String, u32>, piece: &str, out: &mut Vec<u32>) {
1522 use std::cmp::Reverse;
1523 use std::collections::BinaryHeap;
1524 let spans: Vec<(usize, usize)> = piece
1525 .char_indices()
1526 .map(|(i, c)| (i, i + c.len_utf8()))
1527 .collect();
1528 let n = spans.len();
1529 let mut buf = [0u8; 4];
1530 let mut id: Vec<u32> = piece
1531 .chars()
1532 .map(|c| *vocab.get(c.encode_utf8(&mut buf) as &str).unwrap_or(&SP_NONE))
1533 .collect();
1534 let mut next: Vec<usize> = (1..=n).collect();
1535 let mut prev: Vec<usize> = (0..n).map(|i| i.wrapping_sub(1)).collect();
1536 let mut alive = vec![true; n];
1537 let mut heap: BinaryHeap<Reverse<(u32, usize)>> = BinaryHeap::with_capacity(n);
1538 for i in 0..n.saturating_sub(1) {
1539 if let Some((r, _)) = self.rank(id[i], id[i + 1]) {
1540 heap.push(Reverse((r, i)));
1541 }
1542 }
1543 let mut batch: Vec<usize> = Vec::new();
1544 let mut merged: Vec<usize> = Vec::new();
1545 while let Some(&Reverse((r, _))) = heap.peek() {
1546 batch.clear();
1547 while let Some(&Reverse((r2, pos))) = heap.peek() {
1548 if r2 != r {
1549 break;
1550 }
1551 heap.pop();
1552 batch.push(pos);
1553 }
1554 batch.sort_unstable();
1555 batch.dedup();
1556 merged.clear();
1557 for &pos in &batch {
1558 if !alive[pos] || next[pos] >= n {
1559 continue;
1560 }
1561 let j = next[pos];
1562 match self.rank(id[pos], id[j]) {
1563 Some((rr, m)) if rr == r => {
1564 id[pos] = m;
1565 alive[j] = false;
1566 next[pos] = next[j];
1567 if next[j] < n {
1568 prev[next[j]] = pos;
1569 }
1570 merged.push(pos);
1571 }
1572 _ => {}
1573 }
1574 }
1575 for &pos in &merged {
1576 if !alive[pos] {
1577 continue;
1578 }
1579 let pv = prev[pos];
1580 if pv < n && alive[pv] {
1581 if let Some((r2, _)) = self.rank(id[pv], id[pos]) {
1582 heap.push(Reverse((r2, pv)));
1583 }
1584 }
1585 let nx = next[pos];
1586 if nx < n {
1587 if let Some((r2, _)) = self.rank(id[pos], id[nx]) {
1588 heap.push(Reverse((r2, pos)));
1589 }
1590 }
1591 }
1592 }
1593 let mut i = 0usize;
1594 while i < n {
1595 if id[i] != SP_NONE {
1596 out.push(id[i]);
1597 } else {
1598 let (a, b) = spans[i];
1600 let mut ok = true;
1601 for &byte in piece[a..b].as_bytes() {
1602 match self.bytes[byte as usize] {
1603 SP_NONE => {
1604 ok = false;
1605 break;
1606 }
1607 t => out.push(t),
1608 }
1609 }
1610 if !ok {
1611 tracing::error!("tokenizer: no id for SP symbol {:?} — dropped", &piece[a..b]);
1612 }
1613 }
1614 i = next[i];
1615 }
1616 }
1617}
1618
1619#[cfg(test)]
1620mod tests {
1621 use super::*;
1622
1623 #[test]
1624 fn byte_unicode_bijection() {
1625 let (b2c, c2b) = bytes_to_unicode();
1626 for b in 0..=255u8 {
1627 assert_eq!(c2b[&b2c[b as usize]], b);
1628 }
1629 assert_eq!(b2c[b' ' as usize], 'Ġ');
1631 assert_eq!(b2c[b'\n' as usize], 'Ċ');
1632 }
1633
1634 #[test]
1635 fn byte_level_roundtrip_utf8() {
1636 let tok = Tokenizer::byte_level();
1637 let text = "hello 🌍 hi\n";
1638 let ids = tok.encode(text);
1639 assert_eq!(ids.len(), text.len()); assert_eq!(tok.decode(&ids), text);
1641 }
1642
1643 fn mini_json() -> String {
1646 let vocab: Vec<(&str, u32)> = vec![
1648 ("h", 0),
1649 ("e", 1),
1650 ("l", 2),
1651 ("o", 3),
1652 ("Ġ", 4),
1653 ("w", 5),
1654 ("r", 6),
1655 ("d", 7),
1656 ("he", 8),
1657 ("Ġw", 9),
1658 ];
1659 let vocab_json: String = vocab
1660 .iter()
1661 .map(|(t, i)| format!("\"{t}\": {i}"))
1662 .collect::<Vec<_>>()
1663 .join(", ");
1664 format!(
1665 r#"{{
1666 "model": {{
1667 "type": "BPE",
1668 "vocab": {{ {vocab_json} }},
1669 "merges": [["h", "e"], ["Ġ", "w"]]
1670 }},
1671 "added_tokens": [
1672 {{"id": 10, "content": "<|eot|>", "special": true}}
1673 ]
1674 }}"#
1675 )
1676 }
1677
1678 #[test]
1684 fn sp_heap_bpe_matches_the_string_form() {
1685 let toks = [
1686 "a", "b", "c", "▁", "bc", "abc", "ab", "aa", "aab", "ca", "▁a", "aabc",
1687 "aaa", "<0x64>", "<0xC3>", "▁▁",
1688 ];
1689 let vocab: Vec<String> = toks
1690 .iter()
1691 .enumerate()
1692 .map(|(i, t)| format!("\"{t}\": {i}"))
1693 .collect();
1694 let merges = [
1695 ("aab", "c"),
1696 ("b", "c"),
1697 ("a", "bc"),
1698 ("a", "b"),
1699 ("ab", "c"),
1700 ("a", "a"),
1701 ("aa", "b"),
1702 ("c", "a"),
1703 ("▁", "a"),
1704 ("aa", "a"),
1705 ("▁", "▁"),
1706 ];
1707 let merges: Vec<String> = merges
1708 .iter()
1709 .map(|(a, b)| format!("[\"{a}\", \"{b}\"]"))
1710 .collect();
1711 let json = format!(
1712 r#"{{"model": {{"type": "BPE", "byte_fallback": true, "vocab": {{ {} }}, "merges": [{}]}},
1713 "normalizer": {{"type": "Replace", "pattern": {{"String": " "}}, "content": "\u2581"}},
1714 "added_tokens": []}}"#,
1715 vocab.join(", "),
1716 merges.join(", ")
1717 );
1718 let t = Tokenizer::from_json(&json).unwrap();
1719 assert!(t.metaspace, "a byte_fallback vocabulary takes the SentencePiece path");
1720 assert!(t.sp_merges().is_some());
1721 let alphabet = ['a', 'b', 'c', '\u{2581}', 'd', 'é', 'a', 'a', 'b'];
1722 let mut state = 0x2545_f491_4f6c_dd1du64;
1723 for len in 0..40 {
1724 for _ in 0..60 {
1725 let piece: String = (0..len)
1726 .map(|_| {
1727 state ^= state << 13;
1728 state ^= state >> 7;
1729 state ^= state << 17;
1730 alphabet[(state % alphabet.len() as u64) as usize]
1731 })
1732 .collect();
1733 let (mut fast, mut slow) = (Vec::new(), Vec::new());
1734 t.sp_merges().unwrap().encode(&t.vocab, &piece, &mut fast);
1735 t.bpe_piece_sp_strings(&piece, &mut slow);
1736 assert_eq!(fast, slow, "{piece:?}");
1737 }
1738 }
1739 }
1740
1741 #[test]
1745 fn tool_markup_decodes_even_when_special() {
1746 let json = r#"{
1747 "model": {"type": "BPE", "vocab": {"h": 0, "e": 1, "l": 2, "o": 3}, "merges": []},
1748 "added_tokens": [
1749 {"id": 10, "content": "<|im_end|>", "special": true},
1750 {"id": 11, "content": "<function", "special": true},
1751 {"id": 12, "content": "</function>", "special": true},
1752 {"id": 13, "content": "<param", "special": true},
1753 {"id": 14, "content": "</param>", "special": true},
1754 {"id": 15, "content": "<tool_call>", "special": true}
1755 ]
1756 }"#;
1757 let t = Tokenizer::from_json(json).unwrap();
1758 let ids = [11, 0, 1, 13, 2, 14, 12, 15, 10];
1759 assert_eq!(
1760 t.decode(&ids),
1761 "<functionhe<paraml</param></function><tool_call>"
1762 );
1763 let streamed: String = ids.iter().map(|&i| t.decode_token(i)).collect();
1764 assert_eq!(streamed, t.decode(&ids), "streaming must agree with decode");
1765 assert!(
1766 !t.decode(&[10]).contains("im_end"),
1767 "control tokens stay hidden"
1768 );
1769 }
1770
1771 #[test]
1774 fn real_tokenizer_parity_when_available() {
1775 let Ok(path) = std::env::var("CMF_TOK_PARITY") else {
1776 return;
1777 };
1778 let t = Tokenizer::from_file(&path).expect("load");
1779 for (text, want) in [
1780 (
1781 "The capital of France is",
1782 vec![671u32, 6102, 294, 8760, 344],
1783 ),
1784 ("2 + 2 =", vec![20, 940, 223, 20, 438]),
1785 ] {
1786 let got = t.encode(text);
1787 assert_eq!(got, want, "«{text}»");
1788 }
1789 }
1790
1791 #[test]
1796 fn granite_42_chat_template_when_available() {
1797 let Ok(path) = std::env::var("CMF_GRANITE_CHAT_TEMPLATE") else {
1798 return;
1799 };
1800 let mut tok = Tokenizer::byte_level();
1801 tok.chat_template = Some(std::fs::read_to_string(path).expect("read Granite template"));
1802 let messages = vec![("user".to_string(), "Hello".to_string())];
1803
1804 let thinking = tok
1805 .render_chat_opts(&messages, Some(true))
1806 .expect("render Granite thinking prompt");
1807 assert_eq!(
1808 thinking,
1809 "<|im_start|>system\n<|im_end|>\n<|im_start|>user\nHello<|im_end|>\n<|im_start|>assistant\n<think>\n"
1810 );
1811
1812 let direct = tok
1813 .render_chat_opts(&messages, Some(false))
1814 .expect("render Granite direct prompt");
1815 assert_eq!(
1816 direct,
1817 "<|im_start|>system\n<|im_end|>\n<|im_start|>user\nHello<|im_end|>\n<|im_start|>assistant\n<think></think>"
1818 );
1819 }
1820
1821 #[test]
1822 fn json_template_passes_explicit_reasoning_effort_without_overriding_it() {
1823 let mut tok = Tokenizer::byte_level();
1824 tok.chat_template = Some(
1825 "{% if reasoning_effort is defined %}<effort>{{ reasoning_effort }}</effort>{% else %}<default-effort>{% endif %}|{{ enable_thinking }}|{{ messages[0].reasoning_content }}"
1826 .to_string(),
1827 );
1828 let messages = vec![serde_json::json!({
1829 "role": "assistant",
1830 "content": "answer",
1831 "reasoning_content": "prior chain"
1832 })];
1833 let explicit = tok
1834 .render_chat_json_with_reasoning_effort(&messages, None, Some(false), Some("high"))
1835 .expect("render explicit effort");
1836 assert_eq!(explicit, "<effort>high</effort>|False|prior chain");
1837
1838 let native_default = tok
1839 .render_chat_json_with_reasoning_effort(&messages, None, None, None)
1840 .expect("render native default");
1841 assert_eq!(native_default, "<default-effort>||prior chain");
1842 }
1843
1844 #[test]
1845 fn fallible_json_template_render_does_not_use_chatml_on_jinja_error() {
1846 let mut tok = Tokenizer::byte_level();
1847 tok.chat_template = Some("{% for message in messages %}".to_string());
1848 let messages = vec![serde_json::json!({"role": "user", "content": "hello"})];
1849 let error = tok
1850 .try_apply_chat_template_json_with_reasoning_effort(
1851 &messages,
1852 None,
1853 None,
1854 Some("high"),
1855 )
1856 .expect_err("unclosed Jinja must fail instead of producing ChatML");
1857 assert!(!error.is_empty());
1858 }
1859
1860 #[test]
1866 fn minimax_m3_chat_template_when_available() {
1867 let Ok(path) = std::env::var("CMF_MINIMAX_M3_CHAT_TEMPLATE") else {
1868 return;
1869 };
1870 let mut tok = Tokenizer::byte_level();
1871 tok.chat_template = Some(std::fs::read_to_string(path).expect("read MiniMax M3 template"));
1872 let messages = vec![
1873 serde_json::json!({"role": "system", "content": "Keep answers concise."}),
1874 serde_json::json!({"role": "user", "content": "What is 2 + 2?"}),
1875 ];
1876 let tools = vec![serde_json::json!({
1877 "type": "function",
1878 "function": {
1879 "name": "get_weather",
1880 "description": "Return the weather for a city.",
1881 "parameters": {
1882 "type": "object",
1883 "properties": {"city": {"type": "string"}},
1884 "required": ["city"]
1885 }
1886 }
1887 })];
1888 let rendered = tok
1889 .render_chat_json(&messages, Some(&tools), Some(false))
1890 .expect("render MiniMax M3 native prompt");
1891 assert!(rendered.starts_with("]~!b[]~b]system\n"), "{rendered:?}");
1892 assert!(rendered.contains("<effort>low</effort>"), "{rendered:?}");
1893 assert!(
1894 rendered.contains("]<]minimax[>[<tool_call>"),
1895 "{rendered:?}"
1896 );
1897 assert!(rendered.contains("get_weather"), "{rendered:?}");
1898 assert!(rendered.ends_with("]~b]ai\n"), "{rendered:?}");
1899 }
1900
1901 #[test]
1905 fn glm_direct_template_keeps_assistant_special_token_intact() {
1906 let mut tok = Tokenizer::byte_level();
1907 tok.chat_template = Some(
1908 "[gMASK]<sop>{%- set effort = reasoning_effort if reasoning_effort is defined and reasoning_effort in ['low', 'high'] else 'max' -%}<|system|>Reasoning Effort: {{ effort | capitalize }}{%- for m in messages -%}<|user|>{{ m.content }}{%- endfor -%}<|assistant|><think>"
1909 .to_string(),
1910 );
1911 let messages = vec![("user".to_string(), "2+2?".to_string())];
1912 let rendered = tok
1913 .render_chat_opts(&messages, Some(false))
1914 .expect("render GLM direct prompt");
1915 assert!(rendered.contains("Reasoning Effort: Low"));
1916 assert!(rendered.contains("<|assistant|><think></think>"));
1917 assert!(!rendered.contains("<|assistant\n"));
1918
1919 let thinking = tok
1922 .render_chat_opts(&messages, Some(true))
1923 .expect("render GLM thinking prompt");
1924 assert!(thinking.contains("Reasoning Effort: Max"));
1925
1926 let json_messages = vec![serde_json::json!({
1927 "role": "user",
1928 "content": "2+2?"
1929 })];
1930 let rendered_json = tok
1931 .render_chat_json(&json_messages, None, Some(false))
1932 .expect("render GLM JSON direct prompt");
1933 assert!(rendered_json.contains("Reasoning Effort: Low"));
1934 assert!(rendered_json.contains("<|assistant|><think></think>"));
1935 assert!(!rendered_json.contains("<|assistant\n"));
1936 }
1937
1938 #[test]
1944 fn every_split_in_a_sequence_is_applied() {
1945 let pt = serde_json::json!({
1946 "type": "Sequence",
1947 "pretokenizers": [
1948 {"type": "Split", "behavior": "Isolated",
1949 "pattern": {"Regex": r"\p{N}{1,3}"}},
1950 {"type": "Split", "behavior": "Isolated",
1951 "pattern": {"Regex": r" ?[\p{L}]+"}},
1952 {"type": "ByteLevel", "add_prefix_space": false, "use_regex": false}
1953 ]
1954 });
1955 let mut pats = Vec::new();
1956 collect_split_patterns(&pt, &mut pats);
1957 assert_eq!(
1958 pats.len(),
1959 2,
1960 "both Split stages must be collected: {pats:?}"
1961 );
1962 assert!(pats[0].contains("p{N}"), "digit rule first");
1963 assert!(pats[1].contains("p{L}"), "word rule second");
1964
1965 let re: Vec<fancy_regex::Regex> = pats
1969 .iter()
1970 .map(|p| fancy_regex::Regex::new(p).unwrap())
1971 .collect();
1972 let norm = "ab cd12";
1973 let mut pieces: Vec<(usize, usize)> = vec![(0, norm.len())];
1974 for r in &re {
1975 let mut next = Vec::new();
1976 for (ps, pe) in pieces {
1977 let seg = &norm[ps..pe];
1978 let mut last = 0;
1979 for m in r.find_iter(seg).flatten() {
1980 if m.start() > last {
1981 next.push((ps + last, ps + m.start()));
1982 }
1983 if m.end() > m.start() {
1984 next.push((ps + m.start(), ps + m.end()));
1985 }
1986 last = m.end();
1987 }
1988 if last < seg.len() {
1989 next.push((ps + last, pe));
1990 }
1991 }
1992 pieces = next;
1993 }
1994 let got: Vec<&str> = pieces.iter().map(|(a, b)| &norm[*a..*b]).collect();
1995 assert_eq!(
1996 got,
1997 vec!["ab", " cd", "12"],
1998 "staged split produced {got:?}"
1999 );
2000 }
2001
2002 #[test]
2003 fn full_pipeline_merges_and_added_tokens() {
2004 let tok = Tokenizer::from_json(&mini_json()).unwrap();
2005 let ids = tok.encode("hello world");
2007 assert_eq!(ids, vec![8, 2, 2, 3, 9, 3, 6, 2, 7]);
2008 assert_eq!(tok.decode(&ids), "hello world");
2009 let ids2 = tok.encode("he<|eot|>he");
2011 assert_eq!(ids2, vec![8, 10, 8]);
2012 assert_eq!(tok.decode(&ids2), "hehe");
2013 }
2014
2015 #[test]
2016 fn non_ascii_is_never_silently_dropped() {
2017 let tok = Tokenizer::from_json(&mini_json()).unwrap();
2018 let ids = tok.encode("hello");
2021 assert!(!ids.is_empty());
2022 }
2023}
2024
2025#[cfg(test)]
2026mod generation_tag_tests {
2027 use super::strip_generation_tags;
2028
2029 #[test]
2034 fn a_generation_block_becomes_a_no_op_keeping_its_whitespace_control() {
2035 let tpl = "a{%- generation -%}b{%- endgeneration -%}c";
2036 let out = strip_generation_tags(tpl);
2037 assert!(!out.contains("{%- generation"));
2038 assert!(!out.contains("endgeneration"));
2039 assert_eq!(out.matches("{%-").count(), 2);
2041 assert_eq!(out.matches("-%}").count(), 2);
2042 assert!(out.starts_with('a') && out.ends_with('c'));
2043 }
2044
2045 #[test]
2048 fn each_side_keeps_its_own_dash() {
2049 let out = strip_generation_tags("{% generation %}x{%- endgeneration %}");
2050 assert!(out.starts_with("{% set"), "no dash added on the left");
2051 assert!(out.contains("{%- set"), "the right tag keeps its dash");
2052 assert!(!out.contains("-%}"), "no trailing dash invented");
2053 }
2054
2055 #[test]
2058 fn everything_else_is_left_alone() {
2059 let plain = "{%- if x -%}{{ y }}{%- endif -%}";
2060 assert_eq!(strip_generation_tags(plain), plain);
2061 let prose = "{{ 'the generation of tokens' }}";
2063 assert_eq!(strip_generation_tags(prose), prose);
2064 }
2065}