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, HashSet};
22
23const TT_UNKNOWN: i64 = 2;
25const TT_CONTROL: i64 = 3;
26const TT_USER_DEFINED: i64 = 4;
27const TT_BYTE: i64 = 6;
28const MAX_TOKENIZER_ID: u32 = 1_000_000;
31const MAX_TOKENIZER_SPARSE_FACTOR: usize = 16;
32const MAX_TOKENIZER_SPARSE_SLACK: usize = 4096;
33const 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+";
34const 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+";
38const DEEPSEEK_V3_SPLIT_REGEXES: [&str; 3] = [
43 r"\p{N}{1,3}",
44 "[\u{4e00}-\u{9fa5}\u{3040}-\u{309f}\u{30a0}-\u{30ff}]+",
45 "[!\"#$%&'()*+,\\-./:;<=>?@\\[\\\\\\]^_`{|}~][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+",
46];
47
48pub const SUPPORTED_PRETOKENIZERS: &[&str] = &["qwen35", "qwen2", "deepseek-v3", "gemma4"];
51
52pub const ALLOW_UNKNOWN_PRETOKENIZER_ENV: &str = "MEMRA_ALLOW_UNKNOWN_PRETOKENIZER";
55
56fn allow_unknown_pretokenizer() -> bool {
57 std::env::var(ALLOW_UNKNOWN_PRETOKENIZER_ENV).as_deref() == Ok("1")
58}
59
60#[derive(Debug, Clone, PartialEq, Eq)]
68pub struct UnknownPretokenizer {
69 pub pre: String,
72 pub spm_style: bool,
75}
76
77impl std::fmt::Display for UnknownPretokenizer {
78 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79 write!(
80 f,
81 "unsupported tokenizer.ggml.pre '{}' (vocab model is {}) — memra has no exact \
82 pre-tokenizer split for it and token ids would NOT be exact. Supported: {}. \
83 Set {}=1 to load anyway for deliberate experimentation (token ids will be wrong).",
84 self.pre,
85 if self.spm_style { "SPM/gemma4" } else { "gpt2" },
86 SUPPORTED_PRETOKENIZERS.join(", "),
87 ALLOW_UNKNOWN_PRETOKENIZER_ENV,
88 )
89 }
90}
91
92impl std::error::Error for UnknownPretokenizer {}
93
94#[derive(Debug, Clone, Copy, PartialEq, Eq)]
98pub enum PreSplit {
99 Qwen35,
101 DeepseekV3,
103 Spm,
106 UnknownFallbackQwen35,
109}
110
111impl PreSplit {
112 pub fn resolve(pre: &str, spm_style: bool) -> Result<Self, UnknownPretokenizer> {
115 Self::resolve_with(pre, spm_style, allow_unknown_pretokenizer())
116 }
117
118 fn resolve_with(
121 pre: &str,
122 spm_style: bool,
123 allow_unknown: bool,
124 ) -> Result<Self, UnknownPretokenizer> {
125 match (pre, spm_style) {
129 ("qwen35" | "qwen2", false) => Ok(PreSplit::Qwen35),
130 ("deepseek-v3", false) => Ok(PreSplit::DeepseekV3),
131 ("gemma4", true) => Ok(PreSplit::Spm),
132 _ => {
133 let err = UnknownPretokenizer {
134 pre: pre.to_string(),
135 spm_style,
136 };
137 if allow_unknown {
138 eprintln!(
141 "memra-tokenizer: WARNING {ALLOW_UNKNOWN_PRETOKENIZER_ENV}=1 — loading \
142 with {err} FALLING BACK to the qwen35 split. Token ids are NOT exact: \
143 goldens, parity fixtures, acceptance counts and quality numbers taken \
144 on this model are all invalid."
145 );
146 Ok(PreSplit::UnknownFallbackQwen35)
147 } else {
148 Err(err)
149 }
150 }
151 }
152 }
153}
154
155#[derive(Debug, Clone, Copy, PartialEq, Eq)]
156enum TokAttr {
157 Normal,
158 Unknown,
159 Control,
160 UserDefined,
161 Byte,
162 Other,
163}
164
165impl TokAttr {
166 fn from_toktype(t: i64) -> Self {
167 match t {
168 TT_UNKNOWN => TokAttr::Unknown,
169 TT_CONTROL => TokAttr::Control,
170 TT_USER_DEFINED => TokAttr::UserDefined,
171 TT_BYTE => TokAttr::Byte,
172 1 => TokAttr::Normal,
173 _ => TokAttr::Other,
174 }
175 }
176 fn is_special(self) -> bool {
179 matches!(
180 self,
181 TokAttr::Control | TokAttr::UserDefined | TokAttr::Unknown
182 )
183 }
184}
185
186pub struct Tokenizer {
187 id_to_token: Vec<String>,
189 token_to_id: HashMap<String, u32>,
191 attrs: Vec<TokAttr>,
193 bpe_ranks: HashMap<(String, String), i32>,
195 special_tokens: Vec<u32>,
197 eos_id: u32,
198 bos_id: Option<u32>,
199 add_bos: bool,
200 pre: String,
201 split: PreSplit,
204 chat_template: Option<String>,
205 spm_style: bool,
207 dsv4_encoding: Option<chat::Dsv4Encoding>,
214}
215
216#[derive(Clone, Eq, PartialEq)]
221struct Bigram {
222 left: i32,
223 right: i32,
224 rank: i32,
225 text: String,
226}
227
228impl Ord for Bigram {
229 fn cmp(&self, other: &Self) -> Ordering {
230 match other.rank.cmp(&self.rank) {
233 Ordering::Equal => other.left.cmp(&self.left),
234 o => o,
235 }
236 }
237}
238impl PartialOrd for Bigram {
239 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
240 Some(self.cmp(other))
241 }
242}
243
244struct Symbol {
246 text: String,
247 prev: i32,
248 next: i32,
249 n: usize, }
251
252impl Tokenizer {
253 pub fn from_gguf(g: &GgufFile) -> Result<Self, String> {
255 let model = g
256 .metadata
257 .get("tokenizer.ggml.model")
258 .and_then(|v| v.as_str())
259 .ok_or("missing tokenizer.ggml.model")?;
260 if model != "gpt2" && model != "gemma4" {
261 return Err(format!(
262 "unsupported tokenizer model '{model}' (only gpt2/gemma4)"
263 ));
264 }
265 let spm_style = model == "gemma4";
269 let pre = g
270 .metadata
271 .get("tokenizer.ggml.pre")
272 .and_then(|v| v.as_str())
273 .unwrap_or(if spm_style { "gemma4" } else { "default" })
274 .to_string();
275 let split = PreSplit::resolve(&pre, spm_style).map_err(|e| e.to_string())?;
278
279 let tokens = match g.metadata.get("tokenizer.ggml.tokens") {
281 Some(MetaValue::Array(a)) => a,
282 _ => return Err("missing tokenizer.ggml.tokens array".into()),
283 };
284 let n = tokens.len();
285 if n > MAX_TOKENIZER_ID as usize + 1 {
286 return Err(format!(
287 "tokenizer.ggml.tokens has {n} entries; maximum supported vocabulary is {}",
288 MAX_TOKENIZER_ID + 1
289 ));
290 }
291 let mut id_to_token = Vec::with_capacity(n);
292 let mut token_to_id = HashMap::with_capacity(n);
293 for (i, t) in tokens.iter().enumerate() {
294 let s = t.as_str().ok_or("non-string in tokens[]")?.to_string();
295 token_to_id.entry(s.clone()).or_insert(i as u32);
297 id_to_token.push(s);
298 }
299
300 let mut attrs = vec![TokAttr::Normal; n];
302 if let Some(MetaValue::Array(a)) = g.metadata.get("tokenizer.ggml.token_type") {
303 for (i, v) in a.iter().enumerate().take(n) {
304 if let Some(t) = v.as_u64() {
305 attrs[i] = TokAttr::from_toktype(t as i64);
306 } else if let MetaValue::I32(t) = v {
307 attrs[i] = TokAttr::from_toktype(*t as i64);
308 }
309 }
310 }
311
312 let mut bpe_ranks = HashMap::new();
314 if let Some(MetaValue::Array(a)) = g.metadata.get("tokenizer.ggml.merges") {
315 for (i, v) in a.iter().enumerate() {
316 let word = v.as_str().ok_or("non-string in merges[]")?;
317 let bytes = word.as_bytes();
321 if let Some(pos) = bytes.iter().skip(1).position(|&b| b == b' ').map(|p| p + 1) {
322 let first = word[..pos].to_string();
323 let second = word[pos + 1..].to_string();
324 bpe_ranks.insert((first, second), i as i32);
325 }
326 }
327 } else {
328 return Err("missing tokenizer.ggml.merges array".into());
329 }
330
331 let mut special_tokens: Vec<u32> = (0..n as u32)
333 .filter(|&id| attrs[id as usize].is_special())
334 .collect();
335 special_tokens.sort_by(|&a, &b| {
336 id_to_token[b as usize]
337 .len()
338 .cmp(&id_to_token[a as usize].len())
339 });
340
341 let eos_id = g
342 .metadata
343 .get("tokenizer.ggml.eos_token_id")
344 .and_then(|v| v.as_u64())
345 .map(|v| v as u32)
346 .ok_or("missing tokenizer.ggml.eos_token_id")?;
347 let bos_id = g
348 .metadata
349 .get("tokenizer.ggml.bos_token_id")
350 .and_then(|v| v.as_u64())
351 .map(|v| v as u32);
352 let add_bos = g
353 .metadata
354 .get("tokenizer.ggml.add_bos_token")
355 .and_then(|v| match v {
356 MetaValue::Bool(b) => Some(*b),
357 _ => v.as_u64().map(|x| x != 0),
358 })
359 .unwrap_or(false);
360 let add_bos = add_bos || spm_style;
361
362 let chat_template = g
363 .metadata
364 .get("tokenizer.chat_template")
365 .and_then(|v| v.as_str())
366 .map(|s| s.to_string());
367
368 Ok(Tokenizer {
369 id_to_token,
370 token_to_id,
371 attrs,
372 bpe_ranks,
373 special_tokens,
374 eos_id,
375 bos_id,
376 add_bos,
377 pre,
378 split,
379 chat_template,
380 spm_style,
381 dsv4_encoding: None,
385 })
386 }
387
388 pub fn from_hf_dir(dir: &std::path::Path) -> Result<Self, String> {
406 let tj_path = dir.join("tokenizer.json");
407 let text = std::fs::read_to_string(&tj_path)
408 .map_err(|e| format!("read {}: {e}", tj_path.display()))?;
409 let tj = json::parse(&text).map_err(|e| format!("{}: {e}", tj_path.display()))?;
410
411 let model = tj.get("model").ok_or("tokenizer.json: missing model")?;
412 if let Some(t) = model.get("type").and_then(|v| v.as_str()) {
413 if t != "BPE" {
414 return Err(format!(
415 "unsupported tokenizer.json model type '{t}' (only BPE)"
416 ));
417 }
418 }
419 let pre_tok = tj
421 .get("pre_tokenizer")
422 .ok_or("tokenizer.json: missing pre_tokenizer")?;
423 if !pre_tokenizer_is_byte_level(pre_tok) {
424 return Err(
425 "tokenizer.json: pre_tokenizer is not ByteLevel — only byte-level \
426 BPE is supported"
427 .into(),
428 );
429 }
430
431 let vocab = model
433 .get("vocab")
434 .and_then(|v| v.as_obj())
435 .ok_or("tokenizer.json: missing model.vocab")?;
436 let empty: Vec<json::Value> = Vec::new();
437 let added = tj
438 .get("added_tokens")
439 .and_then(|v| v.as_arr())
440 .unwrap_or(&empty);
441 let mut max_id = 0u32;
442 let mut used_ids = HashSet::new();
443 for v in vocab.values() {
444 let id = v
445 .as_u64()
446 .ok_or("tokenizer.json: non-integer id in model.vocab")?;
447 let id = u32::try_from(id).map_err(|_| "tokenizer.json: vocabulary id exceeds u32")?;
448 if id > MAX_TOKENIZER_ID {
449 return Err(format!(
450 "tokenizer.json: vocabulary id {id} exceeds maximum {MAX_TOKENIZER_ID}"
451 ));
452 }
453 max_id = max_id.max(id);
454 used_ids.insert(id);
455 }
456 for a in added {
457 let id = a
458 .get("id")
459 .and_then(|v| v.as_u64())
460 .ok_or("tokenizer.json: added_tokens entry missing id")?;
461 let id = u32::try_from(id).map_err(|_| "tokenizer.json: added token id exceeds u32")?;
462 if id > MAX_TOKENIZER_ID {
463 return Err(format!(
464 "tokenizer.json: added token id {id} exceeds maximum {MAX_TOKENIZER_ID}"
465 ));
466 }
467 max_id = max_id.max(id);
468 used_ids.insert(id);
469 }
470 let n = (max_id as usize)
471 .checked_add(1)
472 .ok_or("tokenizer.json: vocabulary size overflow")?;
473 let entry_count = used_ids.len();
474 let max_dense_len = entry_count
475 .saturating_mul(MAX_TOKENIZER_SPARSE_FACTOR)
476 .saturating_add(MAX_TOKENIZER_SPARSE_SLACK);
477 if n > max_dense_len {
478 return Err(format!(
479 "tokenizer.json: vocabulary ids are too sparse (dense length {n}, {} entries; maximum {max_dense_len})",
480 entry_count
481 ));
482 }
483 let mut id_to_token = vec![String::new(); n];
484 let mut token_to_id: HashMap<String, u32> = HashMap::with_capacity(n);
485 let mut attrs = vec![TokAttr::Normal; n];
486 for (tok, v) in vocab {
487 let id = u32::try_from(
488 v.as_u64()
489 .ok_or("tokenizer.json: non-integer id in model.vocab")?,
490 )
491 .map_err(|_| "tokenizer.json: vocabulary id exceeds u32")?;
492 id_to_token[id as usize] = tok.clone();
493 token_to_id.entry(tok.clone()).or_insert(id);
494 }
495 for a in added {
498 let id = u32::try_from(
499 a.get("id")
500 .and_then(|v| v.as_u64())
501 .ok_or("tokenizer.json: added_tokens entry missing id")?,
502 )
503 .map_err(|_| "tokenizer.json: added token id exceeds u32")?;
504 let content = a
505 .get("content")
506 .and_then(|v| v.as_str())
507 .ok_or("tokenizer.json: added_tokens entry missing content")?;
508 if id_to_token[id as usize].is_empty() {
509 id_to_token[id as usize] = content.to_string();
510 }
511 token_to_id.entry(content.to_string()).or_insert(id);
512 if a.get("special").and_then(|v| v.as_bool()).unwrap_or(false) {
513 attrs[id as usize] = TokAttr::Control;
514 } else {
515 attrs[id as usize] = TokAttr::UserDefined;
521 }
522 }
523
524 let merges = model
526 .get("merges")
527 .and_then(|v| v.as_arr())
528 .ok_or("tokenizer.json: missing model.merges")?;
529 let mut bpe_ranks = HashMap::with_capacity(merges.len());
530 for (i, m) in merges.iter().enumerate() {
531 let (first, second) = match m {
532 json::Value::Str(s) => {
533 let bytes = s.as_bytes();
536 let pos = bytes
537 .iter()
538 .skip(1)
539 .position(|&b| b == b' ')
540 .map(|p| p + 1)
541 .ok_or_else(|| format!("tokenizer.json: merges[{i}] has no space"))?;
542 (s[..pos].to_string(), s[pos + 1..].to_string())
543 }
544 json::Value::Arr(a) if a.len() == 2 => {
545 let f = a[0]
546 .as_str()
547 .ok_or_else(|| format!("tokenizer.json: merges[{i}] non-string pair"))?;
548 let s2 = a[1]
549 .as_str()
550 .ok_or_else(|| format!("tokenizer.json: merges[{i}] non-string pair"))?;
551 (f.to_string(), s2.to_string())
552 }
553 _ => {
554 return Err(format!(
555 "tokenizer.json: merges[{i}] is neither \"a b\" string nor [a, b] pair"
556 ));
557 }
558 };
559 bpe_ranks.insert((first, second), i as i32);
560 }
561
562 let mut special_tokens: Vec<u32> = (0..n as u32)
564 .filter(|&id| attrs[id as usize].is_special())
565 .collect();
566 special_tokens.sort_by(|&a, &b| {
567 id_to_token[b as usize]
568 .len()
569 .cmp(&id_to_token[a as usize].len())
570 });
571
572 let tc = std::fs::read_to_string(dir.join("tokenizer_config.json"))
574 .ok()
575 .and_then(|t| json::parse(&t).ok());
576 let gc = std::fs::read_to_string(dir.join("generation_config.json"))
577 .ok()
578 .and_then(|t| json::parse(&t).ok());
579
580 let tok_content = |v: &json::Value| -> Option<String> {
582 v.as_str().map(|s| s.to_string()).or_else(|| {
583 v.get("content")
584 .and_then(|c| c.as_str())
585 .map(|s| s.to_string())
586 })
587 };
588 let eos_from_cfg = tc
589 .as_ref()
590 .and_then(|c| c.get("eos_token"))
591 .and_then(&tok_content)
592 .and_then(|s| token_to_id.get(&s).copied());
593 let eos_from_gen = gc
595 .as_ref()
596 .and_then(|c| c.get("eos_token_id"))
597 .and_then(|v| match v {
598 json::Value::Num(_) => v.as_u64(),
599 json::Value::Arr(a) => a.first().and_then(|x| x.as_u64()),
600 _ => None,
601 })
602 .map(|v| v as u32);
603 let eos_id = eos_from_cfg.or(eos_from_gen).ok_or(
604 "no eos token: need tokenizer_config.json eos_token or \
605 generation_config.json eos_token_id",
606 )?;
607 let bos_id = tc
608 .as_ref()
609 .and_then(|c| c.get("bos_token"))
610 .and_then(&tok_content)
611 .and_then(|s| token_to_id.get(&s).copied());
612 let add_bos = tc
613 .as_ref()
614 .and_then(|c| c.get("add_bos_token"))
615 .and_then(|v| v.as_bool())
616 .unwrap_or(false);
617
618 let chat_template = tc
620 .as_ref()
621 .and_then(|c| c.get("chat_template"))
622 .and_then(|v| v.as_str())
623 .map(|s| s.to_string())
624 .or_else(|| std::fs::read_to_string(dir.join("chat_template.jinja")).ok());
625 let cfg_regex = tc
633 .as_ref()
634 .and_then(|c| c.get("pretokenize_regex"))
635 .and_then(|v| v.as_str());
636 let mut tj_regexes: Vec<String> = Vec::new();
637 collect_split_regexes(pre_tok, &mut tj_regexes);
638 let pre = cfg_regex
639 .and_then(|r| pre_from_split_regexes(std::slice::from_ref(&r.to_string())))
640 .or_else(|| pre_from_split_regexes(&tj_regexes))
641 .unwrap_or("default");
642 let split = PreSplit::resolve(pre, false).map_err(|e| {
643 if pre == "default" {
644 format!(
645 "{e}\n (HF checkpoint {}: tokenizer_config.json pretokenize_regex = {:?}, \
646 tokenizer.json pre_tokenizer Split regexes = {:?} — neither matched a known \
647 family)",
648 dir.display(),
649 cfg_regex,
650 tj_regexes,
651 )
652 } else {
653 e.to_string()
654 }
655 })?;
656
657 let dsv4_encoding = dsv4_encoding_from_config(dir)?;
663
664 Ok(Tokenizer {
665 id_to_token,
666 token_to_id,
667 attrs,
668 bpe_ranks,
669 special_tokens,
670 eos_id,
671 bos_id,
672 add_bos,
673 pre: pre.to_string(),
674 split,
675 chat_template,
676 spm_style: false,
677 dsv4_encoding,
678 })
679 }
680
681 pub fn eos_id(&self) -> u32 {
682 self.eos_id
683 }
684 pub fn id_of(&self, piece: &str) -> Option<u32> {
686 self.token_to_id.get(piece).copied()
687 }
688 pub fn eog_ids(&self) -> Vec<u32> {
691 let mut ids = vec![self.eos_id];
692 for t in ["<|im_end|>", "<turn|>", "<end_of_turn>"] {
693 if let Some(&id) = self.token_to_id.get(t) {
694 if !ids.contains(&id) {
695 ids.push(id);
696 }
697 }
698 }
699 ids
700 }
701 pub fn bos_id(&self) -> Option<u32> {
702 self.bos_id
703 }
704 pub fn vocab_size(&self) -> usize {
705 self.id_to_token.len()
706 }
707 pub fn pre(&self) -> &str {
708 &self.pre
709 }
710 pub fn split(&self) -> PreSplit {
713 self.split
714 }
715 pub fn chat_template(&self) -> Option<&str> {
716 self.chat_template.as_deref()
717 }
718 pub fn dsv4_encoding(&self) -> Option<chat::Dsv4Encoding> {
721 self.dsv4_encoding
722 }
723
724 #[inline]
725 fn text_to_token(&self, s: &str) -> Option<u32> {
726 self.token_to_id.get(s).copied()
727 }
728
729 fn find_bpe_rank(&self, left: &str, right: &str) -> i32 {
730 self.bpe_ranks
731 .get(&(left.to_string(), right.to_string()))
732 .copied()
733 .unwrap_or(-1)
734 }
735
736 pub fn encode(&self, text: &str, add_special: bool) -> Vec<u32> {
742 self.encode_special(text, add_special, true)
743 }
744
745 pub fn encode_special(&self, text: &str, add_special: bool, parse_special: bool) -> Vec<u32> {
746 let mut output: Vec<u32> = Vec::new();
747 if add_special && self.add_bos {
748 if let Some(b) = self.bos_id {
749 output.push(b);
750 }
751 }
752 if text.is_empty() {
753 return output;
754 }
755
756 for frag in self.st_partition(text, parse_special) {
758 match frag {
759 Fragment::Token(id) => output.push(id),
760 Fragment::Text(span) => self.bpe_tokenize(&span, &mut output),
761 }
762 }
763 output
764 }
765
766 fn st_partition(&self, text: &str, parse_special: bool) -> Vec<Fragment> {
768 let mut frags = vec![Fragment::Text(text.to_string())];
769 for &sid in &self.special_tokens {
770 let attr = self.attrs[sid as usize];
771 if !parse_special && matches!(attr, TokAttr::Control | TokAttr::Unknown) {
773 continue;
774 }
775 let needle = &self.id_to_token[sid as usize];
776 if needle.is_empty() {
777 continue;
778 }
779 let mut next: Vec<Fragment> = Vec::with_capacity(frags.len());
780 for f in frags.drain(..) {
781 match f {
782 Fragment::Token(id) => next.push(Fragment::Token(id)),
783 Fragment::Text(s) => {
784 let mut rest: &str = &s;
785 let mut acc = String::new();
786 while let Some(m) = rest.find(needle.as_str()) {
787 acc.push_str(&rest[..m]);
788 if !acc.is_empty() {
789 next.push(Fragment::Text(std::mem::take(&mut acc)));
790 }
791 next.push(Fragment::Token(sid));
792 rest = &rest[m + needle.len()..];
793 }
794 acc.push_str(rest);
795 if !acc.is_empty() {
796 next.push(Fragment::Text(acc));
797 }
798 }
799 }
800 }
801 frags = next;
802 }
803 frags
804 }
805
806 fn bpe_tokenize(&self, text: &str, output: &mut Vec<u32>) {
808 if self.spm_style {
809 let escaped: String = text
812 .chars()
813 .map(|c| if c == ' ' { '\u{2581}' } else { c })
814 .collect();
815 let mut words: Vec<String> = Vec::new();
816 let mut cur = String::new();
817 let mut cur_nl: Option<bool> = None;
818 for c in escaped.chars() {
819 let nl = c == '\n';
820 if cur_nl != Some(nl) && !cur.is_empty() {
821 words.push(std::mem::take(&mut cur));
822 }
823 cur_nl = Some(nl);
824 cur.push(c);
825 }
826 if !cur.is_empty() {
827 words.push(cur);
828 }
829 for word in &words {
830 if word.chars().all(|c| c == '\n') {
832 if let Some(tok) = self.text_to_token(word) {
833 output.push(tok);
834 continue;
835 }
836 }
837 self.bpe_merge_word(word, output);
838 }
839 return;
840 }
841 let words: Vec<String> = match self.split {
847 PreSplit::Qwen35 => unicode::split_qwen35(text),
850 PreSplit::DeepseekV3 => unicode::split_deepseek_v3(text),
854 PreSplit::UnknownFallbackQwen35 => unicode::split_qwen35(text),
857 PreSplit::Spm => unreachable!("PreSplit::Spm implies spm_style, handled above"),
860 };
861
862 for word in &words {
863 let word = unicode::byte_encode(word);
864 self.bpe_merge_word(&word, output);
865 }
866 }
867
868 fn bpe_merge_word(&self, word: &str, output: &mut Vec<u32>) {
871 {
872 let word = word.to_string();
873
874 let chars: Vec<char> = word.chars().collect();
876 let mut symbols: Vec<Symbol> = Vec::with_capacity(chars.len());
877 for (i, &c) in chars.iter().enumerate() {
878 symbols.push(Symbol {
879 text: c.to_string(),
880 prev: i as i32 - 1,
881 next: if i + 1 == chars.len() {
882 -1
883 } else {
884 i as i32 + 1
885 },
886 n: 1,
887 });
888 }
889
890 let mut queue: BinaryHeap<Bigram> = BinaryHeap::new();
892 for i in 1..symbols.len() {
893 self.add_bigram(&symbols, i as i32 - 1, i as i32, &mut queue);
894 }
895
896 while let Some(bigram) = queue.pop() {
898 let li = bigram.left as usize;
899 let ri = bigram.right as usize;
900 if symbols[li].n == 0 || symbols[ri].n == 0 {
901 continue;
902 }
903 let combined = format!("{}{}", symbols[li].text, symbols[ri].text);
904 if combined != bigram.text {
905 continue; }
907 symbols[li].text = combined;
909 symbols[li].n += symbols[ri].n;
910 symbols[ri].n = 0;
911 let r_next = symbols[ri].next;
912 symbols[li].next = r_next;
913 if r_next >= 0 {
914 symbols[r_next as usize].prev = bigram.left;
915 }
916 let l_prev = symbols[li].prev;
917 let l_next = symbols[li].next;
918 self.add_bigram(&symbols, l_prev, bigram.left, &mut queue);
919 self.add_bigram(&symbols, bigram.left, l_next, &mut queue);
920 }
921
922 for sym in &symbols {
924 if sym.n == 0 {
925 continue;
926 }
927 match self.text_to_token(&sym.text) {
928 Some(tok) => output.push(tok),
929 None => {
930 for b in sym.text.bytes() {
932 let bs = if self.spm_style {
933 format!("<0x{b:02X}>") } else {
935 (b as char).to_string()
936 };
937 if let Some(t) = self.text_to_token(&bs) {
938 output.push(t);
939 }
940 }
941 }
942 }
943 }
944 }
945 }
946
947 fn add_bigram(
948 &self,
949 symbols: &[Symbol],
950 left: i32,
951 right: i32,
952 queue: &mut BinaryHeap<Bigram>,
953 ) {
954 if left == -1 || right == -1 {
955 return;
956 }
957 let lt = &symbols[left as usize].text;
958 let rt = &symbols[right as usize].text;
959 let rank = self.find_bpe_rank(lt, rt);
960 if rank < 0 {
961 return;
962 }
963 queue.push(Bigram {
964 left,
965 right,
966 rank,
967 text: format!("{lt}{rt}"),
968 });
969 }
970
971 pub fn decode(&self, ids: &[u32]) -> String {
974 self.decode_special(ids, true)
975 }
976
977 pub fn token_is_control(&self, id: u32) -> bool {
982 match self.attrs.get(id as usize) {
983 Some(TokAttr::Control) | Some(TokAttr::Unknown) => true,
984 _ => false,
985 }
986 }
987
988 pub fn decode_special(&self, ids: &[u32], special: bool) -> String {
989 String::from_utf8_lossy(&self.decode_bytes_special(ids, special)).into_owned()
990 }
991
992 pub fn decode_bytes_special(&self, ids: &[u32], special: bool) -> Vec<u8> {
995 let mut bytes: Vec<u8> = Vec::new();
996 for &id in ids {
997 let i = id as usize;
998 if i >= self.id_to_token.len() {
999 continue;
1000 }
1001 let attr = self.attrs[i];
1002 let piece = &self.id_to_token[i];
1003 match attr {
1004 TokAttr::Normal | TokAttr::Byte => {
1005 if self.spm_style {
1006 if matches!(attr, TokAttr::Byte)
1008 || (piece.len() == 6
1009 && piece.starts_with("<0x")
1010 && piece.ends_with('>'))
1011 {
1012 if let Ok(b) = u8::from_str_radix(&piece[3..5], 16) {
1013 bytes.push(b);
1014 continue;
1015 }
1016 }
1017 for c in piece.chars() {
1018 if c == '\u{2581}' {
1019 bytes.push(b' ');
1020 } else {
1021 let mut buf = [0u8; 4];
1022 bytes.extend_from_slice(c.encode_utf8(&mut buf).as_bytes());
1023 }
1024 }
1025 } else {
1026 self.piece_to_bytes(piece, &mut bytes);
1028 }
1029 }
1030 TokAttr::UserDefined => {
1031 bytes.extend_from_slice(piece.as_bytes());
1033 }
1034 TokAttr::Control | TokAttr::Unknown => {
1035 if special {
1036 bytes.extend_from_slice(piece.as_bytes());
1037 }
1038 }
1040 TokAttr::Other => {}
1041 }
1042 }
1043 bytes
1044 }
1045
1046 fn piece_to_bytes(&self, piece: &str, out: &mut Vec<u8>) {
1047 for c in piece.chars() {
1048 match unicode::unicode_to_byte(c) {
1049 Some(b) => out.push(b),
1050 None => {
1051 let mut buf = [0u8; 4];
1053 out.extend_from_slice(c.encode_utf8(&mut buf).as_bytes());
1054 }
1055 }
1056 }
1057 }
1058
1059 pub fn apply_chat_template(
1062 &self,
1063 messages: &[(&str, &str)],
1064 add_generation_prompt: bool,
1065 ) -> String {
1066 chat::apply_chat_template_enc(
1067 self.chat_template.as_deref(),
1068 messages,
1069 add_generation_prompt,
1070 self.dsv4_encoding,
1071 )
1072 .expect("plain chat render cannot fail")
1075 }
1076
1077 pub fn has_qwen_effort_ladder(&self) -> bool {
1085 self.chat_template
1086 .as_deref()
1087 .is_some_and(chat::template_has_qwen_effort)
1088 }
1089
1090 pub fn apply_chat_template_tools(
1096 &self,
1097 turns: &[chat::Turn],
1098 add_generation_prompt: bool,
1099 tools_json: &[String],
1100 think: chat::ThinkMode,
1101 reasoning_effort: Option<&str>,
1102 ) -> Result<String, String> {
1103 chat::apply_chat_template_tools_ex(
1104 self.chat_template.as_deref(),
1105 turns,
1106 add_generation_prompt,
1107 tools_json,
1108 &[],
1109 think,
1110 reasoning_effort,
1111 self.dsv4_encoding,
1112 )
1113 }
1114
1115 #[allow(clippy::too_many_arguments)]
1120 pub fn apply_chat_template_tools_ex(
1121 &self,
1122 turns: &[chat::Turn],
1123 add_generation_prompt: bool,
1124 tools_json: &[String],
1125 tools_struct: &[chat::Val],
1126 think: chat::ThinkMode,
1127 reasoning_effort: Option<&str>,
1128 ) -> Result<String, String> {
1129 chat::apply_chat_template_tools_ex(
1130 self.chat_template.as_deref(),
1131 turns,
1132 add_generation_prompt,
1133 tools_json,
1134 tools_struct,
1135 think,
1136 reasoning_effort,
1137 self.dsv4_encoding,
1138 )
1139 }
1140}
1141
1142enum Fragment {
1143 Text(String),
1144 Token(u32),
1145}
1146
1147fn dsv4_encoding_from_config(dir: &std::path::Path) -> Result<Option<chat::Dsv4Encoding>, String> {
1162 const DSPARK_KEYS: [&str; 4] = [
1163 "dspark_block_size",
1164 "dspark_markov_rank",
1165 "dspark_noise_token_id",
1166 "dspark_target_layer_ids",
1167 ];
1168 let cfg_path = dir.join("config.json");
1169 let Ok(text) = std::fs::read_to_string(&cfg_path) else {
1170 return Ok(None);
1171 };
1172 let Ok(cfg) = json::parse(&text) else {
1173 return Ok(None);
1176 };
1177 if cfg.get("model_type").and_then(|v| v.as_str()) != Some("deepseek_v4") {
1178 return Ok(None);
1181 }
1182 let present: Vec<&str> = DSPARK_KEYS
1183 .iter()
1184 .copied()
1185 .filter(|k| cfg.get(k).is_some())
1186 .collect();
1187 match present.len() {
1188 0 => Ok(Some(chat::Dsv4Encoding::Preview)),
1189 4 => Ok(Some(chat::Dsv4Encoding::V0731)),
1190 _ => Err(format!(
1191 "{}: partial dspark_* key set {:?} (expected none or all of {:?}) — cannot \
1192 determine the deepseek-v4 encoding revision; refusing rather than guessing \
1193 the reasoning-effort ladder",
1194 cfg_path.display(),
1195 present,
1196 DSPARK_KEYS
1197 )),
1198 }
1199}
1200
1201fn collect_split_regexes(pt: &json::Value, out: &mut Vec<String>) {
1207 match pt.get("type").and_then(|v| v.as_str()) {
1208 Some("Sequence") => {
1209 if let Some(arr) = pt.get("pretokenizers").and_then(|v| v.as_arr()) {
1210 for step in arr {
1211 collect_split_regexes(step, out);
1212 }
1213 }
1214 }
1215 Some("Split") => {
1216 if let Some(r) = pt
1217 .get("pattern")
1218 .and_then(|p| p.get("Regex"))
1219 .and_then(|v| v.as_str())
1220 {
1221 out.push(r.to_string());
1222 }
1223 }
1224 _ => {}
1225 }
1226}
1227
1228fn pre_from_split_regexes(regexes: &[String]) -> Option<&'static str> {
1233 match regexes {
1234 [one] if one == QWEN35_PRETOKENIZE_REGEX => Some("qwen35"),
1235 [one] if one == QWEN2_PRETOKENIZE_REGEX => Some("qwen2"),
1236 [a, b, c]
1237 if a == DEEPSEEK_V3_SPLIT_REGEXES[0]
1238 && b == DEEPSEEK_V3_SPLIT_REGEXES[1]
1239 && c == DEEPSEEK_V3_SPLIT_REGEXES[2] =>
1240 {
1241 Some("deepseek-v3")
1242 }
1243 _ => None,
1244 }
1245}
1246
1247fn pre_tokenizer_is_byte_level(pt: &json::Value) -> bool {
1248 match pt.get("type").and_then(|v| v.as_str()) {
1249 Some("ByteLevel") => true,
1250 Some("Sequence") => pt
1251 .get("pretokenizers")
1252 .and_then(|v| v.as_arr())
1253 .map(|arr| arr.iter().any(pre_tokenizer_is_byte_level))
1254 .unwrap_or(false),
1255 _ => false,
1256 }
1257}
1258
1259#[cfg(test)]
1260mod pretokenizer_tests {
1261 use super::*;
1262
1263 #[test]
1267 fn every_supported_pre_resolves() {
1268 assert_eq!(
1269 PreSplit::resolve_with("qwen35", false, false),
1270 Ok(PreSplit::Qwen35)
1271 );
1272 assert_eq!(
1273 PreSplit::resolve_with("qwen2", false, false),
1274 Ok(PreSplit::Qwen35)
1275 );
1276 assert_eq!(
1277 PreSplit::resolve_with("deepseek-v3", false, false),
1278 Ok(PreSplit::DeepseekV3)
1279 );
1280 assert_eq!(
1281 PreSplit::resolve_with("gemma4", true, false),
1282 Ok(PreSplit::Spm)
1283 );
1284 assert_eq!(
1287 SUPPORTED_PRETOKENIZERS,
1288 &["qwen35", "qwen2", "deepseek-v3", "gemma4"]
1289 );
1290 }
1291
1292 #[test]
1294 fn unknown_pre_is_a_typed_error() {
1295 let err =
1296 PreSplit::resolve_with("llama4", false, false).expect_err("llama4 has no ported split");
1297 assert_eq!(
1298 err,
1299 UnknownPretokenizer {
1300 pre: "llama4".into(),
1301 spm_style: false
1302 }
1303 );
1304 let msg = err.to_string();
1305 assert!(msg.contains("'llama4'"), "{msg}");
1307 for supported in SUPPORTED_PRETOKENIZERS {
1308 assert!(
1309 msg.contains(supported),
1310 "error must list {supported}: {msg}"
1311 );
1312 }
1313 assert!(msg.contains(ALLOW_UNKNOWN_PRETOKENIZER_ENV), "{msg}");
1314 let _: &dyn std::error::Error = &err;
1316 }
1317
1318 #[test]
1321 fn pre_and_vocab_model_must_agree() {
1322 assert!(PreSplit::resolve_with("qwen35", true, false).is_err());
1323 assert!(PreSplit::resolve_with("gemma4", false, false).is_err());
1324 assert!(PreSplit::resolve_with("default", false, false).is_err());
1326 assert!(PreSplit::resolve_with("", false, false).is_err());
1327 }
1328
1329 #[test]
1331 fn opt_out_loads_with_a_fallback_marker() {
1332 assert_eq!(
1333 PreSplit::resolve_with("llama4", false, true),
1334 Ok(PreSplit::UnknownFallbackQwen35)
1335 );
1336 assert_eq!(
1338 PreSplit::resolve_with("qwen35", true, true),
1339 Ok(PreSplit::UnknownFallbackQwen35)
1340 );
1341 }
1342
1343 #[test]
1346 fn opt_out_env_gate() {
1347 unsafe { std::env::remove_var(ALLOW_UNKNOWN_PRETOKENIZER_ENV) };
1350 assert!(!allow_unknown_pretokenizer());
1351 unsafe { std::env::set_var(ALLOW_UNKNOWN_PRETOKENIZER_ENV, "0") };
1352 assert!(!allow_unknown_pretokenizer());
1353 unsafe { std::env::set_var(ALLOW_UNKNOWN_PRETOKENIZER_ENV, "1") };
1354 assert!(allow_unknown_pretokenizer());
1355 assert_eq!(
1356 PreSplit::resolve("llama4", false),
1357 Ok(PreSplit::UnknownFallbackQwen35)
1358 );
1359 unsafe { std::env::remove_var(ALLOW_UNKNOWN_PRETOKENIZER_ENV) };
1360 assert!(PreSplit::resolve("llama4", false).is_err());
1361 }
1362
1363 #[test]
1366 fn split_regex_identification_is_exact() {
1367 let s = |v: &[&str]| v.iter().map(|x| x.to_string()).collect::<Vec<_>>();
1368 assert_eq!(
1369 pre_from_split_regexes(&s(&[QWEN35_PRETOKENIZE_REGEX])),
1370 Some("qwen35")
1371 );
1372 assert_eq!(
1373 pre_from_split_regexes(&s(&[QWEN2_PRETOKENIZE_REGEX])),
1374 Some("qwen2")
1375 );
1376 assert_eq!(
1377 pre_from_split_regexes(&s(&DEEPSEEK_V3_SPLIT_REGEXES)),
1378 Some("deepseek-v3")
1379 );
1380 assert_eq!(
1382 pre_from_split_regexes(&s(&[
1383 DEEPSEEK_V3_SPLIT_REGEXES[1],
1384 DEEPSEEK_V3_SPLIT_REGEXES[0],
1385 DEEPSEEK_V3_SPLIT_REGEXES[2],
1386 ])),
1387 None
1388 );
1389 assert_eq!(
1391 pre_from_split_regexes(&s(&[
1392 DEEPSEEK_V3_SPLIT_REGEXES[0],
1393 DEEPSEEK_V3_SPLIT_REGEXES[1]
1394 ])),
1395 None
1396 );
1397 let mut near = QWEN35_PRETOKENIZE_REGEX.to_string();
1399 near.push('x');
1400 assert_eq!(pre_from_split_regexes(&s(&[&near])), None);
1401 assert_eq!(pre_from_split_regexes(&[]), None);
1402 assert_ne!(QWEN2_PRETOKENIZE_REGEX, QWEN35_PRETOKENIZE_REGEX);
1404 }
1405
1406 #[test]
1409 fn collect_split_regexes_walks_in_order() {
1410 let src = r#"{"type":"Sequence","pretokenizers":[
1411 {"type":"Split","pattern":{"Regex":"A"},"behavior":"Isolated"},
1412 {"type":"Split","pattern":{"String":" "},"behavior":"Isolated"},
1413 {"type":"Digits","individual_digits":true},
1414 {"type":"Sequence","pretokenizers":[
1415 {"type":"Split","pattern":{"Regex":"B"},"behavior":"Isolated"}
1416 ]},
1417 {"type":"ByteLevel","add_prefix_space":false}
1418 ]}"#;
1419 let v = json::parse(src).unwrap();
1420 let mut out = Vec::new();
1421 collect_split_regexes(&v, &mut out);
1422 assert_eq!(out, vec!["A".to_string(), "B".to_string()]);
1423 }
1424}
1425
1426#[cfg(test)]
1427mod hf_tests {
1428 use super::*;
1429
1430 const TOKENIZER_JSON: &str = r#"{
1439 "version": "1.0",
1440 "added_tokens": [
1441 {"id": 15, "content": "<|end|>", "special": true},
1442 {"id": 16, "content": "<think>", "special": false}
1443 ],
1444 "pre_tokenizer": {
1445 "type": "Sequence",
1446 "pretokenizers": [
1447 {"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"},
1448 {"type": "ByteLevel", "add_prefix_space": false, "trim_offsets": false}
1449 ]
1450 },
1451 "model": {
1452 "type": "BPE",
1453 "vocab": {
1454 "h": 0, "e": 1, "l": 2, "o": 3, "Ġ": 4, "w": 5, "r": 6, "d": 7,
1455 "he": 8, "ll": 9, "hell": 10, "hello": 11, "Ġw": 12, "or": 13, "!": 14
1456 },
1457 "merges": [
1458 "h e",
1459 ["l", "l"],
1460 "he ll",
1461 ["hell", "o"],
1462 ["Ġ", "w"],
1463 "o r"
1464 ]
1465 }
1466 }"#;
1467
1468 fn write_fixture(
1469 name: &str,
1470 tokenizer_config: Option<&str>,
1471 generation_config: Option<&str>,
1472 jinja: Option<&str>,
1473 ) -> std::path::PathBuf {
1474 let dir = std::env::temp_dir().join(format!("memra-tok-hf-{name}-{}", std::process::id()));
1475 let _ = std::fs::remove_dir_all(&dir);
1476 std::fs::create_dir_all(&dir).unwrap();
1477 std::fs::write(dir.join("tokenizer.json"), TOKENIZER_JSON).unwrap();
1478 if let Some(tc) = tokenizer_config {
1479 std::fs::write(dir.join("tokenizer_config.json"), tc).unwrap();
1480 }
1481 if let Some(gc) = generation_config {
1482 std::fs::write(dir.join("generation_config.json"), gc).unwrap();
1483 }
1484 if let Some(j) = jinja {
1485 std::fs::write(dir.join("chat_template.jinja"), j).unwrap();
1486 }
1487 dir
1488 }
1489
1490 #[test]
1491 fn hf_dir_encode_decode_roundtrip_and_specials() {
1492 let tc = r#"{
1494 "eos_token": {"content": "<|end|>", "lstrip": false},
1495 "add_bos_token": false,
1496 "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+",
1497 "chat_template": "{{ messages }}<|end|>"
1498 }"#;
1499 let dir = write_fixture("full", Some(tc), None, None);
1500 let tok = Tokenizer::from_hf_dir(&dir).expect("from_hf_dir");
1501
1502 assert_eq!(tok.eos_id(), 15);
1503 assert_eq!(tok.bos_id(), None);
1504 assert_eq!(tok.pre(), "qwen35");
1505 assert_eq!(tok.vocab_size(), 17); assert_eq!(tok.chat_template(), Some("{{ messages }}<|end|>"));
1507
1508 let ids = tok.encode("hello world", true);
1512 assert_eq!(ids, vec![11, 12, 13, 2, 7]);
1513 assert_eq!(tok.decode(&ids), "hello world");
1514
1515 let ids = tok.encode("hello<|end|> world", true);
1517 assert_eq!(ids, vec![11, 15, 12, 13, 2, 7]);
1518 assert_eq!(tok.decode_special(&ids, true), "hello<|end|> world");
1520 assert_eq!(tok.decode_special(&ids, false), "hello world");
1521
1522 assert_eq!(tok.decode(&[16]), "<think>");
1524 let _ = std::fs::remove_dir_all(&dir);
1525 }
1526
1527 #[test]
1528 fn hf_dir_generation_config_eos_fallback_and_jinja() {
1529 let gc = r#"{"eos_token_id": [15, 14]}"#;
1532 let dir = write_fixture("genconf", None, Some(gc), Some("JINJA {{ messages }}"));
1533 let tok = Tokenizer::from_hf_dir(&dir).expect("from_hf_dir");
1534 assert_eq!(tok.eos_id(), 15);
1535 assert!(!tok.encode("hello", true).is_empty());
1536 assert_eq!(tok.chat_template(), Some("JINJA {{ messages }}"));
1537 let _ = std::fs::remove_dir_all(&dir);
1538 }
1539
1540 #[test]
1544 fn hf_dir_identifies_deepseek_v3_from_tokenizer_json() {
1545 let dsv3_pt = r##""pre_tokenizer": {
1546 "type": "Sequence",
1547 "pretokenizers": [
1548 {"type": "Split", "pattern": {"Regex": "\\p{N}{1,3}"}, "behavior": "Isolated"},
1549 {"type": "Split", "pattern": {"Regex": "[一-龥-ゟ゠-ヿ]+"}, "behavior": "Isolated"},
1550 {"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"},
1551 {"type": "ByteLevel", "add_prefix_space": false, "trim_offsets": true, "use_regex": false}
1552 ]
1553 },"##;
1554 let open = TOKENIZER_JSON.find(r#""pre_tokenizer""#).unwrap();
1556 let close = TOKENIZER_JSON.find(r#""model""#).unwrap();
1557 let json = format!(
1558 "{}{}\n {}",
1559 &TOKENIZER_JSON[..open],
1560 dsv3_pt,
1561 &TOKENIZER_JSON[close..]
1562 );
1563 let dir = std::env::temp_dir().join(format!("memra-tok-hf-dsv3-{}", std::process::id()));
1564 let _ = std::fs::remove_dir_all(&dir);
1565 std::fs::create_dir_all(&dir).unwrap();
1566 std::fs::write(dir.join("tokenizer.json"), &json).unwrap();
1567 std::fs::write(
1568 dir.join("generation_config.json"),
1569 r#"{"eos_token_id": 15}"#,
1570 )
1571 .unwrap();
1572 let tok = Tokenizer::from_hf_dir(&dir).expect("from_hf_dir");
1573 assert_eq!(tok.pre(), "deepseek-v3");
1574 assert_eq!(tok.split(), PreSplit::DeepseekV3);
1575 let _ = std::fs::remove_dir_all(&dir);
1576 }
1577
1578 #[test]
1581 fn hf_dir_identifies_qwen2_regex() {
1582 let json = TOKENIZER_JSON.replace(
1583 r"[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+",
1584 r"[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+",
1585 );
1586 assert_ne!(json, TOKENIZER_JSON, "the qwen2 substitution must apply");
1587 let dir = std::env::temp_dir().join(format!("memra-tok-hf-qwen2-{}", std::process::id()));
1588 let _ = std::fs::remove_dir_all(&dir);
1589 std::fs::create_dir_all(&dir).unwrap();
1590 std::fs::write(dir.join("tokenizer.json"), &json).unwrap();
1591 std::fs::write(
1592 dir.join("generation_config.json"),
1593 r#"{"eos_token_id": 15}"#,
1594 )
1595 .unwrap();
1596 let tok = Tokenizer::from_hf_dir(&dir).expect("from_hf_dir");
1597 assert_eq!(tok.pre(), "qwen2");
1598 assert_eq!(
1599 tok.split(),
1600 PreSplit::Qwen35,
1601 "qwen2 rides the qwen35 split"
1602 );
1603 let _ = std::fs::remove_dir_all(&dir);
1604 }
1605
1606 #[test]
1609 fn hf_dir_refuses_unidentifiable_pretokenizer() {
1610 let json = TOKENIZER_JSON.replace(r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|", "SOMETHING-ELSE|");
1611 assert_ne!(json, TOKENIZER_JSON);
1612 let dir = std::env::temp_dir().join(format!("memra-tok-hf-unk-{}", std::process::id()));
1613 let _ = std::fs::remove_dir_all(&dir);
1614 std::fs::create_dir_all(&dir).unwrap();
1615 std::fs::write(dir.join("tokenizer.json"), &json).unwrap();
1616 std::fs::write(
1617 dir.join("generation_config.json"),
1618 r#"{"eos_token_id": 15}"#,
1619 )
1620 .unwrap();
1621 let err = match Tokenizer::from_hf_dir(&dir) {
1622 Ok(_) => panic!("unidentifiable pre must refuse to load"),
1623 Err(e) => e,
1624 };
1625 assert!(
1626 err.contains("unsupported tokenizer.ggml.pre 'default'"),
1627 "{err}"
1628 );
1629 assert!(
1630 err.contains("SOMETHING-ELSE"),
1631 "error must quote the regex: {err}"
1632 );
1633 assert!(err.contains("MEMRA_ALLOW_UNKNOWN_PRETOKENIZER"), "{err}");
1634 let _ = std::fs::remove_dir_all(&dir);
1635 }
1636
1637 #[test]
1643 fn staged_checkpoints_resolve_their_own_pretokenizer() {
1644 let cases: &[(&str, &str)] = &[
1645 (
1647 "/data/ai-ml/hf-models/hy3-layer103p5-sparse-source",
1648 "deepseek-v3",
1649 ),
1650 ("/data/ai-ml/hf-models/qwen3-1.7b-blk128fp8-synth", "qwen2"),
1652 ("/data/ai-ml/hf-models/qwen35-9b-hf", "qwen35"),
1654 ];
1655 let mut ran = 0;
1656 for (path, want) in cases {
1657 let dir = std::path::Path::new(path);
1658 if !dir.join("tokenizer.json").exists() {
1659 eprintln!("skip: {path} not staged");
1660 continue;
1661 }
1662 let tok = Tokenizer::from_hf_dir(dir).unwrap_or_else(|e| panic!("{path}: {e}"));
1663 assert_eq!(tok.pre(), *want, "{path}");
1664 ran += 1;
1665 }
1666 eprintln!("staged_checkpoints_resolve_their_own_pretokenizer: {ran}/3 cases ran");
1667 }
1668
1669 #[test]
1670 fn hf_dir_rejects_non_byte_level() {
1671 let dir = std::env::temp_dir().join(format!("memra-tok-hf-nonbl-{}", std::process::id()));
1672 let _ = std::fs::remove_dir_all(&dir);
1673 std::fs::create_dir_all(&dir).unwrap();
1674 let bad = TOKENIZER_JSON.replace("\"ByteLevel\"", "\"Metaspace\"");
1675 std::fs::write(dir.join("tokenizer.json"), bad).unwrap();
1676 assert!(Tokenizer::from_hf_dir(&dir).is_err());
1677 let _ = std::fs::remove_dir_all(&dir);
1678 }
1679
1680 #[test]
1683 fn hf_dir_dsv4_encoding_detection() {
1684 let gc = r#"{"eos_token_id": [15]}"#;
1685 let full_dspark = r#""dspark_block_size": 5, "dspark_markov_rank": 256,
1686 "dspark_noise_token_id": 128799, "dspark_target_layer_ids": [40, 41, 42]"#;
1687
1688 let dir = write_fixture("dsv4-none", None, Some(gc), None);
1690 let tok = Tokenizer::from_hf_dir(&dir).unwrap();
1691 assert_eq!(tok.dsv4_encoding(), None);
1692 let _ = std::fs::remove_dir_all(&dir);
1693
1694 let dir = write_fixture("dsv4-preview", None, Some(gc), None);
1696 std::fs::write(
1697 dir.join("config.json"),
1698 r#"{"model_type": "deepseek_v4", "num_hidden_layers": 43}"#,
1699 )
1700 .unwrap();
1701 let tok = Tokenizer::from_hf_dir(&dir).unwrap();
1702 assert_eq!(tok.dsv4_encoding(), Some(chat::Dsv4Encoding::Preview));
1703 let _ = std::fs::remove_dir_all(&dir);
1704
1705 let dir = write_fixture("dsv4-0731", None, Some(gc), None);
1707 std::fs::write(
1708 dir.join("config.json"),
1709 format!(r#"{{"model_type": "deepseek_v4", {full_dspark}}}"#),
1710 )
1711 .unwrap();
1712 let tok = Tokenizer::from_hf_dir(&dir).unwrap();
1713 assert_eq!(tok.dsv4_encoding(), Some(chat::Dsv4Encoding::V0731));
1714 let _ = std::fs::remove_dir_all(&dir);
1715
1716 let dir = write_fixture("dsv4-partial", None, Some(gc), None);
1718 std::fs::write(
1719 dir.join("config.json"),
1720 r#"{"model_type": "deepseek_v4", "dspark_block_size": 5}"#,
1721 )
1722 .unwrap();
1723 let err = match Tokenizer::from_hf_dir(&dir) {
1724 Err(e) => e,
1725 Ok(_) => panic!("a partial dspark_* config must refuse the load"),
1726 };
1727 assert!(err.contains("partial dspark_*"), "{err}");
1728 let _ = std::fs::remove_dir_all(&dir);
1729
1730 let dir = write_fixture("dsv4-foreign", None, Some(gc), None);
1732 std::fs::write(dir.join("config.json"), r#"{"model_type": "qwen3"}"#).unwrap();
1733 let tok = Tokenizer::from_hf_dir(&dir).unwrap();
1734 assert_eq!(tok.dsv4_encoding(), None);
1735 let _ = std::fs::remove_dir_all(&dir);
1736 }
1737}