1pub mod chat;
13mod json;
14mod unicode;
15mod unicode_data;
16
17pub use chat::apply_chat_template_str;
18
19use memra_gguf::{GgufFile, MetaValue};
20use std::cmp::Ordering;
21use std::collections::{BinaryHeap, HashMap};
22
23const TT_UNKNOWN: i64 = 2;
25const TT_CONTROL: i64 = 3;
26const TT_USER_DEFINED: i64 = 4;
27const TT_BYTE: i64 = 6;
28const QWEN35_PRETOKENIZE_REGEX: &str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?[\p{L}\p{M}]+|\p{N}| ?[^\s\p{L}\p{M}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+";
29const QWEN2_PRETOKENIZE_REGEX: &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+";
33const DEEPSEEK_V3_SPLIT_REGEXES: [&str; 3] = [
38 r"\p{N}{1,3}",
39 "[\u{4e00}-\u{9fa5}\u{3040}-\u{309f}\u{30a0}-\u{30ff}]+",
40 "[!\"#$%&'()*+,\\-./:;<=>?@\\[\\\\\\]^_`{|}~][A-Za-z]+|[^\r\n\\p{L}\\p{P}\\p{S}]?[\\p{L}\\p{M}]+| ?[\\p{P}\\p{S}]+[\r\n]*|\\s*[\r\n]+|\\s+(?!\\S)|\\s+",
41];
42
43pub const SUPPORTED_PRETOKENIZERS: &[&str] = &["qwen35", "qwen2", "deepseek-v3", "gemma4"];
46
47pub const ALLOW_UNKNOWN_PRETOKENIZER_ENV: &str = "MEMRA_ALLOW_UNKNOWN_PRETOKENIZER";
50
51fn allow_unknown_pretokenizer() -> bool {
52 std::env::var(ALLOW_UNKNOWN_PRETOKENIZER_ENV).as_deref() == Ok("1")
53}
54
55#[derive(Debug, Clone, PartialEq, Eq)]
63pub struct UnknownPretokenizer {
64 pub pre: String,
67 pub spm_style: bool,
70}
71
72impl std::fmt::Display for UnknownPretokenizer {
73 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74 write!(
75 f,
76 "unsupported tokenizer.ggml.pre '{}' (vocab model is {}) — memra has no exact \
77 pre-tokenizer split for it and token ids would NOT be exact. Supported: {}. \
78 Set {}=1 to load anyway for deliberate experimentation (token ids will be wrong).",
79 self.pre,
80 if self.spm_style { "SPM/gemma4" } else { "gpt2" },
81 SUPPORTED_PRETOKENIZERS.join(", "),
82 ALLOW_UNKNOWN_PRETOKENIZER_ENV,
83 )
84 }
85}
86
87impl std::error::Error for UnknownPretokenizer {}
88
89#[derive(Debug, Clone, Copy, PartialEq, Eq)]
93pub enum PreSplit {
94 Qwen35,
96 DeepseekV3,
98 Spm,
101 UnknownFallbackQwen35,
104}
105
106impl PreSplit {
107 pub fn resolve(pre: &str, spm_style: bool) -> Result<Self, UnknownPretokenizer> {
110 Self::resolve_with(pre, spm_style, allow_unknown_pretokenizer())
111 }
112
113 fn resolve_with(
116 pre: &str,
117 spm_style: bool,
118 allow_unknown: bool,
119 ) -> Result<Self, UnknownPretokenizer> {
120 match (pre, spm_style) {
124 ("qwen35" | "qwen2", false) => Ok(PreSplit::Qwen35),
125 ("deepseek-v3", false) => Ok(PreSplit::DeepseekV3),
126 ("gemma4", true) => Ok(PreSplit::Spm),
127 _ => {
128 let err = UnknownPretokenizer {
129 pre: pre.to_string(),
130 spm_style,
131 };
132 if allow_unknown {
133 eprintln!(
136 "memra-tokenizer: WARNING {ALLOW_UNKNOWN_PRETOKENIZER_ENV}=1 — loading \
137 with {err} FALLING BACK to the qwen35 split. Token ids are NOT exact: \
138 goldens, parity fixtures, acceptance counts and quality numbers taken \
139 on this model are all invalid."
140 );
141 Ok(PreSplit::UnknownFallbackQwen35)
142 } else {
143 Err(err)
144 }
145 }
146 }
147 }
148}
149
150#[derive(Debug, Clone, Copy, PartialEq, Eq)]
151enum TokAttr {
152 Normal,
153 Unknown,
154 Control,
155 UserDefined,
156 Byte,
157 Other,
158}
159
160impl TokAttr {
161 fn from_toktype(t: i64) -> Self {
162 match t {
163 TT_UNKNOWN => TokAttr::Unknown,
164 TT_CONTROL => TokAttr::Control,
165 TT_USER_DEFINED => TokAttr::UserDefined,
166 TT_BYTE => TokAttr::Byte,
167 1 => TokAttr::Normal,
168 _ => TokAttr::Other,
169 }
170 }
171 fn is_special(self) -> bool {
174 matches!(
175 self,
176 TokAttr::Control | TokAttr::UserDefined | TokAttr::Unknown
177 )
178 }
179}
180
181pub struct Tokenizer {
182 id_to_token: Vec<String>,
184 token_to_id: HashMap<String, u32>,
186 attrs: Vec<TokAttr>,
188 bpe_ranks: HashMap<(String, String), i32>,
190 special_tokens: Vec<u32>,
192 eos_id: u32,
193 bos_id: Option<u32>,
194 add_bos: bool,
195 pre: String,
196 split: PreSplit,
199 chat_template: Option<String>,
200 spm_style: bool,
202 dsv4_encoding: Option<chat::Dsv4Encoding>,
209}
210
211#[derive(Clone, Eq, PartialEq)]
216struct Bigram {
217 left: i32,
218 right: i32,
219 rank: i32,
220 text: String,
221}
222
223impl Ord for Bigram {
224 fn cmp(&self, other: &Self) -> Ordering {
225 match other.rank.cmp(&self.rank) {
228 Ordering::Equal => other.left.cmp(&self.left),
229 o => o,
230 }
231 }
232}
233impl PartialOrd for Bigram {
234 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
235 Some(self.cmp(other))
236 }
237}
238
239struct Symbol {
241 text: String,
242 prev: i32,
243 next: i32,
244 n: usize, }
246
247impl Tokenizer {
248 pub fn from_gguf(g: &GgufFile) -> Result<Self, String> {
250 let model = g
251 .metadata
252 .get("tokenizer.ggml.model")
253 .and_then(|v| v.as_str())
254 .ok_or("missing tokenizer.ggml.model")?;
255 if model != "gpt2" && model != "gemma4" {
256 return Err(format!(
257 "unsupported tokenizer model '{model}' (only gpt2/gemma4)"
258 ));
259 }
260 let spm_style = model == "gemma4";
264 let pre = g
265 .metadata
266 .get("tokenizer.ggml.pre")
267 .and_then(|v| v.as_str())
268 .unwrap_or(if spm_style { "gemma4" } else { "default" })
269 .to_string();
270 let split = PreSplit::resolve(&pre, spm_style).map_err(|e| e.to_string())?;
273
274 let tokens = match g.metadata.get("tokenizer.ggml.tokens") {
276 Some(MetaValue::Array(a)) => a,
277 _ => return Err("missing tokenizer.ggml.tokens array".into()),
278 };
279 let n = tokens.len();
280 let mut id_to_token = Vec::with_capacity(n);
281 let mut token_to_id = HashMap::with_capacity(n);
282 for (i, t) in tokens.iter().enumerate() {
283 let s = t.as_str().ok_or("non-string in tokens[]")?.to_string();
284 token_to_id.entry(s.clone()).or_insert(i as u32);
286 id_to_token.push(s);
287 }
288
289 let mut attrs = vec![TokAttr::Normal; n];
291 if let Some(MetaValue::Array(a)) = g.metadata.get("tokenizer.ggml.token_type") {
292 for (i, v) in a.iter().enumerate().take(n) {
293 if let Some(t) = v.as_u64() {
294 attrs[i] = TokAttr::from_toktype(t as i64);
295 } else if let MetaValue::I32(t) = v {
296 attrs[i] = TokAttr::from_toktype(*t as i64);
297 }
298 }
299 }
300
301 let mut bpe_ranks = HashMap::new();
303 if let Some(MetaValue::Array(a)) = g.metadata.get("tokenizer.ggml.merges") {
304 for (i, v) in a.iter().enumerate() {
305 let word = v.as_str().ok_or("non-string in merges[]")?;
306 let bytes = word.as_bytes();
310 if let Some(pos) = bytes.iter().skip(1).position(|&b| b == b' ').map(|p| p + 1) {
311 let first = word[..pos].to_string();
312 let second = word[pos + 1..].to_string();
313 bpe_ranks.insert((first, second), i as i32);
314 }
315 }
316 } else {
317 return Err("missing tokenizer.ggml.merges array".into());
318 }
319
320 let mut special_tokens: Vec<u32> = (0..n as u32)
322 .filter(|&id| attrs[id as usize].is_special())
323 .collect();
324 special_tokens.sort_by(|&a, &b| {
325 id_to_token[b as usize]
326 .len()
327 .cmp(&id_to_token[a as usize].len())
328 });
329
330 let eos_id = g
331 .metadata
332 .get("tokenizer.ggml.eos_token_id")
333 .and_then(|v| v.as_u64())
334 .map(|v| v as u32)
335 .ok_or("missing tokenizer.ggml.eos_token_id")?;
336 let bos_id = g
337 .metadata
338 .get("tokenizer.ggml.bos_token_id")
339 .and_then(|v| v.as_u64())
340 .map(|v| v as u32);
341 let add_bos = g
342 .metadata
343 .get("tokenizer.ggml.add_bos_token")
344 .and_then(|v| match v {
345 MetaValue::Bool(b) => Some(*b),
346 _ => v.as_u64().map(|x| x != 0),
347 })
348 .unwrap_or(false);
349 let add_bos = add_bos || spm_style;
350
351 let chat_template = g
352 .metadata
353 .get("tokenizer.chat_template")
354 .and_then(|v| v.as_str())
355 .map(|s| s.to_string());
356
357 Ok(Tokenizer {
358 id_to_token,
359 token_to_id,
360 attrs,
361 bpe_ranks,
362 special_tokens,
363 eos_id,
364 bos_id,
365 add_bos,
366 pre,
367 split,
368 chat_template,
369 spm_style,
370 dsv4_encoding: None,
374 })
375 }
376
377 pub fn from_hf_dir(dir: &std::path::Path) -> Result<Self, String> {
395 let tj_path = dir.join("tokenizer.json");
396 let text = std::fs::read_to_string(&tj_path)
397 .map_err(|e| format!("read {}: {e}", tj_path.display()))?;
398 let tj = json::parse(&text).map_err(|e| format!("{}: {e}", tj_path.display()))?;
399
400 let model = tj.get("model").ok_or("tokenizer.json: missing model")?;
401 if let Some(t) = model.get("type").and_then(|v| v.as_str()) {
402 if t != "BPE" {
403 return Err(format!(
404 "unsupported tokenizer.json model type '{t}' (only BPE)"
405 ));
406 }
407 }
408 let pre_tok = tj
410 .get("pre_tokenizer")
411 .ok_or("tokenizer.json: missing pre_tokenizer")?;
412 if !pre_tokenizer_is_byte_level(pre_tok) {
413 return Err(
414 "tokenizer.json: pre_tokenizer is not ByteLevel — only byte-level \
415 BPE is supported"
416 .into(),
417 );
418 }
419
420 let vocab = model
422 .get("vocab")
423 .and_then(|v| v.as_obj())
424 .ok_or("tokenizer.json: missing model.vocab")?;
425 let empty: Vec<json::Value> = Vec::new();
426 let added = tj
427 .get("added_tokens")
428 .and_then(|v| v.as_arr())
429 .unwrap_or(&empty);
430 let mut max_id = 0u32;
431 for v in vocab.values() {
432 let id =
433 v.as_u64()
434 .ok_or("tokenizer.json: non-integer id in model.vocab")? as u32;
435 max_id = max_id.max(id);
436 }
437 for a in added {
438 if let Some(id) = a.get("id").and_then(|v| v.as_u64()) {
439 max_id = max_id.max(id as u32);
440 }
441 }
442 let n = max_id as usize + 1;
443 let mut id_to_token = vec![String::new(); n];
444 let mut token_to_id: HashMap<String, u32> = HashMap::with_capacity(n);
445 let mut attrs = vec![TokAttr::Normal; n];
446 for (tok, v) in vocab {
447 let id = v.as_u64().unwrap() as u32;
448 id_to_token[id as usize] = tok.clone();
449 token_to_id.entry(tok.clone()).or_insert(id);
450 }
451 for a in added {
454 let id =
455 a.get("id")
456 .and_then(|v| v.as_u64())
457 .ok_or("tokenizer.json: added_tokens entry missing id")? as u32;
458 let content = a
459 .get("content")
460 .and_then(|v| v.as_str())
461 .ok_or("tokenizer.json: added_tokens entry missing content")?;
462 if id_to_token[id as usize].is_empty() {
463 id_to_token[id as usize] = content.to_string();
464 }
465 token_to_id.entry(content.to_string()).or_insert(id);
466 if a.get("special").and_then(|v| v.as_bool()).unwrap_or(false) {
467 attrs[id as usize] = TokAttr::Control;
468 } else {
469 attrs[id as usize] = TokAttr::UserDefined;
475 }
476 }
477
478 let merges = model
480 .get("merges")
481 .and_then(|v| v.as_arr())
482 .ok_or("tokenizer.json: missing model.merges")?;
483 let mut bpe_ranks = HashMap::with_capacity(merges.len());
484 for (i, m) in merges.iter().enumerate() {
485 let (first, second) = match m {
486 json::Value::Str(s) => {
487 let bytes = s.as_bytes();
490 let pos = bytes
491 .iter()
492 .skip(1)
493 .position(|&b| b == b' ')
494 .map(|p| p + 1)
495 .ok_or_else(|| format!("tokenizer.json: merges[{i}] has no space"))?;
496 (s[..pos].to_string(), s[pos + 1..].to_string())
497 }
498 json::Value::Arr(a) if a.len() == 2 => {
499 let f = a[0]
500 .as_str()
501 .ok_or_else(|| format!("tokenizer.json: merges[{i}] non-string pair"))?;
502 let s2 = a[1]
503 .as_str()
504 .ok_or_else(|| format!("tokenizer.json: merges[{i}] non-string pair"))?;
505 (f.to_string(), s2.to_string())
506 }
507 _ => {
508 return Err(format!(
509 "tokenizer.json: merges[{i}] is neither \"a b\" string nor [a, b] pair"
510 ));
511 }
512 };
513 bpe_ranks.insert((first, second), i as i32);
514 }
515
516 let mut special_tokens: Vec<u32> = (0..n as u32)
518 .filter(|&id| attrs[id as usize].is_special())
519 .collect();
520 special_tokens.sort_by(|&a, &b| {
521 id_to_token[b as usize]
522 .len()
523 .cmp(&id_to_token[a as usize].len())
524 });
525
526 let tc = std::fs::read_to_string(dir.join("tokenizer_config.json"))
528 .ok()
529 .and_then(|t| json::parse(&t).ok());
530 let gc = std::fs::read_to_string(dir.join("generation_config.json"))
531 .ok()
532 .and_then(|t| json::parse(&t).ok());
533
534 let tok_content = |v: &json::Value| -> Option<String> {
536 v.as_str().map(|s| s.to_string()).or_else(|| {
537 v.get("content")
538 .and_then(|c| c.as_str())
539 .map(|s| s.to_string())
540 })
541 };
542 let eos_from_cfg = tc
543 .as_ref()
544 .and_then(|c| c.get("eos_token"))
545 .and_then(&tok_content)
546 .and_then(|s| token_to_id.get(&s).copied());
547 let eos_from_gen = gc
549 .as_ref()
550 .and_then(|c| c.get("eos_token_id"))
551 .and_then(|v| match v {
552 json::Value::Num(_) => v.as_u64(),
553 json::Value::Arr(a) => a.first().and_then(|x| x.as_u64()),
554 _ => None,
555 })
556 .map(|v| v as u32);
557 let eos_id = eos_from_cfg.or(eos_from_gen).ok_or(
558 "no eos token: need tokenizer_config.json eos_token or \
559 generation_config.json eos_token_id",
560 )?;
561 let bos_id = tc
562 .as_ref()
563 .and_then(|c| c.get("bos_token"))
564 .and_then(&tok_content)
565 .and_then(|s| token_to_id.get(&s).copied());
566 let add_bos = tc
567 .as_ref()
568 .and_then(|c| c.get("add_bos_token"))
569 .and_then(|v| v.as_bool())
570 .unwrap_or(false);
571
572 let chat_template = tc
574 .as_ref()
575 .and_then(|c| c.get("chat_template"))
576 .and_then(|v| v.as_str())
577 .map(|s| s.to_string())
578 .or_else(|| std::fs::read_to_string(dir.join("chat_template.jinja")).ok());
579 let cfg_regex = tc
587 .as_ref()
588 .and_then(|c| c.get("pretokenize_regex"))
589 .and_then(|v| v.as_str());
590 let mut tj_regexes: Vec<String> = Vec::new();
591 collect_split_regexes(pre_tok, &mut tj_regexes);
592 let pre = cfg_regex
593 .and_then(|r| pre_from_split_regexes(std::slice::from_ref(&r.to_string())))
594 .or_else(|| pre_from_split_regexes(&tj_regexes))
595 .unwrap_or("default");
596 let split = PreSplit::resolve(pre, false).map_err(|e| {
597 if pre == "default" {
598 format!(
599 "{e}\n (HF checkpoint {}: tokenizer_config.json pretokenize_regex = {:?}, \
600 tokenizer.json pre_tokenizer Split regexes = {:?} — neither matched a known \
601 family)",
602 dir.display(),
603 cfg_regex,
604 tj_regexes,
605 )
606 } else {
607 e.to_string()
608 }
609 })?;
610
611 let dsv4_encoding = dsv4_encoding_from_config(dir)?;
617
618 Ok(Tokenizer {
619 id_to_token,
620 token_to_id,
621 attrs,
622 bpe_ranks,
623 special_tokens,
624 eos_id,
625 bos_id,
626 add_bos,
627 pre: pre.to_string(),
628 split,
629 chat_template,
630 spm_style: false,
631 dsv4_encoding,
632 })
633 }
634
635 pub fn eos_id(&self) -> u32 {
636 self.eos_id
637 }
638 pub fn id_of(&self, piece: &str) -> Option<u32> {
640 self.token_to_id.get(piece).copied()
641 }
642 pub fn eog_ids(&self) -> Vec<u32> {
645 let mut ids = vec![self.eos_id];
646 for t in ["<|im_end|>", "<turn|>", "<end_of_turn>"] {
647 if let Some(&id) = self.token_to_id.get(t) {
648 if !ids.contains(&id) {
649 ids.push(id);
650 }
651 }
652 }
653 ids
654 }
655 pub fn bos_id(&self) -> Option<u32> {
656 self.bos_id
657 }
658 pub fn vocab_size(&self) -> usize {
659 self.id_to_token.len()
660 }
661 pub fn pre(&self) -> &str {
662 &self.pre
663 }
664 pub fn split(&self) -> PreSplit {
667 self.split
668 }
669 pub fn chat_template(&self) -> Option<&str> {
670 self.chat_template.as_deref()
671 }
672 pub fn dsv4_encoding(&self) -> Option<chat::Dsv4Encoding> {
675 self.dsv4_encoding
676 }
677
678 #[inline]
679 fn text_to_token(&self, s: &str) -> Option<u32> {
680 self.token_to_id.get(s).copied()
681 }
682
683 fn find_bpe_rank(&self, left: &str, right: &str) -> i32 {
684 self.bpe_ranks
685 .get(&(left.to_string(), right.to_string()))
686 .copied()
687 .unwrap_or(-1)
688 }
689
690 pub fn encode(&self, text: &str, add_special: bool) -> Vec<u32> {
696 self.encode_special(text, add_special, true)
697 }
698
699 pub fn encode_special(&self, text: &str, add_special: bool, parse_special: bool) -> Vec<u32> {
700 let mut output: Vec<u32> = Vec::new();
701 if add_special && self.add_bos {
702 if let Some(b) = self.bos_id {
703 output.push(b);
704 }
705 }
706 if text.is_empty() {
707 return output;
708 }
709
710 for frag in self.st_partition(text, parse_special) {
712 match frag {
713 Fragment::Token(id) => output.push(id),
714 Fragment::Text(span) => self.bpe_tokenize(&span, &mut output),
715 }
716 }
717 output
718 }
719
720 fn st_partition(&self, text: &str, parse_special: bool) -> Vec<Fragment> {
722 let mut frags = vec![Fragment::Text(text.to_string())];
723 for &sid in &self.special_tokens {
724 let attr = self.attrs[sid as usize];
725 if !parse_special && matches!(attr, TokAttr::Control | TokAttr::Unknown) {
727 continue;
728 }
729 let needle = &self.id_to_token[sid as usize];
730 if needle.is_empty() {
731 continue;
732 }
733 let mut next: Vec<Fragment> = Vec::with_capacity(frags.len());
734 for f in frags.drain(..) {
735 match f {
736 Fragment::Token(id) => next.push(Fragment::Token(id)),
737 Fragment::Text(s) => {
738 let mut rest: &str = &s;
739 let mut acc = String::new();
740 while let Some(m) = rest.find(needle.as_str()) {
741 acc.push_str(&rest[..m]);
742 if !acc.is_empty() {
743 next.push(Fragment::Text(std::mem::take(&mut acc)));
744 }
745 next.push(Fragment::Token(sid));
746 rest = &rest[m + needle.len()..];
747 }
748 acc.push_str(rest);
749 if !acc.is_empty() {
750 next.push(Fragment::Text(acc));
751 }
752 }
753 }
754 }
755 frags = next;
756 }
757 frags
758 }
759
760 fn bpe_tokenize(&self, text: &str, output: &mut Vec<u32>) {
762 if self.spm_style {
763 let escaped: String = text
766 .chars()
767 .map(|c| if c == ' ' { '\u{2581}' } else { c })
768 .collect();
769 let mut words: Vec<String> = Vec::new();
770 let mut cur = String::new();
771 let mut cur_nl: Option<bool> = None;
772 for c in escaped.chars() {
773 let nl = c == '\n';
774 if cur_nl != Some(nl) && !cur.is_empty() {
775 words.push(std::mem::take(&mut cur));
776 }
777 cur_nl = Some(nl);
778 cur.push(c);
779 }
780 if !cur.is_empty() {
781 words.push(cur);
782 }
783 for word in &words {
784 if word.chars().all(|c| c == '\n') {
786 if let Some(tok) = self.text_to_token(word) {
787 output.push(tok);
788 continue;
789 }
790 }
791 self.bpe_merge_word(word, output);
792 }
793 return;
794 }
795 let words: Vec<String> = match self.split {
801 PreSplit::Qwen35 => unicode::split_qwen35(text),
804 PreSplit::DeepseekV3 => unicode::split_deepseek_v3(text),
808 PreSplit::UnknownFallbackQwen35 => unicode::split_qwen35(text),
811 PreSplit::Spm => unreachable!("PreSplit::Spm implies spm_style, handled above"),
814 };
815
816 for word in &words {
817 let word = unicode::byte_encode(word);
818 self.bpe_merge_word(&word, output);
819 }
820 }
821
822 fn bpe_merge_word(&self, word: &str, output: &mut Vec<u32>) {
825 {
826 let word = word.to_string();
827
828 let chars: Vec<char> = word.chars().collect();
830 let mut symbols: Vec<Symbol> = Vec::with_capacity(chars.len());
831 for (i, &c) in chars.iter().enumerate() {
832 symbols.push(Symbol {
833 text: c.to_string(),
834 prev: i as i32 - 1,
835 next: if i + 1 == chars.len() {
836 -1
837 } else {
838 i as i32 + 1
839 },
840 n: 1,
841 });
842 }
843
844 let mut queue: BinaryHeap<Bigram> = BinaryHeap::new();
846 for i in 1..symbols.len() {
847 self.add_bigram(&symbols, i as i32 - 1, i as i32, &mut queue);
848 }
849
850 while let Some(bigram) = queue.pop() {
852 let li = bigram.left as usize;
853 let ri = bigram.right as usize;
854 if symbols[li].n == 0 || symbols[ri].n == 0 {
855 continue;
856 }
857 let combined = format!("{}{}", symbols[li].text, symbols[ri].text);
858 if combined != bigram.text {
859 continue; }
861 symbols[li].text = combined;
863 symbols[li].n += symbols[ri].n;
864 symbols[ri].n = 0;
865 let r_next = symbols[ri].next;
866 symbols[li].next = r_next;
867 if r_next >= 0 {
868 symbols[r_next as usize].prev = bigram.left;
869 }
870 let l_prev = symbols[li].prev;
871 let l_next = symbols[li].next;
872 self.add_bigram(&symbols, l_prev, bigram.left, &mut queue);
873 self.add_bigram(&symbols, bigram.left, l_next, &mut queue);
874 }
875
876 for sym in &symbols {
878 if sym.n == 0 {
879 continue;
880 }
881 match self.text_to_token(&sym.text) {
882 Some(tok) => output.push(tok),
883 None => {
884 for b in sym.text.bytes() {
886 let bs = if self.spm_style {
887 format!("<0x{b:02X}>") } else {
889 (b as char).to_string()
890 };
891 if let Some(t) = self.text_to_token(&bs) {
892 output.push(t);
893 }
894 }
895 }
896 }
897 }
898 }
899 }
900
901 fn add_bigram(
902 &self,
903 symbols: &[Symbol],
904 left: i32,
905 right: i32,
906 queue: &mut BinaryHeap<Bigram>,
907 ) {
908 if left == -1 || right == -1 {
909 return;
910 }
911 let lt = &symbols[left as usize].text;
912 let rt = &symbols[right as usize].text;
913 let rank = self.find_bpe_rank(lt, rt);
914 if rank < 0 {
915 return;
916 }
917 queue.push(Bigram {
918 left,
919 right,
920 rank,
921 text: format!("{lt}{rt}"),
922 });
923 }
924
925 pub fn decode(&self, ids: &[u32]) -> String {
928 self.decode_special(ids, true)
929 }
930
931 pub fn token_is_control(&self, id: u32) -> bool {
936 match self.attrs.get(id as usize) {
937 Some(TokAttr::Control) | Some(TokAttr::Unknown) => true,
938 _ => false,
939 }
940 }
941
942 pub fn decode_special(&self, ids: &[u32], special: bool) -> String {
943 String::from_utf8_lossy(&self.decode_bytes_special(ids, special)).into_owned()
944 }
945
946 pub fn decode_bytes_special(&self, ids: &[u32], special: bool) -> Vec<u8> {
949 let mut bytes: Vec<u8> = Vec::new();
950 for &id in ids {
951 let i = id as usize;
952 if i >= self.id_to_token.len() {
953 continue;
954 }
955 let attr = self.attrs[i];
956 let piece = &self.id_to_token[i];
957 match attr {
958 TokAttr::Normal | TokAttr::Byte => {
959 if self.spm_style {
960 if matches!(attr, TokAttr::Byte)
962 || (piece.len() == 6
963 && piece.starts_with("<0x")
964 && piece.ends_with('>'))
965 {
966 if let Ok(b) = u8::from_str_radix(&piece[3..5], 16) {
967 bytes.push(b);
968 continue;
969 }
970 }
971 for c in piece.chars() {
972 if c == '\u{2581}' {
973 bytes.push(b' ');
974 } else {
975 let mut buf = [0u8; 4];
976 bytes.extend_from_slice(c.encode_utf8(&mut buf).as_bytes());
977 }
978 }
979 } else {
980 self.piece_to_bytes(piece, &mut bytes);
982 }
983 }
984 TokAttr::UserDefined => {
985 bytes.extend_from_slice(piece.as_bytes());
987 }
988 TokAttr::Control | TokAttr::Unknown => {
989 if special {
990 bytes.extend_from_slice(piece.as_bytes());
991 }
992 }
994 TokAttr::Other => {}
995 }
996 }
997 bytes
998 }
999
1000 fn piece_to_bytes(&self, piece: &str, out: &mut Vec<u8>) {
1001 for c in piece.chars() {
1002 match unicode::unicode_to_byte(c) {
1003 Some(b) => out.push(b),
1004 None => {
1005 let mut buf = [0u8; 4];
1007 out.extend_from_slice(c.encode_utf8(&mut buf).as_bytes());
1008 }
1009 }
1010 }
1011 }
1012
1013 pub fn apply_chat_template(
1016 &self,
1017 messages: &[(&str, &str)],
1018 add_generation_prompt: bool,
1019 ) -> String {
1020 chat::apply_chat_template_enc(
1021 self.chat_template.as_deref(),
1022 messages,
1023 add_generation_prompt,
1024 self.dsv4_encoding,
1025 )
1026 .expect("plain chat render cannot fail")
1029 }
1030
1031 pub fn apply_chat_template_tools(
1037 &self,
1038 turns: &[chat::Turn],
1039 add_generation_prompt: bool,
1040 tools_json: &[String],
1041 think: chat::ThinkMode,
1042 reasoning_effort: Option<&str>,
1043 ) -> Result<String, String> {
1044 chat::apply_chat_template_tools_ex(
1045 self.chat_template.as_deref(),
1046 turns,
1047 add_generation_prompt,
1048 tools_json,
1049 &[],
1050 think,
1051 reasoning_effort,
1052 self.dsv4_encoding,
1053 )
1054 }
1055
1056 #[allow(clippy::too_many_arguments)]
1061 pub fn apply_chat_template_tools_ex(
1062 &self,
1063 turns: &[chat::Turn],
1064 add_generation_prompt: bool,
1065 tools_json: &[String],
1066 tools_struct: &[chat::Val],
1067 think: chat::ThinkMode,
1068 reasoning_effort: Option<&str>,
1069 ) -> Result<String, String> {
1070 chat::apply_chat_template_tools_ex(
1071 self.chat_template.as_deref(),
1072 turns,
1073 add_generation_prompt,
1074 tools_json,
1075 tools_struct,
1076 think,
1077 reasoning_effort,
1078 self.dsv4_encoding,
1079 )
1080 }
1081}
1082
1083enum Fragment {
1084 Text(String),
1085 Token(u32),
1086}
1087
1088fn dsv4_encoding_from_config(dir: &std::path::Path) -> Result<Option<chat::Dsv4Encoding>, String> {
1103 const DSPARK_KEYS: [&str; 4] = [
1104 "dspark_block_size",
1105 "dspark_markov_rank",
1106 "dspark_noise_token_id",
1107 "dspark_target_layer_ids",
1108 ];
1109 let cfg_path = dir.join("config.json");
1110 let Ok(text) = std::fs::read_to_string(&cfg_path) else {
1111 return Ok(None);
1112 };
1113 let Ok(cfg) = json::parse(&text) else {
1114 return Ok(None);
1117 };
1118 if cfg.get("model_type").and_then(|v| v.as_str()) != Some("deepseek_v4") {
1119 return Ok(None);
1122 }
1123 let present: Vec<&str> = DSPARK_KEYS
1124 .iter()
1125 .copied()
1126 .filter(|k| cfg.get(k).is_some())
1127 .collect();
1128 match present.len() {
1129 0 => Ok(Some(chat::Dsv4Encoding::Preview)),
1130 4 => Ok(Some(chat::Dsv4Encoding::V0731)),
1131 _ => Err(format!(
1132 "{}: partial dspark_* key set {:?} (expected none or all of {:?}) — cannot \
1133 determine the deepseek-v4 encoding revision; refusing rather than guessing \
1134 the reasoning-effort ladder",
1135 cfg_path.display(),
1136 present,
1137 DSPARK_KEYS
1138 )),
1139 }
1140}
1141
1142fn collect_split_regexes(pt: &json::Value, out: &mut Vec<String>) {
1148 match pt.get("type").and_then(|v| v.as_str()) {
1149 Some("Sequence") => {
1150 if let Some(arr) = pt.get("pretokenizers").and_then(|v| v.as_arr()) {
1151 for step in arr {
1152 collect_split_regexes(step, out);
1153 }
1154 }
1155 }
1156 Some("Split") => {
1157 if let Some(r) = pt
1158 .get("pattern")
1159 .and_then(|p| p.get("Regex"))
1160 .and_then(|v| v.as_str())
1161 {
1162 out.push(r.to_string());
1163 }
1164 }
1165 _ => {}
1166 }
1167}
1168
1169fn pre_from_split_regexes(regexes: &[String]) -> Option<&'static str> {
1174 match regexes {
1175 [one] if one == QWEN35_PRETOKENIZE_REGEX => Some("qwen35"),
1176 [one] if one == QWEN2_PRETOKENIZE_REGEX => Some("qwen2"),
1177 [a, b, c]
1178 if a == DEEPSEEK_V3_SPLIT_REGEXES[0]
1179 && b == DEEPSEEK_V3_SPLIT_REGEXES[1]
1180 && c == DEEPSEEK_V3_SPLIT_REGEXES[2] =>
1181 {
1182 Some("deepseek-v3")
1183 }
1184 _ => None,
1185 }
1186}
1187
1188fn pre_tokenizer_is_byte_level(pt: &json::Value) -> bool {
1189 match pt.get("type").and_then(|v| v.as_str()) {
1190 Some("ByteLevel") => true,
1191 Some("Sequence") => pt
1192 .get("pretokenizers")
1193 .and_then(|v| v.as_arr())
1194 .map(|arr| arr.iter().any(pre_tokenizer_is_byte_level))
1195 .unwrap_or(false),
1196 _ => false,
1197 }
1198}
1199
1200#[cfg(test)]
1201mod pretokenizer_tests {
1202 use super::*;
1203
1204 #[test]
1208 fn every_supported_pre_resolves() {
1209 assert_eq!(
1210 PreSplit::resolve_with("qwen35", false, false),
1211 Ok(PreSplit::Qwen35)
1212 );
1213 assert_eq!(
1214 PreSplit::resolve_with("qwen2", false, false),
1215 Ok(PreSplit::Qwen35)
1216 );
1217 assert_eq!(
1218 PreSplit::resolve_with("deepseek-v3", false, false),
1219 Ok(PreSplit::DeepseekV3)
1220 );
1221 assert_eq!(
1222 PreSplit::resolve_with("gemma4", true, false),
1223 Ok(PreSplit::Spm)
1224 );
1225 assert_eq!(
1228 SUPPORTED_PRETOKENIZERS,
1229 &["qwen35", "qwen2", "deepseek-v3", "gemma4"]
1230 );
1231 }
1232
1233 #[test]
1235 fn unknown_pre_is_a_typed_error() {
1236 let err =
1237 PreSplit::resolve_with("llama4", false, false).expect_err("llama4 has no ported split");
1238 assert_eq!(
1239 err,
1240 UnknownPretokenizer {
1241 pre: "llama4".into(),
1242 spm_style: false
1243 }
1244 );
1245 let msg = err.to_string();
1246 assert!(msg.contains("'llama4'"), "{msg}");
1248 for supported in SUPPORTED_PRETOKENIZERS {
1249 assert!(
1250 msg.contains(supported),
1251 "error must list {supported}: {msg}"
1252 );
1253 }
1254 assert!(msg.contains(ALLOW_UNKNOWN_PRETOKENIZER_ENV), "{msg}");
1255 let _: &dyn std::error::Error = &err;
1257 }
1258
1259 #[test]
1262 fn pre_and_vocab_model_must_agree() {
1263 assert!(PreSplit::resolve_with("qwen35", true, false).is_err());
1264 assert!(PreSplit::resolve_with("gemma4", false, false).is_err());
1265 assert!(PreSplit::resolve_with("default", false, false).is_err());
1267 assert!(PreSplit::resolve_with("", false, false).is_err());
1268 }
1269
1270 #[test]
1272 fn opt_out_loads_with_a_fallback_marker() {
1273 assert_eq!(
1274 PreSplit::resolve_with("llama4", false, true),
1275 Ok(PreSplit::UnknownFallbackQwen35)
1276 );
1277 assert_eq!(
1279 PreSplit::resolve_with("qwen35", true, true),
1280 Ok(PreSplit::UnknownFallbackQwen35)
1281 );
1282 }
1283
1284 #[test]
1287 fn opt_out_env_gate() {
1288 unsafe { std::env::remove_var(ALLOW_UNKNOWN_PRETOKENIZER_ENV) };
1291 assert!(!allow_unknown_pretokenizer());
1292 unsafe { std::env::set_var(ALLOW_UNKNOWN_PRETOKENIZER_ENV, "0") };
1293 assert!(!allow_unknown_pretokenizer());
1294 unsafe { std::env::set_var(ALLOW_UNKNOWN_PRETOKENIZER_ENV, "1") };
1295 assert!(allow_unknown_pretokenizer());
1296 assert_eq!(
1297 PreSplit::resolve("llama4", false),
1298 Ok(PreSplit::UnknownFallbackQwen35)
1299 );
1300 unsafe { std::env::remove_var(ALLOW_UNKNOWN_PRETOKENIZER_ENV) };
1301 assert!(PreSplit::resolve("llama4", false).is_err());
1302 }
1303
1304 #[test]
1307 fn split_regex_identification_is_exact() {
1308 let s = |v: &[&str]| v.iter().map(|x| x.to_string()).collect::<Vec<_>>();
1309 assert_eq!(
1310 pre_from_split_regexes(&s(&[QWEN35_PRETOKENIZE_REGEX])),
1311 Some("qwen35")
1312 );
1313 assert_eq!(
1314 pre_from_split_regexes(&s(&[QWEN2_PRETOKENIZE_REGEX])),
1315 Some("qwen2")
1316 );
1317 assert_eq!(
1318 pre_from_split_regexes(&s(&DEEPSEEK_V3_SPLIT_REGEXES)),
1319 Some("deepseek-v3")
1320 );
1321 assert_eq!(
1323 pre_from_split_regexes(&s(&[
1324 DEEPSEEK_V3_SPLIT_REGEXES[1],
1325 DEEPSEEK_V3_SPLIT_REGEXES[0],
1326 DEEPSEEK_V3_SPLIT_REGEXES[2],
1327 ])),
1328 None
1329 );
1330 assert_eq!(
1332 pre_from_split_regexes(&s(&[
1333 DEEPSEEK_V3_SPLIT_REGEXES[0],
1334 DEEPSEEK_V3_SPLIT_REGEXES[1]
1335 ])),
1336 None
1337 );
1338 let mut near = QWEN35_PRETOKENIZE_REGEX.to_string();
1340 near.push('x');
1341 assert_eq!(pre_from_split_regexes(&s(&[&near])), None);
1342 assert_eq!(pre_from_split_regexes(&[]), None);
1343 assert_ne!(QWEN2_PRETOKENIZE_REGEX, QWEN35_PRETOKENIZE_REGEX);
1345 }
1346
1347 #[test]
1350 fn collect_split_regexes_walks_in_order() {
1351 let src = r#"{"type":"Sequence","pretokenizers":[
1352 {"type":"Split","pattern":{"Regex":"A"},"behavior":"Isolated"},
1353 {"type":"Split","pattern":{"String":" "},"behavior":"Isolated"},
1354 {"type":"Digits","individual_digits":true},
1355 {"type":"Sequence","pretokenizers":[
1356 {"type":"Split","pattern":{"Regex":"B"},"behavior":"Isolated"}
1357 ]},
1358 {"type":"ByteLevel","add_prefix_space":false}
1359 ]}"#;
1360 let v = json::parse(src).unwrap();
1361 let mut out = Vec::new();
1362 collect_split_regexes(&v, &mut out);
1363 assert_eq!(out, vec!["A".to_string(), "B".to_string()]);
1364 }
1365}
1366
1367#[cfg(test)]
1368mod hf_tests {
1369 use super::*;
1370
1371 const TOKENIZER_JSON: &str = r#"{
1380 "version": "1.0",
1381 "added_tokens": [
1382 {"id": 15, "content": "<|end|>", "special": true},
1383 {"id": 16, "content": "<think>", "special": false}
1384 ],
1385 "pre_tokenizer": {
1386 "type": "Sequence",
1387 "pretokenizers": [
1388 {"type": "Split", "pattern": {"Regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"}, "behavior": "Isolated"},
1389 {"type": "ByteLevel", "add_prefix_space": false, "trim_offsets": false}
1390 ]
1391 },
1392 "model": {
1393 "type": "BPE",
1394 "vocab": {
1395 "h": 0, "e": 1, "l": 2, "o": 3, "Ġ": 4, "w": 5, "r": 6, "d": 7,
1396 "he": 8, "ll": 9, "hell": 10, "hello": 11, "Ġw": 12, "or": 13, "!": 14
1397 },
1398 "merges": [
1399 "h e",
1400 ["l", "l"],
1401 "he ll",
1402 ["hell", "o"],
1403 ["Ġ", "w"],
1404 "o r"
1405 ]
1406 }
1407 }"#;
1408
1409 fn write_fixture(
1410 name: &str,
1411 tokenizer_config: Option<&str>,
1412 generation_config: Option<&str>,
1413 jinja: Option<&str>,
1414 ) -> std::path::PathBuf {
1415 let dir = std::env::temp_dir().join(format!("memra-tok-hf-{name}-{}", std::process::id()));
1416 let _ = std::fs::remove_dir_all(&dir);
1417 std::fs::create_dir_all(&dir).unwrap();
1418 std::fs::write(dir.join("tokenizer.json"), TOKENIZER_JSON).unwrap();
1419 if let Some(tc) = tokenizer_config {
1420 std::fs::write(dir.join("tokenizer_config.json"), tc).unwrap();
1421 }
1422 if let Some(gc) = generation_config {
1423 std::fs::write(dir.join("generation_config.json"), gc).unwrap();
1424 }
1425 if let Some(j) = jinja {
1426 std::fs::write(dir.join("chat_template.jinja"), j).unwrap();
1427 }
1428 dir
1429 }
1430
1431 #[test]
1432 fn hf_dir_encode_decode_roundtrip_and_specials() {
1433 let tc = r#"{
1435 "eos_token": {"content": "<|end|>", "lstrip": false},
1436 "add_bos_token": false,
1437 "pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
1438 "chat_template": "{{ messages }}<|end|>"
1439 }"#;
1440 let dir = write_fixture("full", Some(tc), None, None);
1441 let tok = Tokenizer::from_hf_dir(&dir).expect("from_hf_dir");
1442
1443 assert_eq!(tok.eos_id(), 15);
1444 assert_eq!(tok.bos_id(), None);
1445 assert_eq!(tok.pre(), "qwen35");
1446 assert_eq!(tok.vocab_size(), 17); assert_eq!(tok.chat_template(), Some("{{ messages }}<|end|>"));
1448
1449 let ids = tok.encode("hello world", true);
1453 assert_eq!(ids, vec![11, 12, 13, 2, 7]);
1454 assert_eq!(tok.decode(&ids), "hello world");
1455
1456 let ids = tok.encode("hello<|end|> world", true);
1458 assert_eq!(ids, vec![11, 15, 12, 13, 2, 7]);
1459 assert_eq!(tok.decode_special(&ids, true), "hello<|end|> world");
1461 assert_eq!(tok.decode_special(&ids, false), "hello world");
1462
1463 assert_eq!(tok.decode(&[16]), "<think>");
1465 let _ = std::fs::remove_dir_all(&dir);
1466 }
1467
1468 #[test]
1469 fn hf_dir_generation_config_eos_fallback_and_jinja() {
1470 let gc = r#"{"eos_token_id": [15, 14]}"#;
1473 let dir = write_fixture("genconf", None, Some(gc), Some("JINJA {{ messages }}"));
1474 let tok = Tokenizer::from_hf_dir(&dir).expect("from_hf_dir");
1475 assert_eq!(tok.eos_id(), 15);
1476 assert!(!tok.encode("hello", true).is_empty());
1477 assert_eq!(tok.chat_template(), Some("JINJA {{ messages }}"));
1478 let _ = std::fs::remove_dir_all(&dir);
1479 }
1480
1481 #[test]
1485 fn hf_dir_identifies_deepseek_v3_from_tokenizer_json() {
1486 let dsv3_pt = r##""pre_tokenizer": {
1487 "type": "Sequence",
1488 "pretokenizers": [
1489 {"type": "Split", "pattern": {"Regex": "\\p{N}{1,3}"}, "behavior": "Isolated"},
1490 {"type": "Split", "pattern": {"Regex": "[一-龥-ゟ゠-ヿ]+"}, "behavior": "Isolated"},
1491 {"type": "Split", "pattern": {"Regex": "[!\"#$%&'()*+,\\-./:;<=>?@\\[\\\\\\]^_`{|}~][A-Za-z]+|[^\r\n\\p{L}\\p{P}\\p{S}]?[\\p{L}\\p{M}]+| ?[\\p{P}\\p{S}]+[\r\n]*|\\s*[\r\n]+|\\s+(?!\\S)|\\s+"}, "behavior": "Isolated"},
1492 {"type": "ByteLevel", "add_prefix_space": false, "trim_offsets": true, "use_regex": false}
1493 ]
1494 },"##;
1495 let open = TOKENIZER_JSON.find(r#""pre_tokenizer""#).unwrap();
1497 let close = TOKENIZER_JSON.find(r#""model""#).unwrap();
1498 let json = format!(
1499 "{}{}\n {}",
1500 &TOKENIZER_JSON[..open],
1501 dsv3_pt,
1502 &TOKENIZER_JSON[close..]
1503 );
1504 let dir = std::env::temp_dir().join(format!("memra-tok-hf-dsv3-{}", std::process::id()));
1505 let _ = std::fs::remove_dir_all(&dir);
1506 std::fs::create_dir_all(&dir).unwrap();
1507 std::fs::write(dir.join("tokenizer.json"), &json).unwrap();
1508 std::fs::write(
1509 dir.join("generation_config.json"),
1510 r#"{"eos_token_id": 15}"#,
1511 )
1512 .unwrap();
1513 let tok = Tokenizer::from_hf_dir(&dir).expect("from_hf_dir");
1514 assert_eq!(tok.pre(), "deepseek-v3");
1515 assert_eq!(tok.split(), PreSplit::DeepseekV3);
1516 let _ = std::fs::remove_dir_all(&dir);
1517 }
1518
1519 #[test]
1522 fn hf_dir_identifies_qwen2_regex() {
1523 let json = TOKENIZER_JSON.replace(
1524 r"[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+",
1525 r"[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+",
1526 );
1527 assert_ne!(json, TOKENIZER_JSON, "the qwen2 substitution must apply");
1528 let dir = std::env::temp_dir().join(format!("memra-tok-hf-qwen2-{}", std::process::id()));
1529 let _ = std::fs::remove_dir_all(&dir);
1530 std::fs::create_dir_all(&dir).unwrap();
1531 std::fs::write(dir.join("tokenizer.json"), &json).unwrap();
1532 std::fs::write(
1533 dir.join("generation_config.json"),
1534 r#"{"eos_token_id": 15}"#,
1535 )
1536 .unwrap();
1537 let tok = Tokenizer::from_hf_dir(&dir).expect("from_hf_dir");
1538 assert_eq!(tok.pre(), "qwen2");
1539 assert_eq!(
1540 tok.split(),
1541 PreSplit::Qwen35,
1542 "qwen2 rides the qwen35 split"
1543 );
1544 let _ = std::fs::remove_dir_all(&dir);
1545 }
1546
1547 #[test]
1550 fn hf_dir_refuses_unidentifiable_pretokenizer() {
1551 let json = TOKENIZER_JSON.replace(r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|", "SOMETHING-ELSE|");
1552 assert_ne!(json, TOKENIZER_JSON);
1553 let dir = std::env::temp_dir().join(format!("memra-tok-hf-unk-{}", std::process::id()));
1554 let _ = std::fs::remove_dir_all(&dir);
1555 std::fs::create_dir_all(&dir).unwrap();
1556 std::fs::write(dir.join("tokenizer.json"), &json).unwrap();
1557 std::fs::write(
1558 dir.join("generation_config.json"),
1559 r#"{"eos_token_id": 15}"#,
1560 )
1561 .unwrap();
1562 let err = match Tokenizer::from_hf_dir(&dir) {
1563 Ok(_) => panic!("unidentifiable pre must refuse to load"),
1564 Err(e) => e,
1565 };
1566 assert!(
1567 err.contains("unsupported tokenizer.ggml.pre 'default'"),
1568 "{err}"
1569 );
1570 assert!(
1571 err.contains("SOMETHING-ELSE"),
1572 "error must quote the regex: {err}"
1573 );
1574 assert!(err.contains("MEMRA_ALLOW_UNKNOWN_PRETOKENIZER"), "{err}");
1575 let _ = std::fs::remove_dir_all(&dir);
1576 }
1577
1578 #[test]
1584 fn staged_checkpoints_resolve_their_own_pretokenizer() {
1585 let cases: &[(&str, &str)] = &[
1586 (
1588 "/data/ai-ml/hf-models/hy3-layer103p5-sparse-source",
1589 "deepseek-v3",
1590 ),
1591 ("/data/ai-ml/hf-models/qwen3-1.7b-blk128fp8-synth", "qwen2"),
1593 ("/data/ai-ml/hf-models/qwen35-9b-hf", "qwen35"),
1595 ];
1596 let mut ran = 0;
1597 for (path, want) in cases {
1598 let dir = std::path::Path::new(path);
1599 if !dir.join("tokenizer.json").exists() {
1600 eprintln!("skip: {path} not staged");
1601 continue;
1602 }
1603 let tok = Tokenizer::from_hf_dir(dir).unwrap_or_else(|e| panic!("{path}: {e}"));
1604 assert_eq!(tok.pre(), *want, "{path}");
1605 ran += 1;
1606 }
1607 eprintln!("staged_checkpoints_resolve_their_own_pretokenizer: {ran}/3 cases ran");
1608 }
1609
1610 #[test]
1611 fn hf_dir_rejects_non_byte_level() {
1612 let dir = std::env::temp_dir().join(format!("memra-tok-hf-nonbl-{}", std::process::id()));
1613 let _ = std::fs::remove_dir_all(&dir);
1614 std::fs::create_dir_all(&dir).unwrap();
1615 let bad = TOKENIZER_JSON.replace("\"ByteLevel\"", "\"Metaspace\"");
1616 std::fs::write(dir.join("tokenizer.json"), bad).unwrap();
1617 assert!(Tokenizer::from_hf_dir(&dir).is_err());
1618 let _ = std::fs::remove_dir_all(&dir);
1619 }
1620
1621 #[test]
1624 fn hf_dir_dsv4_encoding_detection() {
1625 let gc = r#"{"eos_token_id": [15]}"#;
1626 let full_dspark = r#""dspark_block_size": 5, "dspark_markov_rank": 256,
1627 "dspark_noise_token_id": 128799, "dspark_target_layer_ids": [40, 41, 42]"#;
1628
1629 let dir = write_fixture("dsv4-none", None, Some(gc), None);
1631 let tok = Tokenizer::from_hf_dir(&dir).unwrap();
1632 assert_eq!(tok.dsv4_encoding(), None);
1633 let _ = std::fs::remove_dir_all(&dir);
1634
1635 let dir = write_fixture("dsv4-preview", None, Some(gc), None);
1637 std::fs::write(
1638 dir.join("config.json"),
1639 r#"{"model_type": "deepseek_v4", "num_hidden_layers": 43}"#,
1640 )
1641 .unwrap();
1642 let tok = Tokenizer::from_hf_dir(&dir).unwrap();
1643 assert_eq!(tok.dsv4_encoding(), Some(chat::Dsv4Encoding::Preview));
1644 let _ = std::fs::remove_dir_all(&dir);
1645
1646 let dir = write_fixture("dsv4-0731", None, Some(gc), None);
1648 std::fs::write(
1649 dir.join("config.json"),
1650 format!(r#"{{"model_type": "deepseek_v4", {full_dspark}}}"#),
1651 )
1652 .unwrap();
1653 let tok = Tokenizer::from_hf_dir(&dir).unwrap();
1654 assert_eq!(tok.dsv4_encoding(), Some(chat::Dsv4Encoding::V0731));
1655 let _ = std::fs::remove_dir_all(&dir);
1656
1657 let dir = write_fixture("dsv4-partial", None, Some(gc), None);
1659 std::fs::write(
1660 dir.join("config.json"),
1661 r#"{"model_type": "deepseek_v4", "dspark_block_size": 5}"#,
1662 )
1663 .unwrap();
1664 let err = match Tokenizer::from_hf_dir(&dir) {
1665 Err(e) => e,
1666 Ok(_) => panic!("a partial dspark_* config must refuse the load"),
1667 };
1668 assert!(err.contains("partial dspark_*"), "{err}");
1669 let _ = std::fs::remove_dir_all(&dir);
1670
1671 let dir = write_fixture("dsv4-foreign", None, Some(gc), None);
1673 std::fs::write(dir.join("config.json"), r#"{"model_type": "qwen3"}"#).unwrap();
1674 let tok = Tokenizer::from_hf_dir(&dir).unwrap();
1675 assert_eq!(tok.dsv4_encoding(), None);
1676 let _ = std::fs::remove_dir_all(&dir);
1677 }
1678}