pub mod chat;
mod json;
mod unicode;
mod unicode_data;
pub use chat::apply_chat_template_str;
use memra_gguf::{GgufFile, MetaValue};
use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap};
const TT_UNKNOWN: i64 = 2;
const TT_CONTROL: i64 = 3;
const TT_USER_DEFINED: i64 = 4;
const TT_BYTE: i64 = 6;
const 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+";
const 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+";
const DEEPSEEK_V3_SPLIT_REGEXES: [&str; 3] = [
r"\p{N}{1,3}",
"[\u{4e00}-\u{9fa5}\u{3040}-\u{309f}\u{30a0}-\u{30ff}]+",
"[!\"#$%&'()*+,\\-./:;<=>?@\\[\\\\\\]^_`{|}~][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+",
];
pub const SUPPORTED_PRETOKENIZERS: &[&str] = &["qwen35", "qwen2", "deepseek-v3", "gemma4"];
pub const ALLOW_UNKNOWN_PRETOKENIZER_ENV: &str = "MEMRA_ALLOW_UNKNOWN_PRETOKENIZER";
fn allow_unknown_pretokenizer() -> bool {
std::env::var(ALLOW_UNKNOWN_PRETOKENIZER_ENV).as_deref() == Ok("1")
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UnknownPretokenizer {
pub pre: String,
pub spm_style: bool,
}
impl std::fmt::Display for UnknownPretokenizer {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"unsupported tokenizer.ggml.pre '{}' (vocab model is {}) — memra has no exact \
pre-tokenizer split for it and token ids would NOT be exact. Supported: {}. \
Set {}=1 to load anyway for deliberate experimentation (token ids will be wrong).",
self.pre,
if self.spm_style { "SPM/gemma4" } else { "gpt2" },
SUPPORTED_PRETOKENIZERS.join(", "),
ALLOW_UNKNOWN_PRETOKENIZER_ENV,
)
}
}
impl std::error::Error for UnknownPretokenizer {}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PreSplit {
Qwen35,
DeepseekV3,
Spm,
UnknownFallbackQwen35,
}
impl PreSplit {
pub fn resolve(pre: &str, spm_style: bool) -> Result<Self, UnknownPretokenizer> {
Self::resolve_with(pre, spm_style, allow_unknown_pretokenizer())
}
fn resolve_with(
pre: &str,
spm_style: bool,
allow_unknown: bool,
) -> Result<Self, UnknownPretokenizer> {
match (pre, spm_style) {
("qwen35" | "qwen2", false) => Ok(PreSplit::Qwen35),
("deepseek-v3", false) => Ok(PreSplit::DeepseekV3),
("gemma4", true) => Ok(PreSplit::Spm),
_ => {
let err = UnknownPretokenizer {
pre: pre.to_string(),
spm_style,
};
if allow_unknown {
eprintln!(
"memra-tokenizer: WARNING {ALLOW_UNKNOWN_PRETOKENIZER_ENV}=1 — loading \
with {err} FALLING BACK to the qwen35 split. Token ids are NOT exact: \
goldens, parity fixtures, acceptance counts and quality numbers taken \
on this model are all invalid."
);
Ok(PreSplit::UnknownFallbackQwen35)
} else {
Err(err)
}
}
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum TokAttr {
Normal,
Unknown,
Control,
UserDefined,
Byte,
Other,
}
impl TokAttr {
fn from_toktype(t: i64) -> Self {
match t {
TT_UNKNOWN => TokAttr::Unknown,
TT_CONTROL => TokAttr::Control,
TT_USER_DEFINED => TokAttr::UserDefined,
TT_BYTE => TokAttr::Byte,
1 => TokAttr::Normal,
_ => TokAttr::Other,
}
}
fn is_special(self) -> bool {
matches!(
self,
TokAttr::Control | TokAttr::UserDefined | TokAttr::Unknown
)
}
}
pub struct Tokenizer {
id_to_token: Vec<String>,
token_to_id: HashMap<String, u32>,
attrs: Vec<TokAttr>,
bpe_ranks: HashMap<(String, String), i32>,
special_tokens: Vec<u32>,
eos_id: u32,
bos_id: Option<u32>,
add_bos: bool,
pre: String,
split: PreSplit,
chat_template: Option<String>,
spm_style: bool,
}
#[derive(Clone, Eq, PartialEq)]
struct Bigram {
left: i32,
right: i32,
rank: i32,
text: String,
}
impl Ord for Bigram {
fn cmp(&self, other: &Self) -> Ordering {
match other.rank.cmp(&self.rank) {
Ordering::Equal => other.left.cmp(&self.left),
o => o,
}
}
}
impl PartialOrd for Bigram {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
struct Symbol {
text: String,
prev: i32,
next: i32,
n: usize, }
impl Tokenizer {
pub fn from_gguf(g: &GgufFile) -> Result<Self, String> {
let model = g
.metadata
.get("tokenizer.ggml.model")
.and_then(|v| v.as_str())
.ok_or("missing tokenizer.ggml.model")?;
if model != "gpt2" && model != "gemma4" {
return Err(format!(
"unsupported tokenizer model '{model}' (only gpt2/gemma4)"
));
}
let spm_style = model == "gemma4";
let pre = g
.metadata
.get("tokenizer.ggml.pre")
.and_then(|v| v.as_str())
.unwrap_or(if spm_style { "gemma4" } else { "default" })
.to_string();
let split = PreSplit::resolve(&pre, spm_style).map_err(|e| e.to_string())?;
let tokens = match g.metadata.get("tokenizer.ggml.tokens") {
Some(MetaValue::Array(a)) => a,
_ => return Err("missing tokenizer.ggml.tokens array".into()),
};
let n = tokens.len();
let mut id_to_token = Vec::with_capacity(n);
let mut token_to_id = HashMap::with_capacity(n);
for (i, t) in tokens.iter().enumerate() {
let s = t.as_str().ok_or("non-string in tokens[]")?.to_string();
token_to_id.entry(s.clone()).or_insert(i as u32);
id_to_token.push(s);
}
let mut attrs = vec![TokAttr::Normal; n];
if let Some(MetaValue::Array(a)) = g.metadata.get("tokenizer.ggml.token_type") {
for (i, v) in a.iter().enumerate().take(n) {
if let Some(t) = v.as_u64() {
attrs[i] = TokAttr::from_toktype(t as i64);
} else if let MetaValue::I32(t) = v {
attrs[i] = TokAttr::from_toktype(*t as i64);
}
}
}
let mut bpe_ranks = HashMap::new();
if let Some(MetaValue::Array(a)) = g.metadata.get("tokenizer.ggml.merges") {
for (i, v) in a.iter().enumerate() {
let word = v.as_str().ok_or("non-string in merges[]")?;
let bytes = word.as_bytes();
if let Some(pos) = bytes.iter().skip(1).position(|&b| b == b' ').map(|p| p + 1) {
let first = word[..pos].to_string();
let second = word[pos + 1..].to_string();
bpe_ranks.insert((first, second), i as i32);
}
}
} else {
return Err("missing tokenizer.ggml.merges array".into());
}
let mut special_tokens: Vec<u32> = (0..n as u32)
.filter(|&id| attrs[id as usize].is_special())
.collect();
special_tokens.sort_by(|&a, &b| {
id_to_token[b as usize]
.len()
.cmp(&id_to_token[a as usize].len())
});
let eos_id = g
.metadata
.get("tokenizer.ggml.eos_token_id")
.and_then(|v| v.as_u64())
.map(|v| v as u32)
.ok_or("missing tokenizer.ggml.eos_token_id")?;
let bos_id = g
.metadata
.get("tokenizer.ggml.bos_token_id")
.and_then(|v| v.as_u64())
.map(|v| v as u32);
let add_bos = g
.metadata
.get("tokenizer.ggml.add_bos_token")
.and_then(|v| match v {
MetaValue::Bool(b) => Some(*b),
_ => v.as_u64().map(|x| x != 0),
})
.unwrap_or(false);
let add_bos = add_bos || spm_style;
let chat_template = g
.metadata
.get("tokenizer.chat_template")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
Ok(Tokenizer {
id_to_token,
token_to_id,
attrs,
bpe_ranks,
special_tokens,
eos_id,
bos_id,
add_bos,
pre,
split,
chat_template,
spm_style,
})
}
pub fn from_hf_dir(dir: &std::path::Path) -> Result<Self, String> {
let tj_path = dir.join("tokenizer.json");
let text = std::fs::read_to_string(&tj_path)
.map_err(|e| format!("read {}: {e}", tj_path.display()))?;
let tj = json::parse(&text).map_err(|e| format!("{}: {e}", tj_path.display()))?;
let model = tj.get("model").ok_or("tokenizer.json: missing model")?;
if let Some(t) = model.get("type").and_then(|v| v.as_str()) {
if t != "BPE" {
return Err(format!(
"unsupported tokenizer.json model type '{t}' (only BPE)"
));
}
}
let pre_tok = tj
.get("pre_tokenizer")
.ok_or("tokenizer.json: missing pre_tokenizer")?;
if !pre_tokenizer_is_byte_level(pre_tok) {
return Err(
"tokenizer.json: pre_tokenizer is not ByteLevel — only byte-level \
BPE is supported"
.into(),
);
}
let vocab = model
.get("vocab")
.and_then(|v| v.as_obj())
.ok_or("tokenizer.json: missing model.vocab")?;
let empty: Vec<json::Value> = Vec::new();
let added = tj
.get("added_tokens")
.and_then(|v| v.as_arr())
.unwrap_or(&empty);
let mut max_id = 0u32;
for v in vocab.values() {
let id =
v.as_u64()
.ok_or("tokenizer.json: non-integer id in model.vocab")? as u32;
max_id = max_id.max(id);
}
for a in added {
if let Some(id) = a.get("id").and_then(|v| v.as_u64()) {
max_id = max_id.max(id as u32);
}
}
let n = max_id as usize + 1;
let mut id_to_token = vec![String::new(); n];
let mut token_to_id: HashMap<String, u32> = HashMap::with_capacity(n);
let mut attrs = vec![TokAttr::Normal; n];
for (tok, v) in vocab {
let id = v.as_u64().unwrap() as u32;
id_to_token[id as usize] = tok.clone();
token_to_id.entry(tok.clone()).or_insert(id);
}
for a in added {
let id =
a.get("id")
.and_then(|v| v.as_u64())
.ok_or("tokenizer.json: added_tokens entry missing id")? as u32;
let content = a
.get("content")
.and_then(|v| v.as_str())
.ok_or("tokenizer.json: added_tokens entry missing content")?;
if id_to_token[id as usize].is_empty() {
id_to_token[id as usize] = content.to_string();
}
token_to_id.entry(content.to_string()).or_insert(id);
if a.get("special").and_then(|v| v.as_bool()).unwrap_or(false) {
attrs[id as usize] = TokAttr::Control;
} else {
attrs[id as usize] = TokAttr::UserDefined;
}
}
let merges = model
.get("merges")
.and_then(|v| v.as_arr())
.ok_or("tokenizer.json: missing model.merges")?;
let mut bpe_ranks = HashMap::with_capacity(merges.len());
for (i, m) in merges.iter().enumerate() {
let (first, second) = match m {
json::Value::Str(s) => {
let bytes = s.as_bytes();
let pos = bytes
.iter()
.skip(1)
.position(|&b| b == b' ')
.map(|p| p + 1)
.ok_or_else(|| format!("tokenizer.json: merges[{i}] has no space"))?;
(s[..pos].to_string(), s[pos + 1..].to_string())
}
json::Value::Arr(a) if a.len() == 2 => {
let f = a[0]
.as_str()
.ok_or_else(|| format!("tokenizer.json: merges[{i}] non-string pair"))?;
let s2 = a[1]
.as_str()
.ok_or_else(|| format!("tokenizer.json: merges[{i}] non-string pair"))?;
(f.to_string(), s2.to_string())
}
_ => {
return Err(format!(
"tokenizer.json: merges[{i}] is neither \"a b\" string nor [a, b] pair"
));
}
};
bpe_ranks.insert((first, second), i as i32);
}
let mut special_tokens: Vec<u32> = (0..n as u32)
.filter(|&id| attrs[id as usize].is_special())
.collect();
special_tokens.sort_by(|&a, &b| {
id_to_token[b as usize]
.len()
.cmp(&id_to_token[a as usize].len())
});
let tc = std::fs::read_to_string(dir.join("tokenizer_config.json"))
.ok()
.and_then(|t| json::parse(&t).ok());
let gc = std::fs::read_to_string(dir.join("generation_config.json"))
.ok()
.and_then(|t| json::parse(&t).ok());
let tok_content = |v: &json::Value| -> Option<String> {
v.as_str().map(|s| s.to_string()).or_else(|| {
v.get("content")
.and_then(|c| c.as_str())
.map(|s| s.to_string())
})
};
let eos_from_cfg = tc
.as_ref()
.and_then(|c| c.get("eos_token"))
.and_then(&tok_content)
.and_then(|s| token_to_id.get(&s).copied());
let eos_from_gen = gc
.as_ref()
.and_then(|c| c.get("eos_token_id"))
.and_then(|v| match v {
json::Value::Num(_) => v.as_u64(),
json::Value::Arr(a) => a.first().and_then(|x| x.as_u64()),
_ => None,
})
.map(|v| v as u32);
let eos_id = eos_from_cfg.or(eos_from_gen).ok_or(
"no eos token: need tokenizer_config.json eos_token or \
generation_config.json eos_token_id",
)?;
let bos_id = tc
.as_ref()
.and_then(|c| c.get("bos_token"))
.and_then(&tok_content)
.and_then(|s| token_to_id.get(&s).copied());
let add_bos = tc
.as_ref()
.and_then(|c| c.get("add_bos_token"))
.and_then(|v| v.as_bool())
.unwrap_or(false);
let chat_template = tc
.as_ref()
.and_then(|c| c.get("chat_template"))
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.or_else(|| std::fs::read_to_string(dir.join("chat_template.jinja")).ok());
let cfg_regex = tc
.as_ref()
.and_then(|c| c.get("pretokenize_regex"))
.and_then(|v| v.as_str());
let mut tj_regexes: Vec<String> = Vec::new();
collect_split_regexes(pre_tok, &mut tj_regexes);
let pre = cfg_regex
.and_then(|r| pre_from_split_regexes(std::slice::from_ref(&r.to_string())))
.or_else(|| pre_from_split_regexes(&tj_regexes))
.unwrap_or("default");
let split = PreSplit::resolve(pre, false).map_err(|e| {
if pre == "default" {
format!(
"{e}\n (HF checkpoint {}: tokenizer_config.json pretokenize_regex = {:?}, \
tokenizer.json pre_tokenizer Split regexes = {:?} — neither matched a known \
family)",
dir.display(),
cfg_regex,
tj_regexes,
)
} else {
e.to_string()
}
})?;
Ok(Tokenizer {
id_to_token,
token_to_id,
attrs,
bpe_ranks,
special_tokens,
eos_id,
bos_id,
add_bos,
pre: pre.to_string(),
split,
chat_template,
spm_style: false,
})
}
pub fn eos_id(&self) -> u32 {
self.eos_id
}
pub fn id_of(&self, piece: &str) -> Option<u32> {
self.token_to_id.get(piece).copied()
}
pub fn eog_ids(&self) -> Vec<u32> {
let mut ids = vec![self.eos_id];
for t in ["<|im_end|>", "<turn|>", "<end_of_turn>"] {
if let Some(&id) = self.token_to_id.get(t) {
if !ids.contains(&id) {
ids.push(id);
}
}
}
ids
}
pub fn bos_id(&self) -> Option<u32> {
self.bos_id
}
pub fn vocab_size(&self) -> usize {
self.id_to_token.len()
}
pub fn pre(&self) -> &str {
&self.pre
}
pub fn split(&self) -> PreSplit {
self.split
}
pub fn chat_template(&self) -> Option<&str> {
self.chat_template.as_deref()
}
#[inline]
fn text_to_token(&self, s: &str) -> Option<u32> {
self.token_to_id.get(s).copied()
}
fn find_bpe_rank(&self, left: &str, right: &str) -> i32 {
self.bpe_ranks
.get(&(left.to_string(), right.to_string()))
.copied()
.unwrap_or(-1)
}
pub fn encode(&self, text: &str, add_special: bool) -> Vec<u32> {
self.encode_special(text, add_special, true)
}
pub fn encode_special(&self, text: &str, add_special: bool, parse_special: bool) -> Vec<u32> {
let mut output: Vec<u32> = Vec::new();
if add_special && self.add_bos {
if let Some(b) = self.bos_id {
output.push(b);
}
}
if text.is_empty() {
return output;
}
for frag in self.st_partition(text, parse_special) {
match frag {
Fragment::Token(id) => output.push(id),
Fragment::Text(span) => self.bpe_tokenize(&span, &mut output),
}
}
output
}
fn st_partition(&self, text: &str, parse_special: bool) -> Vec<Fragment> {
let mut frags = vec![Fragment::Text(text.to_string())];
for &sid in &self.special_tokens {
let attr = self.attrs[sid as usize];
if !parse_special && matches!(attr, TokAttr::Control | TokAttr::Unknown) {
continue;
}
let needle = &self.id_to_token[sid as usize];
if needle.is_empty() {
continue;
}
let mut next: Vec<Fragment> = Vec::with_capacity(frags.len());
for f in frags.drain(..) {
match f {
Fragment::Token(id) => next.push(Fragment::Token(id)),
Fragment::Text(s) => {
let mut rest: &str = &s;
let mut acc = String::new();
while let Some(m) = rest.find(needle.as_str()) {
acc.push_str(&rest[..m]);
if !acc.is_empty() {
next.push(Fragment::Text(std::mem::take(&mut acc)));
}
next.push(Fragment::Token(sid));
rest = &rest[m + needle.len()..];
}
acc.push_str(rest);
if !acc.is_empty() {
next.push(Fragment::Text(acc));
}
}
}
}
frags = next;
}
frags
}
fn bpe_tokenize(&self, text: &str, output: &mut Vec<u32>) {
if self.spm_style {
let escaped: String = text
.chars()
.map(|c| if c == ' ' { '\u{2581}' } else { c })
.collect();
let mut words: Vec<String> = Vec::new();
let mut cur = String::new();
let mut cur_nl: Option<bool> = None;
for c in escaped.chars() {
let nl = c == '\n';
if cur_nl != Some(nl) && !cur.is_empty() {
words.push(std::mem::take(&mut cur));
}
cur_nl = Some(nl);
cur.push(c);
}
if !cur.is_empty() {
words.push(cur);
}
for word in &words {
if word.chars().all(|c| c == '\n') {
if let Some(tok) = self.text_to_token(word) {
output.push(tok);
continue;
}
}
self.bpe_merge_word(word, output);
}
return;
}
let words: Vec<String> = match self.split {
PreSplit::Qwen35 => unicode::split_qwen35(text),
PreSplit::DeepseekV3 => unicode::split_deepseek_v3(text),
PreSplit::UnknownFallbackQwen35 => unicode::split_qwen35(text),
PreSplit::Spm => unreachable!("PreSplit::Spm implies spm_style, handled above"),
};
for word in &words {
let word = unicode::byte_encode(word);
self.bpe_merge_word(&word, output);
}
}
fn bpe_merge_word(&self, word: &str, output: &mut Vec<u32>) {
{
let word = word.to_string();
let chars: Vec<char> = word.chars().collect();
let mut symbols: Vec<Symbol> = Vec::with_capacity(chars.len());
for (i, &c) in chars.iter().enumerate() {
symbols.push(Symbol {
text: c.to_string(),
prev: i as i32 - 1,
next: if i + 1 == chars.len() {
-1
} else {
i as i32 + 1
},
n: 1,
});
}
let mut queue: BinaryHeap<Bigram> = BinaryHeap::new();
for i in 1..symbols.len() {
self.add_bigram(&symbols, i as i32 - 1, i as i32, &mut queue);
}
while let Some(bigram) = queue.pop() {
let li = bigram.left as usize;
let ri = bigram.right as usize;
if symbols[li].n == 0 || symbols[ri].n == 0 {
continue;
}
let combined = format!("{}{}", symbols[li].text, symbols[ri].text);
if combined != bigram.text {
continue; }
symbols[li].text = combined;
symbols[li].n += symbols[ri].n;
symbols[ri].n = 0;
let r_next = symbols[ri].next;
symbols[li].next = r_next;
if r_next >= 0 {
symbols[r_next as usize].prev = bigram.left;
}
let l_prev = symbols[li].prev;
let l_next = symbols[li].next;
self.add_bigram(&symbols, l_prev, bigram.left, &mut queue);
self.add_bigram(&symbols, bigram.left, l_next, &mut queue);
}
for sym in &symbols {
if sym.n == 0 {
continue;
}
match self.text_to_token(&sym.text) {
Some(tok) => output.push(tok),
None => {
for b in sym.text.bytes() {
let bs = if self.spm_style {
format!("<0x{b:02X}>") } else {
(b as char).to_string()
};
if let Some(t) = self.text_to_token(&bs) {
output.push(t);
}
}
}
}
}
}
}
fn add_bigram(
&self,
symbols: &[Symbol],
left: i32,
right: i32,
queue: &mut BinaryHeap<Bigram>,
) {
if left == -1 || right == -1 {
return;
}
let lt = &symbols[left as usize].text;
let rt = &symbols[right as usize].text;
let rank = self.find_bpe_rank(lt, rt);
if rank < 0 {
return;
}
queue.push(Bigram {
left,
right,
rank,
text: format!("{lt}{rt}"),
});
}
pub fn decode(&self, ids: &[u32]) -> String {
self.decode_special(ids, true)
}
pub fn token_is_control(&self, id: u32) -> bool {
match self.attrs.get(id as usize) {
Some(TokAttr::Control) | Some(TokAttr::Unknown) => true,
_ => false,
}
}
pub fn decode_special(&self, ids: &[u32], special: bool) -> String {
String::from_utf8_lossy(&self.decode_bytes_special(ids, special)).into_owned()
}
pub fn decode_bytes_special(&self, ids: &[u32], special: bool) -> Vec<u8> {
let mut bytes: Vec<u8> = Vec::new();
for &id in ids {
let i = id as usize;
if i >= self.id_to_token.len() {
continue;
}
let attr = self.attrs[i];
let piece = &self.id_to_token[i];
match attr {
TokAttr::Normal | TokAttr::Byte => {
if self.spm_style {
if matches!(attr, TokAttr::Byte)
|| (piece.len() == 6
&& piece.starts_with("<0x")
&& piece.ends_with('>'))
{
if let Ok(b) = u8::from_str_radix(&piece[3..5], 16) {
bytes.push(b);
continue;
}
}
for c in piece.chars() {
if c == '\u{2581}' {
bytes.push(b' ');
} else {
let mut buf = [0u8; 4];
bytes.extend_from_slice(c.encode_utf8(&mut buf).as_bytes());
}
}
} else {
self.piece_to_bytes(piece, &mut bytes);
}
}
TokAttr::UserDefined => {
bytes.extend_from_slice(piece.as_bytes());
}
TokAttr::Control | TokAttr::Unknown => {
if special {
bytes.extend_from_slice(piece.as_bytes());
}
}
TokAttr::Other => {}
}
}
bytes
}
fn piece_to_bytes(&self, piece: &str, out: &mut Vec<u8>) {
for c in piece.chars() {
match unicode::unicode_to_byte(c) {
Some(b) => out.push(b),
None => {
let mut buf = [0u8; 4];
out.extend_from_slice(c.encode_utf8(&mut buf).as_bytes());
}
}
}
}
pub fn apply_chat_template(
&self,
messages: &[(&str, &str)],
add_generation_prompt: bool,
) -> String {
chat::apply_chat_template_str(
self.chat_template.as_deref(),
messages,
add_generation_prompt,
)
}
pub fn apply_chat_template_tools(
&self,
turns: &[chat::Turn],
add_generation_prompt: bool,
tools_json: &[String],
think: chat::ThinkMode,
reasoning_effort: Option<&str>,
) -> Result<String, String> {
chat::apply_chat_template_tools(
self.chat_template.as_deref(),
turns,
add_generation_prompt,
tools_json,
think,
reasoning_effort,
)
}
#[allow(clippy::too_many_arguments)]
pub fn apply_chat_template_tools_ex(
&self,
turns: &[chat::Turn],
add_generation_prompt: bool,
tools_json: &[String],
tools_struct: &[chat::Val],
think: chat::ThinkMode,
reasoning_effort: Option<&str>,
) -> Result<String, String> {
chat::apply_chat_template_tools_ex(
self.chat_template.as_deref(),
turns,
add_generation_prompt,
tools_json,
tools_struct,
think,
reasoning_effort,
)
}
}
enum Fragment {
Text(String),
Token(u32),
}
fn collect_split_regexes(pt: &json::Value, out: &mut Vec<String>) {
match pt.get("type").and_then(|v| v.as_str()) {
Some("Sequence") => {
if let Some(arr) = pt.get("pretokenizers").and_then(|v| v.as_arr()) {
for step in arr {
collect_split_regexes(step, out);
}
}
}
Some("Split") => {
if let Some(r) = pt
.get("pattern")
.and_then(|p| p.get("Regex"))
.and_then(|v| v.as_str())
{
out.push(r.to_string());
}
}
_ => {}
}
}
fn pre_from_split_regexes(regexes: &[String]) -> Option<&'static str> {
match regexes {
[one] if one == QWEN35_PRETOKENIZE_REGEX => Some("qwen35"),
[one] if one == QWEN2_PRETOKENIZE_REGEX => Some("qwen2"),
[a, b, c]
if a == DEEPSEEK_V3_SPLIT_REGEXES[0]
&& b == DEEPSEEK_V3_SPLIT_REGEXES[1]
&& c == DEEPSEEK_V3_SPLIT_REGEXES[2] =>
{
Some("deepseek-v3")
}
_ => None,
}
}
fn pre_tokenizer_is_byte_level(pt: &json::Value) -> bool {
match pt.get("type").and_then(|v| v.as_str()) {
Some("ByteLevel") => true,
Some("Sequence") => pt
.get("pretokenizers")
.and_then(|v| v.as_arr())
.map(|arr| arr.iter().any(pre_tokenizer_is_byte_level))
.unwrap_or(false),
_ => false,
}
}
#[cfg(test)]
mod pretokenizer_tests {
use super::*;
#[test]
fn every_supported_pre_resolves() {
assert_eq!(
PreSplit::resolve_with("qwen35", false, false),
Ok(PreSplit::Qwen35)
);
assert_eq!(
PreSplit::resolve_with("qwen2", false, false),
Ok(PreSplit::Qwen35)
);
assert_eq!(
PreSplit::resolve_with("deepseek-v3", false, false),
Ok(PreSplit::DeepseekV3)
);
assert_eq!(
PreSplit::resolve_with("gemma4", true, false),
Ok(PreSplit::Spm)
);
assert_eq!(
SUPPORTED_PRETOKENIZERS,
&["qwen35", "qwen2", "deepseek-v3", "gemma4"]
);
}
#[test]
fn unknown_pre_is_a_typed_error() {
let err =
PreSplit::resolve_with("llama4", false, false).expect_err("llama4 has no ported split");
assert_eq!(
err,
UnknownPretokenizer {
pre: "llama4".into(),
spm_style: false
}
);
let msg = err.to_string();
assert!(msg.contains("'llama4'"), "{msg}");
for supported in SUPPORTED_PRETOKENIZERS {
assert!(
msg.contains(supported),
"error must list {supported}: {msg}"
);
}
assert!(msg.contains(ALLOW_UNKNOWN_PRETOKENIZER_ENV), "{msg}");
let _: &dyn std::error::Error = &err;
}
#[test]
fn pre_and_vocab_model_must_agree() {
assert!(PreSplit::resolve_with("qwen35", true, false).is_err());
assert!(PreSplit::resolve_with("gemma4", false, false).is_err());
assert!(PreSplit::resolve_with("default", false, false).is_err());
assert!(PreSplit::resolve_with("", false, false).is_err());
}
#[test]
fn opt_out_loads_with_a_fallback_marker() {
assert_eq!(
PreSplit::resolve_with("llama4", false, true),
Ok(PreSplit::UnknownFallbackQwen35)
);
assert_eq!(
PreSplit::resolve_with("qwen35", true, true),
Ok(PreSplit::UnknownFallbackQwen35)
);
}
#[test]
fn opt_out_env_gate() {
unsafe { std::env::remove_var(ALLOW_UNKNOWN_PRETOKENIZER_ENV) };
assert!(!allow_unknown_pretokenizer());
unsafe { std::env::set_var(ALLOW_UNKNOWN_PRETOKENIZER_ENV, "0") };
assert!(!allow_unknown_pretokenizer());
unsafe { std::env::set_var(ALLOW_UNKNOWN_PRETOKENIZER_ENV, "1") };
assert!(allow_unknown_pretokenizer());
assert_eq!(
PreSplit::resolve("llama4", false),
Ok(PreSplit::UnknownFallbackQwen35)
);
unsafe { std::env::remove_var(ALLOW_UNKNOWN_PRETOKENIZER_ENV) };
assert!(PreSplit::resolve("llama4", false).is_err());
}
#[test]
fn split_regex_identification_is_exact() {
let s = |v: &[&str]| v.iter().map(|x| x.to_string()).collect::<Vec<_>>();
assert_eq!(
pre_from_split_regexes(&s(&[QWEN35_PRETOKENIZE_REGEX])),
Some("qwen35")
);
assert_eq!(
pre_from_split_regexes(&s(&[QWEN2_PRETOKENIZE_REGEX])),
Some("qwen2")
);
assert_eq!(
pre_from_split_regexes(&s(&DEEPSEEK_V3_SPLIT_REGEXES)),
Some("deepseek-v3")
);
assert_eq!(
pre_from_split_regexes(&s(&[
DEEPSEEK_V3_SPLIT_REGEXES[1],
DEEPSEEK_V3_SPLIT_REGEXES[0],
DEEPSEEK_V3_SPLIT_REGEXES[2],
])),
None
);
assert_eq!(
pre_from_split_regexes(&s(&[
DEEPSEEK_V3_SPLIT_REGEXES[0],
DEEPSEEK_V3_SPLIT_REGEXES[1]
])),
None
);
let mut near = QWEN35_PRETOKENIZE_REGEX.to_string();
near.push('x');
assert_eq!(pre_from_split_regexes(&s(&[&near])), None);
assert_eq!(pre_from_split_regexes(&[]), None);
assert_ne!(QWEN2_PRETOKENIZE_REGEX, QWEN35_PRETOKENIZE_REGEX);
}
#[test]
fn collect_split_regexes_walks_in_order() {
let src = r#"{"type":"Sequence","pretokenizers":[
{"type":"Split","pattern":{"Regex":"A"},"behavior":"Isolated"},
{"type":"Split","pattern":{"String":" "},"behavior":"Isolated"},
{"type":"Digits","individual_digits":true},
{"type":"Sequence","pretokenizers":[
{"type":"Split","pattern":{"Regex":"B"},"behavior":"Isolated"}
]},
{"type":"ByteLevel","add_prefix_space":false}
]}"#;
let v = json::parse(src).unwrap();
let mut out = Vec::new();
collect_split_regexes(&v, &mut out);
assert_eq!(out, vec!["A".to_string(), "B".to_string()]);
}
}
#[cfg(test)]
mod hf_tests {
use super::*;
const TOKENIZER_JSON: &str = r#"{
"version": "1.0",
"added_tokens": [
{"id": 15, "content": "<|end|>", "special": true},
{"id": 16, "content": "<think>", "special": false}
],
"pre_tokenizer": {
"type": "Sequence",
"pretokenizers": [
{"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"},
{"type": "ByteLevel", "add_prefix_space": false, "trim_offsets": false}
]
},
"model": {
"type": "BPE",
"vocab": {
"h": 0, "e": 1, "l": 2, "o": 3, "Ġ": 4, "w": 5, "r": 6, "d": 7,
"he": 8, "ll": 9, "hell": 10, "hello": 11, "Ġw": 12, "or": 13, "!": 14
},
"merges": [
"h e",
["l", "l"],
"he ll",
["hell", "o"],
["Ġ", "w"],
"o r"
]
}
}"#;
fn write_fixture(
name: &str,
tokenizer_config: Option<&str>,
generation_config: Option<&str>,
jinja: Option<&str>,
) -> std::path::PathBuf {
let dir = std::env::temp_dir().join(format!("memra-tok-hf-{name}-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join("tokenizer.json"), TOKENIZER_JSON).unwrap();
if let Some(tc) = tokenizer_config {
std::fs::write(dir.join("tokenizer_config.json"), tc).unwrap();
}
if let Some(gc) = generation_config {
std::fs::write(dir.join("generation_config.json"), gc).unwrap();
}
if let Some(j) = jinja {
std::fs::write(dir.join("chat_template.jinja"), j).unwrap();
}
dir
}
#[test]
fn hf_dir_encode_decode_roundtrip_and_specials() {
let tc = r#"{
"eos_token": {"content": "<|end|>", "lstrip": false},
"add_bos_token": false,
"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+",
"chat_template": "{{ messages }}<|end|>"
}"#;
let dir = write_fixture("full", Some(tc), None, None);
let tok = Tokenizer::from_hf_dir(&dir).expect("from_hf_dir");
assert_eq!(tok.eos_id(), 15);
assert_eq!(tok.bos_id(), None);
assert_eq!(tok.pre(), "qwen35");
assert_eq!(tok.vocab_size(), 17); assert_eq!(tok.chat_template(), Some("{{ messages }}<|end|>"));
let ids = tok.encode("hello world", true);
assert_eq!(ids, vec![11, 12, 13, 2, 7]);
assert_eq!(tok.decode(&ids), "hello world");
let ids = tok.encode("hello<|end|> world", true);
assert_eq!(ids, vec![11, 15, 12, 13, 2, 7]);
assert_eq!(tok.decode_special(&ids, true), "hello<|end|> world");
assert_eq!(tok.decode_special(&ids, false), "hello world");
assert_eq!(tok.decode(&[16]), "<think>");
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn hf_dir_generation_config_eos_fallback_and_jinja() {
let gc = r#"{"eos_token_id": [15, 14]}"#;
let dir = write_fixture("genconf", None, Some(gc), Some("JINJA {{ messages }}"));
let tok = Tokenizer::from_hf_dir(&dir).expect("from_hf_dir");
assert_eq!(tok.eos_id(), 15);
assert!(!tok.encode("hello", true).is_empty());
assert_eq!(tok.chat_template(), Some("JINJA {{ messages }}"));
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn hf_dir_identifies_deepseek_v3_from_tokenizer_json() {
let dsv3_pt = r##""pre_tokenizer": {
"type": "Sequence",
"pretokenizers": [
{"type": "Split", "pattern": {"Regex": "\\p{N}{1,3}"}, "behavior": "Isolated"},
{"type": "Split", "pattern": {"Regex": "[一-龥-ゟ゠-ヿ]+"}, "behavior": "Isolated"},
{"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"},
{"type": "ByteLevel", "add_prefix_space": false, "trim_offsets": true, "use_regex": false}
]
},"##;
let open = TOKENIZER_JSON.find(r#""pre_tokenizer""#).unwrap();
let close = TOKENIZER_JSON.find(r#""model""#).unwrap();
let json = format!(
"{}{}\n {}",
&TOKENIZER_JSON[..open],
dsv3_pt,
&TOKENIZER_JSON[close..]
);
let dir = std::env::temp_dir().join(format!("memra-tok-hf-dsv3-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join("tokenizer.json"), &json).unwrap();
std::fs::write(
dir.join("generation_config.json"),
r#"{"eos_token_id": 15}"#,
)
.unwrap();
let tok = Tokenizer::from_hf_dir(&dir).expect("from_hf_dir");
assert_eq!(tok.pre(), "deepseek-v3");
assert_eq!(tok.split(), PreSplit::DeepseekV3);
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn hf_dir_identifies_qwen2_regex() {
let json = TOKENIZER_JSON.replace(
r"[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+",
r"[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+",
);
assert_ne!(json, TOKENIZER_JSON, "the qwen2 substitution must apply");
let dir = std::env::temp_dir().join(format!("memra-tok-hf-qwen2-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join("tokenizer.json"), &json).unwrap();
std::fs::write(
dir.join("generation_config.json"),
r#"{"eos_token_id": 15}"#,
)
.unwrap();
let tok = Tokenizer::from_hf_dir(&dir).expect("from_hf_dir");
assert_eq!(tok.pre(), "qwen2");
assert_eq!(
tok.split(),
PreSplit::Qwen35,
"qwen2 rides the qwen35 split"
);
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn hf_dir_refuses_unidentifiable_pretokenizer() {
let json = TOKENIZER_JSON.replace(r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|", "SOMETHING-ELSE|");
assert_ne!(json, TOKENIZER_JSON);
let dir = std::env::temp_dir().join(format!("memra-tok-hf-unk-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join("tokenizer.json"), &json).unwrap();
std::fs::write(
dir.join("generation_config.json"),
r#"{"eos_token_id": 15}"#,
)
.unwrap();
let err = match Tokenizer::from_hf_dir(&dir) {
Ok(_) => panic!("unidentifiable pre must refuse to load"),
Err(e) => e,
};
assert!(
err.contains("unsupported tokenizer.ggml.pre 'default'"),
"{err}"
);
assert!(
err.contains("SOMETHING-ELSE"),
"error must quote the regex: {err}"
);
assert!(err.contains("MEMRA_ALLOW_UNKNOWN_PRETOKENIZER"), "{err}");
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn staged_checkpoints_resolve_their_own_pretokenizer() {
let cases: &[(&str, &str)] = &[
(
"/data/ai-ml/hf-models/hy3-layer103p5-sparse-source",
"deepseek-v3",
),
("/data/ai-ml/hf-models/qwen3-1.7b-blk128fp8-synth", "qwen2"),
("/data/ai-ml/hf-models/qwen35-9b-hf", "qwen35"),
];
let mut ran = 0;
for (path, want) in cases {
let dir = std::path::Path::new(path);
if !dir.join("tokenizer.json").exists() {
eprintln!("skip: {path} not staged");
continue;
}
let tok = Tokenizer::from_hf_dir(dir).unwrap_or_else(|e| panic!("{path}: {e}"));
assert_eq!(tok.pre(), *want, "{path}");
ran += 1;
}
eprintln!("staged_checkpoints_resolve_their_own_pretokenizer: {ran}/3 cases ran");
}
#[test]
fn hf_dir_rejects_non_byte_level() {
let dir = std::env::temp_dir().join(format!("memra-tok-hf-nonbl-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
let bad = TOKENIZER_JSON.replace("\"ByteLevel\"", "\"Metaspace\"");
std::fs::write(dir.join("tokenizer.json"), bad).unwrap();
assert!(Tokenizer::from_hf_dir(&dir).is_err());
let _ = std::fs::remove_dir_all(&dir);
}
}