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 has_qwen_effort_ladder(&self) -> bool {
1039 self.chat_template
1040 .as_deref()
1041 .is_some_and(chat::template_has_qwen_effort)
1042 }
1043
1044 pub fn apply_chat_template_tools(
1050 &self,
1051 turns: &[chat::Turn],
1052 add_generation_prompt: bool,
1053 tools_json: &[String],
1054 think: chat::ThinkMode,
1055 reasoning_effort: Option<&str>,
1056 ) -> Result<String, String> {
1057 chat::apply_chat_template_tools_ex(
1058 self.chat_template.as_deref(),
1059 turns,
1060 add_generation_prompt,
1061 tools_json,
1062 &[],
1063 think,
1064 reasoning_effort,
1065 self.dsv4_encoding,
1066 )
1067 }
1068
1069 #[allow(clippy::too_many_arguments)]
1074 pub fn apply_chat_template_tools_ex(
1075 &self,
1076 turns: &[chat::Turn],
1077 add_generation_prompt: bool,
1078 tools_json: &[String],
1079 tools_struct: &[chat::Val],
1080 think: chat::ThinkMode,
1081 reasoning_effort: Option<&str>,
1082 ) -> Result<String, String> {
1083 chat::apply_chat_template_tools_ex(
1084 self.chat_template.as_deref(),
1085 turns,
1086 add_generation_prompt,
1087 tools_json,
1088 tools_struct,
1089 think,
1090 reasoning_effort,
1091 self.dsv4_encoding,
1092 )
1093 }
1094}
1095
1096enum Fragment {
1097 Text(String),
1098 Token(u32),
1099}
1100
1101fn dsv4_encoding_from_config(dir: &std::path::Path) -> Result<Option<chat::Dsv4Encoding>, String> {
1116 const DSPARK_KEYS: [&str; 4] = [
1117 "dspark_block_size",
1118 "dspark_markov_rank",
1119 "dspark_noise_token_id",
1120 "dspark_target_layer_ids",
1121 ];
1122 let cfg_path = dir.join("config.json");
1123 let Ok(text) = std::fs::read_to_string(&cfg_path) else {
1124 return Ok(None);
1125 };
1126 let Ok(cfg) = json::parse(&text) else {
1127 return Ok(None);
1130 };
1131 if cfg.get("model_type").and_then(|v| v.as_str()) != Some("deepseek_v4") {
1132 return Ok(None);
1135 }
1136 let present: Vec<&str> = DSPARK_KEYS
1137 .iter()
1138 .copied()
1139 .filter(|k| cfg.get(k).is_some())
1140 .collect();
1141 match present.len() {
1142 0 => Ok(Some(chat::Dsv4Encoding::Preview)),
1143 4 => Ok(Some(chat::Dsv4Encoding::V0731)),
1144 _ => Err(format!(
1145 "{}: partial dspark_* key set {:?} (expected none or all of {:?}) — cannot \
1146 determine the deepseek-v4 encoding revision; refusing rather than guessing \
1147 the reasoning-effort ladder",
1148 cfg_path.display(),
1149 present,
1150 DSPARK_KEYS
1151 )),
1152 }
1153}
1154
1155fn collect_split_regexes(pt: &json::Value, out: &mut Vec<String>) {
1161 match pt.get("type").and_then(|v| v.as_str()) {
1162 Some("Sequence") => {
1163 if let Some(arr) = pt.get("pretokenizers").and_then(|v| v.as_arr()) {
1164 for step in arr {
1165 collect_split_regexes(step, out);
1166 }
1167 }
1168 }
1169 Some("Split") => {
1170 if let Some(r) = pt
1171 .get("pattern")
1172 .and_then(|p| p.get("Regex"))
1173 .and_then(|v| v.as_str())
1174 {
1175 out.push(r.to_string());
1176 }
1177 }
1178 _ => {}
1179 }
1180}
1181
1182fn pre_from_split_regexes(regexes: &[String]) -> Option<&'static str> {
1187 match regexes {
1188 [one] if one == QWEN35_PRETOKENIZE_REGEX => Some("qwen35"),
1189 [one] if one == QWEN2_PRETOKENIZE_REGEX => Some("qwen2"),
1190 [a, b, c]
1191 if a == DEEPSEEK_V3_SPLIT_REGEXES[0]
1192 && b == DEEPSEEK_V3_SPLIT_REGEXES[1]
1193 && c == DEEPSEEK_V3_SPLIT_REGEXES[2] =>
1194 {
1195 Some("deepseek-v3")
1196 }
1197 _ => None,
1198 }
1199}
1200
1201fn pre_tokenizer_is_byte_level(pt: &json::Value) -> bool {
1202 match pt.get("type").and_then(|v| v.as_str()) {
1203 Some("ByteLevel") => true,
1204 Some("Sequence") => pt
1205 .get("pretokenizers")
1206 .and_then(|v| v.as_arr())
1207 .map(|arr| arr.iter().any(pre_tokenizer_is_byte_level))
1208 .unwrap_or(false),
1209 _ => false,
1210 }
1211}
1212
1213#[cfg(test)]
1214mod pretokenizer_tests {
1215 use super::*;
1216
1217 #[test]
1221 fn every_supported_pre_resolves() {
1222 assert_eq!(
1223 PreSplit::resolve_with("qwen35", false, false),
1224 Ok(PreSplit::Qwen35)
1225 );
1226 assert_eq!(
1227 PreSplit::resolve_with("qwen2", false, false),
1228 Ok(PreSplit::Qwen35)
1229 );
1230 assert_eq!(
1231 PreSplit::resolve_with("deepseek-v3", false, false),
1232 Ok(PreSplit::DeepseekV3)
1233 );
1234 assert_eq!(
1235 PreSplit::resolve_with("gemma4", true, false),
1236 Ok(PreSplit::Spm)
1237 );
1238 assert_eq!(
1241 SUPPORTED_PRETOKENIZERS,
1242 &["qwen35", "qwen2", "deepseek-v3", "gemma4"]
1243 );
1244 }
1245
1246 #[test]
1248 fn unknown_pre_is_a_typed_error() {
1249 let err =
1250 PreSplit::resolve_with("llama4", false, false).expect_err("llama4 has no ported split");
1251 assert_eq!(
1252 err,
1253 UnknownPretokenizer {
1254 pre: "llama4".into(),
1255 spm_style: false
1256 }
1257 );
1258 let msg = err.to_string();
1259 assert!(msg.contains("'llama4'"), "{msg}");
1261 for supported in SUPPORTED_PRETOKENIZERS {
1262 assert!(
1263 msg.contains(supported),
1264 "error must list {supported}: {msg}"
1265 );
1266 }
1267 assert!(msg.contains(ALLOW_UNKNOWN_PRETOKENIZER_ENV), "{msg}");
1268 let _: &dyn std::error::Error = &err;
1270 }
1271
1272 #[test]
1275 fn pre_and_vocab_model_must_agree() {
1276 assert!(PreSplit::resolve_with("qwen35", true, false).is_err());
1277 assert!(PreSplit::resolve_with("gemma4", false, false).is_err());
1278 assert!(PreSplit::resolve_with("default", false, false).is_err());
1280 assert!(PreSplit::resolve_with("", false, false).is_err());
1281 }
1282
1283 #[test]
1285 fn opt_out_loads_with_a_fallback_marker() {
1286 assert_eq!(
1287 PreSplit::resolve_with("llama4", false, true),
1288 Ok(PreSplit::UnknownFallbackQwen35)
1289 );
1290 assert_eq!(
1292 PreSplit::resolve_with("qwen35", true, true),
1293 Ok(PreSplit::UnknownFallbackQwen35)
1294 );
1295 }
1296
1297 #[test]
1300 fn opt_out_env_gate() {
1301 unsafe { std::env::remove_var(ALLOW_UNKNOWN_PRETOKENIZER_ENV) };
1304 assert!(!allow_unknown_pretokenizer());
1305 unsafe { std::env::set_var(ALLOW_UNKNOWN_PRETOKENIZER_ENV, "0") };
1306 assert!(!allow_unknown_pretokenizer());
1307 unsafe { std::env::set_var(ALLOW_UNKNOWN_PRETOKENIZER_ENV, "1") };
1308 assert!(allow_unknown_pretokenizer());
1309 assert_eq!(
1310 PreSplit::resolve("llama4", false),
1311 Ok(PreSplit::UnknownFallbackQwen35)
1312 );
1313 unsafe { std::env::remove_var(ALLOW_UNKNOWN_PRETOKENIZER_ENV) };
1314 assert!(PreSplit::resolve("llama4", false).is_err());
1315 }
1316
1317 #[test]
1320 fn split_regex_identification_is_exact() {
1321 let s = |v: &[&str]| v.iter().map(|x| x.to_string()).collect::<Vec<_>>();
1322 assert_eq!(
1323 pre_from_split_regexes(&s(&[QWEN35_PRETOKENIZE_REGEX])),
1324 Some("qwen35")
1325 );
1326 assert_eq!(
1327 pre_from_split_regexes(&s(&[QWEN2_PRETOKENIZE_REGEX])),
1328 Some("qwen2")
1329 );
1330 assert_eq!(
1331 pre_from_split_regexes(&s(&DEEPSEEK_V3_SPLIT_REGEXES)),
1332 Some("deepseek-v3")
1333 );
1334 assert_eq!(
1336 pre_from_split_regexes(&s(&[
1337 DEEPSEEK_V3_SPLIT_REGEXES[1],
1338 DEEPSEEK_V3_SPLIT_REGEXES[0],
1339 DEEPSEEK_V3_SPLIT_REGEXES[2],
1340 ])),
1341 None
1342 );
1343 assert_eq!(
1345 pre_from_split_regexes(&s(&[
1346 DEEPSEEK_V3_SPLIT_REGEXES[0],
1347 DEEPSEEK_V3_SPLIT_REGEXES[1]
1348 ])),
1349 None
1350 );
1351 let mut near = QWEN35_PRETOKENIZE_REGEX.to_string();
1353 near.push('x');
1354 assert_eq!(pre_from_split_regexes(&s(&[&near])), None);
1355 assert_eq!(pre_from_split_regexes(&[]), None);
1356 assert_ne!(QWEN2_PRETOKENIZE_REGEX, QWEN35_PRETOKENIZE_REGEX);
1358 }
1359
1360 #[test]
1363 fn collect_split_regexes_walks_in_order() {
1364 let src = r#"{"type":"Sequence","pretokenizers":[
1365 {"type":"Split","pattern":{"Regex":"A"},"behavior":"Isolated"},
1366 {"type":"Split","pattern":{"String":" "},"behavior":"Isolated"},
1367 {"type":"Digits","individual_digits":true},
1368 {"type":"Sequence","pretokenizers":[
1369 {"type":"Split","pattern":{"Regex":"B"},"behavior":"Isolated"}
1370 ]},
1371 {"type":"ByteLevel","add_prefix_space":false}
1372 ]}"#;
1373 let v = json::parse(src).unwrap();
1374 let mut out = Vec::new();
1375 collect_split_regexes(&v, &mut out);
1376 assert_eq!(out, vec!["A".to_string(), "B".to_string()]);
1377 }
1378}
1379
1380#[cfg(test)]
1381mod hf_tests {
1382 use super::*;
1383
1384 const TOKENIZER_JSON: &str = r#"{
1393 "version": "1.0",
1394 "added_tokens": [
1395 {"id": 15, "content": "<|end|>", "special": true},
1396 {"id": 16, "content": "<think>", "special": false}
1397 ],
1398 "pre_tokenizer": {
1399 "type": "Sequence",
1400 "pretokenizers": [
1401 {"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"},
1402 {"type": "ByteLevel", "add_prefix_space": false, "trim_offsets": false}
1403 ]
1404 },
1405 "model": {
1406 "type": "BPE",
1407 "vocab": {
1408 "h": 0, "e": 1, "l": 2, "o": 3, "Ġ": 4, "w": 5, "r": 6, "d": 7,
1409 "he": 8, "ll": 9, "hell": 10, "hello": 11, "Ġw": 12, "or": 13, "!": 14
1410 },
1411 "merges": [
1412 "h e",
1413 ["l", "l"],
1414 "he ll",
1415 ["hell", "o"],
1416 ["Ġ", "w"],
1417 "o r"
1418 ]
1419 }
1420 }"#;
1421
1422 fn write_fixture(
1423 name: &str,
1424 tokenizer_config: Option<&str>,
1425 generation_config: Option<&str>,
1426 jinja: Option<&str>,
1427 ) -> std::path::PathBuf {
1428 let dir = std::env::temp_dir().join(format!("memra-tok-hf-{name}-{}", std::process::id()));
1429 let _ = std::fs::remove_dir_all(&dir);
1430 std::fs::create_dir_all(&dir).unwrap();
1431 std::fs::write(dir.join("tokenizer.json"), TOKENIZER_JSON).unwrap();
1432 if let Some(tc) = tokenizer_config {
1433 std::fs::write(dir.join("tokenizer_config.json"), tc).unwrap();
1434 }
1435 if let Some(gc) = generation_config {
1436 std::fs::write(dir.join("generation_config.json"), gc).unwrap();
1437 }
1438 if let Some(j) = jinja {
1439 std::fs::write(dir.join("chat_template.jinja"), j).unwrap();
1440 }
1441 dir
1442 }
1443
1444 #[test]
1445 fn hf_dir_encode_decode_roundtrip_and_specials() {
1446 let tc = r#"{
1448 "eos_token": {"content": "<|end|>", "lstrip": false},
1449 "add_bos_token": false,
1450 "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+",
1451 "chat_template": "{{ messages }}<|end|>"
1452 }"#;
1453 let dir = write_fixture("full", Some(tc), None, None);
1454 let tok = Tokenizer::from_hf_dir(&dir).expect("from_hf_dir");
1455
1456 assert_eq!(tok.eos_id(), 15);
1457 assert_eq!(tok.bos_id(), None);
1458 assert_eq!(tok.pre(), "qwen35");
1459 assert_eq!(tok.vocab_size(), 17); assert_eq!(tok.chat_template(), Some("{{ messages }}<|end|>"));
1461
1462 let ids = tok.encode("hello world", true);
1466 assert_eq!(ids, vec![11, 12, 13, 2, 7]);
1467 assert_eq!(tok.decode(&ids), "hello world");
1468
1469 let ids = tok.encode("hello<|end|> world", true);
1471 assert_eq!(ids, vec![11, 15, 12, 13, 2, 7]);
1472 assert_eq!(tok.decode_special(&ids, true), "hello<|end|> world");
1474 assert_eq!(tok.decode_special(&ids, false), "hello world");
1475
1476 assert_eq!(tok.decode(&[16]), "<think>");
1478 let _ = std::fs::remove_dir_all(&dir);
1479 }
1480
1481 #[test]
1482 fn hf_dir_generation_config_eos_fallback_and_jinja() {
1483 let gc = r#"{"eos_token_id": [15, 14]}"#;
1486 let dir = write_fixture("genconf", None, Some(gc), Some("JINJA {{ messages }}"));
1487 let tok = Tokenizer::from_hf_dir(&dir).expect("from_hf_dir");
1488 assert_eq!(tok.eos_id(), 15);
1489 assert!(!tok.encode("hello", true).is_empty());
1490 assert_eq!(tok.chat_template(), Some("JINJA {{ messages }}"));
1491 let _ = std::fs::remove_dir_all(&dir);
1492 }
1493
1494 #[test]
1498 fn hf_dir_identifies_deepseek_v3_from_tokenizer_json() {
1499 let dsv3_pt = r##""pre_tokenizer": {
1500 "type": "Sequence",
1501 "pretokenizers": [
1502 {"type": "Split", "pattern": {"Regex": "\\p{N}{1,3}"}, "behavior": "Isolated"},
1503 {"type": "Split", "pattern": {"Regex": "[一-龥-ゟ゠-ヿ]+"}, "behavior": "Isolated"},
1504 {"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"},
1505 {"type": "ByteLevel", "add_prefix_space": false, "trim_offsets": true, "use_regex": false}
1506 ]
1507 },"##;
1508 let open = TOKENIZER_JSON.find(r#""pre_tokenizer""#).unwrap();
1510 let close = TOKENIZER_JSON.find(r#""model""#).unwrap();
1511 let json = format!(
1512 "{}{}\n {}",
1513 &TOKENIZER_JSON[..open],
1514 dsv3_pt,
1515 &TOKENIZER_JSON[close..]
1516 );
1517 let dir = std::env::temp_dir().join(format!("memra-tok-hf-dsv3-{}", std::process::id()));
1518 let _ = std::fs::remove_dir_all(&dir);
1519 std::fs::create_dir_all(&dir).unwrap();
1520 std::fs::write(dir.join("tokenizer.json"), &json).unwrap();
1521 std::fs::write(
1522 dir.join("generation_config.json"),
1523 r#"{"eos_token_id": 15}"#,
1524 )
1525 .unwrap();
1526 let tok = Tokenizer::from_hf_dir(&dir).expect("from_hf_dir");
1527 assert_eq!(tok.pre(), "deepseek-v3");
1528 assert_eq!(tok.split(), PreSplit::DeepseekV3);
1529 let _ = std::fs::remove_dir_all(&dir);
1530 }
1531
1532 #[test]
1535 fn hf_dir_identifies_qwen2_regex() {
1536 let json = TOKENIZER_JSON.replace(
1537 r"[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+",
1538 r"[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+",
1539 );
1540 assert_ne!(json, TOKENIZER_JSON, "the qwen2 substitution must apply");
1541 let dir = std::env::temp_dir().join(format!("memra-tok-hf-qwen2-{}", std::process::id()));
1542 let _ = std::fs::remove_dir_all(&dir);
1543 std::fs::create_dir_all(&dir).unwrap();
1544 std::fs::write(dir.join("tokenizer.json"), &json).unwrap();
1545 std::fs::write(
1546 dir.join("generation_config.json"),
1547 r#"{"eos_token_id": 15}"#,
1548 )
1549 .unwrap();
1550 let tok = Tokenizer::from_hf_dir(&dir).expect("from_hf_dir");
1551 assert_eq!(tok.pre(), "qwen2");
1552 assert_eq!(
1553 tok.split(),
1554 PreSplit::Qwen35,
1555 "qwen2 rides the qwen35 split"
1556 );
1557 let _ = std::fs::remove_dir_all(&dir);
1558 }
1559
1560 #[test]
1563 fn hf_dir_refuses_unidentifiable_pretokenizer() {
1564 let json = TOKENIZER_JSON.replace(r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|", "SOMETHING-ELSE|");
1565 assert_ne!(json, TOKENIZER_JSON);
1566 let dir = std::env::temp_dir().join(format!("memra-tok-hf-unk-{}", std::process::id()));
1567 let _ = std::fs::remove_dir_all(&dir);
1568 std::fs::create_dir_all(&dir).unwrap();
1569 std::fs::write(dir.join("tokenizer.json"), &json).unwrap();
1570 std::fs::write(
1571 dir.join("generation_config.json"),
1572 r#"{"eos_token_id": 15}"#,
1573 )
1574 .unwrap();
1575 let err = match Tokenizer::from_hf_dir(&dir) {
1576 Ok(_) => panic!("unidentifiable pre must refuse to load"),
1577 Err(e) => e,
1578 };
1579 assert!(
1580 err.contains("unsupported tokenizer.ggml.pre 'default'"),
1581 "{err}"
1582 );
1583 assert!(
1584 err.contains("SOMETHING-ELSE"),
1585 "error must quote the regex: {err}"
1586 );
1587 assert!(err.contains("MEMRA_ALLOW_UNKNOWN_PRETOKENIZER"), "{err}");
1588 let _ = std::fs::remove_dir_all(&dir);
1589 }
1590
1591 #[test]
1597 fn staged_checkpoints_resolve_their_own_pretokenizer() {
1598 let cases: &[(&str, &str)] = &[
1599 (
1601 "/data/ai-ml/hf-models/hy3-layer103p5-sparse-source",
1602 "deepseek-v3",
1603 ),
1604 ("/data/ai-ml/hf-models/qwen3-1.7b-blk128fp8-synth", "qwen2"),
1606 ("/data/ai-ml/hf-models/qwen35-9b-hf", "qwen35"),
1608 ];
1609 let mut ran = 0;
1610 for (path, want) in cases {
1611 let dir = std::path::Path::new(path);
1612 if !dir.join("tokenizer.json").exists() {
1613 eprintln!("skip: {path} not staged");
1614 continue;
1615 }
1616 let tok = Tokenizer::from_hf_dir(dir).unwrap_or_else(|e| panic!("{path}: {e}"));
1617 assert_eq!(tok.pre(), *want, "{path}");
1618 ran += 1;
1619 }
1620 eprintln!("staged_checkpoints_resolve_their_own_pretokenizer: {ran}/3 cases ran");
1621 }
1622
1623 #[test]
1624 fn hf_dir_rejects_non_byte_level() {
1625 let dir = std::env::temp_dir().join(format!("memra-tok-hf-nonbl-{}", std::process::id()));
1626 let _ = std::fs::remove_dir_all(&dir);
1627 std::fs::create_dir_all(&dir).unwrap();
1628 let bad = TOKENIZER_JSON.replace("\"ByteLevel\"", "\"Metaspace\"");
1629 std::fs::write(dir.join("tokenizer.json"), bad).unwrap();
1630 assert!(Tokenizer::from_hf_dir(&dir).is_err());
1631 let _ = std::fs::remove_dir_all(&dir);
1632 }
1633
1634 #[test]
1637 fn hf_dir_dsv4_encoding_detection() {
1638 let gc = r#"{"eos_token_id": [15]}"#;
1639 let full_dspark = r#""dspark_block_size": 5, "dspark_markov_rank": 256,
1640 "dspark_noise_token_id": 128799, "dspark_target_layer_ids": [40, 41, 42]"#;
1641
1642 let dir = write_fixture("dsv4-none", None, Some(gc), None);
1644 let tok = Tokenizer::from_hf_dir(&dir).unwrap();
1645 assert_eq!(tok.dsv4_encoding(), None);
1646 let _ = std::fs::remove_dir_all(&dir);
1647
1648 let dir = write_fixture("dsv4-preview", None, Some(gc), None);
1650 std::fs::write(
1651 dir.join("config.json"),
1652 r#"{"model_type": "deepseek_v4", "num_hidden_layers": 43}"#,
1653 )
1654 .unwrap();
1655 let tok = Tokenizer::from_hf_dir(&dir).unwrap();
1656 assert_eq!(tok.dsv4_encoding(), Some(chat::Dsv4Encoding::Preview));
1657 let _ = std::fs::remove_dir_all(&dir);
1658
1659 let dir = write_fixture("dsv4-0731", None, Some(gc), None);
1661 std::fs::write(
1662 dir.join("config.json"),
1663 format!(r#"{{"model_type": "deepseek_v4", {full_dspark}}}"#),
1664 )
1665 .unwrap();
1666 let tok = Tokenizer::from_hf_dir(&dir).unwrap();
1667 assert_eq!(tok.dsv4_encoding(), Some(chat::Dsv4Encoding::V0731));
1668 let _ = std::fs::remove_dir_all(&dir);
1669
1670 let dir = write_fixture("dsv4-partial", None, Some(gc), None);
1672 std::fs::write(
1673 dir.join("config.json"),
1674 r#"{"model_type": "deepseek_v4", "dspark_block_size": 5}"#,
1675 )
1676 .unwrap();
1677 let err = match Tokenizer::from_hf_dir(&dir) {
1678 Err(e) => e,
1679 Ok(_) => panic!("a partial dspark_* config must refuse the load"),
1680 };
1681 assert!(err.contains("partial dspark_*"), "{err}");
1682 let _ = std::fs::remove_dir_all(&dir);
1683
1684 let dir = write_fixture("dsv4-foreign", None, Some(gc), None);
1686 std::fs::write(dir.join("config.json"), r#"{"model_type": "qwen3"}"#).unwrap();
1687 let tok = Tokenizer::from_hf_dir(&dir).unwrap();
1688 assert_eq!(tok.dsv4_encoding(), None);
1689 let _ = std::fs::remove_dir_all(&dir);
1690 }
1691}