377a04b81f
Parse the model's `pre_tokenizer` section to extract its Split regex instead of hardcoding the GPT-2 pattern. The gpt-oss-20b model uses a GPT-4-style regex that produces different word boundaries, causing a 1-token prompt mismatch vs HuggingFace (136 → 135 tokens, now aligned). Unsupported lookahead `(?!\S)` is stripped — it only affects trailing whitespace edge cases. Falls back to the old GPT-2/Qwen heuristic if the model regex fails to compile or is absent. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
442 lines
16 KiB
Rust
442 lines
16 KiB
Rust
use regex::Regex;
|
|
use serde::Deserialize;
|
|
use std::collections::HashMap;
|
|
use std::path::Path;
|
|
|
|
pub struct Tokenizer {
|
|
encoder: HashMap<Vec<u8>, u32>,
|
|
decoder: Vec<Vec<u8>>,
|
|
merge_ranks: HashMap<(u32, u32), usize>,
|
|
special_tokens: HashMap<String, u32>,
|
|
#[allow(dead_code)]
|
|
special_token_ids: HashMap<u32, String>,
|
|
pre_tokenize_re: Regex,
|
|
eos_token_id: Option<u32>,
|
|
eos_token_ids: Vec<u32>,
|
|
byte_fallback: bool,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct TokenizerJson {
|
|
model: ModelSection,
|
|
#[serde(default)]
|
|
added_tokens: Vec<AddedToken>,
|
|
#[serde(default)]
|
|
pre_tokenizer: Option<PreTokenizerSection>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct PreTokenizerSection {
|
|
#[serde(default, rename = "type")]
|
|
kind: Option<String>,
|
|
#[serde(default)]
|
|
pattern: Option<PatternSpec>,
|
|
#[serde(default)]
|
|
pretokenizers: Option<Vec<PreTokenizerSection>>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct PatternSpec {
|
|
#[serde(rename = "Regex")]
|
|
regex: Option<String>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct ModelSection {
|
|
vocab: HashMap<String, u32>,
|
|
merges: Vec<MergeEntry>,
|
|
#[serde(default)]
|
|
byte_fallback: bool,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
#[serde(untagged)]
|
|
enum MergeEntry {
|
|
Str(String),
|
|
Pair(Vec<String>),
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct AddedToken {
|
|
id: u32,
|
|
content: String,
|
|
#[allow(dead_code)]
|
|
special: bool,
|
|
}
|
|
|
|
impl Tokenizer {
|
|
pub fn from_file(path: &Path) -> Self {
|
|
let data = std::fs::read_to_string(path)
|
|
.unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display()));
|
|
let tj: TokenizerJson = serde_json::from_str(&data)
|
|
.unwrap_or_else(|e| panic!("failed to parse tokenizer.json: {e}"));
|
|
|
|
// Build encoder: token bytes → ID
|
|
// All HF tokenizers use GPT-2 byte-to-unicode mapping for vocab keys.
|
|
let mut encoder = HashMap::new();
|
|
for (token_str, &id) in &tj.model.vocab {
|
|
let bytes = token_str_to_bytes(token_str);
|
|
encoder.insert(bytes, id);
|
|
}
|
|
|
|
// Build decoder: ID → token bytes
|
|
let max_id = tj.model.vocab.values().copied().max().unwrap_or(0);
|
|
let added_max = tj.added_tokens.iter().map(|t| t.id).max().unwrap_or(0);
|
|
let vocab_size = (max_id.max(added_max) + 1) as usize;
|
|
let mut decoder = vec![vec![]; vocab_size];
|
|
for (token_str, &id) in &tj.model.vocab {
|
|
decoder[id as usize] = token_str_to_bytes(token_str);
|
|
}
|
|
|
|
// Parse merges (supports both "a b" string format and ["a", "b"] array format)
|
|
let byte_fallback = tj.model.byte_fallback;
|
|
let mut merge_ranks = HashMap::new();
|
|
for (rank, entry) in tj.model.merges.iter().enumerate() {
|
|
let (a_str, b_str) = match entry {
|
|
MergeEntry::Str(s) => {
|
|
let parts: Vec<&str> = s.splitn(2, ' ').collect();
|
|
if parts.len() != 2 { continue; }
|
|
(parts[0].to_string(), parts[1].to_string())
|
|
}
|
|
MergeEntry::Pair(v) => {
|
|
if v.len() != 2 { continue; }
|
|
(v[0].clone(), v[1].clone())
|
|
}
|
|
};
|
|
let a_bytes = token_str_to_bytes(&a_str);
|
|
let b_bytes = token_str_to_bytes(&b_str);
|
|
if let (Some(&a_id), Some(&b_id)) = (encoder.get(&a_bytes), encoder.get(&b_bytes)) {
|
|
merge_ranks.insert((a_id, b_id), rank);
|
|
}
|
|
}
|
|
|
|
// Added tokens are matched as indivisible tokens by HF tokenizers,
|
|
// even when their `special` flag is false (for example Qwen3's
|
|
// <think> and </think> tokens).
|
|
let mut special_tokens = HashMap::new();
|
|
let mut special_token_ids = HashMap::new();
|
|
for at in &tj.added_tokens {
|
|
special_tokens.insert(at.content.clone(), at.id);
|
|
special_token_ids.insert(at.id, at.content.clone());
|
|
decoder.resize(decoder.len().max(at.id as usize + 1), vec![]);
|
|
decoder[at.id as usize] = at.content.as_bytes().to_vec();
|
|
}
|
|
// End-of-generation tokens, in priority order. Families differ:
|
|
// Qwen uses <|im_end|>, Llama <|end_of_text|>, GPT-2 <|endoftext|>.
|
|
// gpt-oss (harmony) ends the assistant turn with <|return|> and also
|
|
// treats <|call|> (tool call) and <|endoftext|> as terminators
|
|
// (see generation_config.json eos_token_id = [200002, 199999, 200012]).
|
|
let eos_names = [
|
|
"<|im_end|>",
|
|
"<|end_of_text|>",
|
|
"<|return|>",
|
|
"<|call|>",
|
|
"<|endoftext|>",
|
|
];
|
|
let mut eos_token_ids: Vec<u32> = Vec::new();
|
|
for name in eos_names {
|
|
if let Some(&id) = special_tokens.get(name) {
|
|
if !eos_token_ids.contains(&id) {
|
|
eos_token_ids.push(id);
|
|
}
|
|
}
|
|
}
|
|
let eos_token_id = eos_token_ids.first().copied();
|
|
|
|
// Pre-tokenization regex: prefer the model's own regex from tokenizer.json,
|
|
// fall back to GPT-2/Qwen heuristic if not present or unsupported.
|
|
let model_regex = tj.pre_tokenizer.as_ref().and_then(|pt| {
|
|
// Direct Split with regex
|
|
if pt.kind.as_deref() == Some("Split") {
|
|
return pt.pattern.as_ref().and_then(|p| p.regex.clone());
|
|
}
|
|
// Sequence → find the Split entry
|
|
if let Some(subs) = &pt.pretokenizers {
|
|
for sub in subs {
|
|
if sub.kind.as_deref() == Some("Split") {
|
|
if let Some(r) = sub.pattern.as_ref().and_then(|p| p.regex.clone()) {
|
|
return Some(r);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
None
|
|
});
|
|
|
|
let pre_tokenize_re = if let Some(ref pat) = model_regex {
|
|
// Strip unsupported lookahead (?!\S) — Rust regex doesn't support it.
|
|
// The lookahead only affects trailing-whitespace edge cases.
|
|
let cleaned = pat.replace(r"(?!\S)", "");
|
|
match Regex::new(&cleaned) {
|
|
Ok(re) => re,
|
|
Err(e) => {
|
|
eprintln!("warning: model pre_tokenizer regex failed ({e}), using fallback");
|
|
if byte_fallback {
|
|
Regex::new(r"[\p{L}\p{N}]+|[^\s\p{L}\p{N}]|\s+").unwrap()
|
|
} else {
|
|
Regex::new(r"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+").unwrap()
|
|
}
|
|
}
|
|
}
|
|
} else if byte_fallback {
|
|
Regex::new(r"[\p{L}\p{N}]+|[^\s\p{L}\p{N}]|\s+").unwrap()
|
|
} else {
|
|
Regex::new(r"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+").unwrap()
|
|
};
|
|
|
|
Self {
|
|
encoder,
|
|
decoder,
|
|
merge_ranks,
|
|
special_tokens,
|
|
special_token_ids,
|
|
pre_tokenize_re,
|
|
eos_token_id,
|
|
eos_token_ids,
|
|
byte_fallback,
|
|
}
|
|
}
|
|
|
|
pub fn encode(&self, text: &str) -> Vec<u32> {
|
|
let mut tokens = Vec::new();
|
|
|
|
// Check for special tokens first (split around them)
|
|
let mut remaining = text;
|
|
while !remaining.is_empty() {
|
|
// Find earliest special token
|
|
let mut earliest: Option<(usize, &str, u32)> = None;
|
|
for (st, &id) in &self.special_tokens {
|
|
if let Some(pos) = remaining.find(st.as_str()) {
|
|
if earliest.is_none() || pos < earliest.unwrap().0 {
|
|
earliest = Some((pos, st, id));
|
|
}
|
|
}
|
|
}
|
|
|
|
if let Some((pos, st, id)) = earliest {
|
|
if pos > 0 {
|
|
self.encode_ordinary(&remaining[..pos], &mut tokens);
|
|
}
|
|
tokens.push(id);
|
|
remaining = &remaining[pos + st.len()..];
|
|
} else {
|
|
self.encode_ordinary(remaining, &mut tokens);
|
|
break;
|
|
}
|
|
}
|
|
|
|
tokens
|
|
}
|
|
|
|
fn encode_ordinary(&self, text: &str, out: &mut Vec<u32>) {
|
|
for mat in self.pre_tokenize_re.find_iter(text) {
|
|
let word = mat.as_str();
|
|
// Try to encode the whole word first
|
|
if let Some(&id) = self.encoder.get(word.as_bytes()) {
|
|
out.push(id);
|
|
continue;
|
|
}
|
|
// Fall back to per-byte encoding
|
|
let word_bytes: Vec<u8> = word.bytes().collect();
|
|
let mut token_ids: Vec<u32> = word_bytes.iter().filter_map(|&b| {
|
|
if let Some(&id) = self.encoder.get(&vec![b]) {
|
|
Some(id)
|
|
} else if self.byte_fallback {
|
|
let hex_token = format!("<0x{:02X}>", b);
|
|
if let Some(&id) = self.special_tokens.get(&hex_token) {
|
|
Some(id)
|
|
} else if let Some(&id) = self.encoder.get(hex_token.as_bytes()) {
|
|
Some(id)
|
|
} else if let Some(&unk_id) = self.special_tokens.get("<unk>") {
|
|
eprintln!("warning: byte 0x{b:02X} not in vocab, using <unk> token");
|
|
Some(unk_id)
|
|
} else {
|
|
eprintln!("warning: byte 0x{b:02X} not in vocab and no fallback token, using token 0");
|
|
Some(0)
|
|
}
|
|
} else {
|
|
eprintln!("warning: byte {b} (0x{b:02X}) not in vocab, skipping");
|
|
None
|
|
}
|
|
}).collect();
|
|
|
|
// BPE merges
|
|
loop {
|
|
if token_ids.len() < 2 { break; }
|
|
let mut best_rank = usize::MAX;
|
|
let mut best_idx = 0;
|
|
for i in 0..token_ids.len() - 1 {
|
|
if let Some(&rank) = self.merge_ranks.get(&(token_ids[i], token_ids[i + 1])) {
|
|
if rank < best_rank {
|
|
best_rank = rank;
|
|
best_idx = i;
|
|
}
|
|
}
|
|
}
|
|
if best_rank == usize::MAX { break; }
|
|
|
|
let merged_bytes = [
|
|
self.decoder[token_ids[best_idx] as usize].as_slice(),
|
|
self.decoder[token_ids[best_idx + 1] as usize].as_slice(),
|
|
].concat();
|
|
let merged_id = *self.encoder.get(&merged_bytes).unwrap_or_else(|| {
|
|
panic!("merged token not in vocab");
|
|
});
|
|
token_ids[best_idx] = merged_id;
|
|
token_ids.remove(best_idx + 1);
|
|
}
|
|
|
|
out.extend_from_slice(&token_ids);
|
|
}
|
|
}
|
|
|
|
pub fn decode(&self, token_ids: &[u32]) -> String {
|
|
let mut bytes = Vec::new();
|
|
for &id in token_ids {
|
|
if let Some(b) = self.decoder.get(id as usize) {
|
|
bytes.extend_from_slice(b);
|
|
}
|
|
}
|
|
String::from_utf8_lossy(&bytes).into_owned()
|
|
}
|
|
|
|
pub fn decode_token_stream(&self, token_id: u32, pending: &mut Vec<u8>) -> String {
|
|
if let Some(bytes) = self.decoder.get(token_id as usize) {
|
|
pending.extend_from_slice(bytes);
|
|
}
|
|
take_valid_utf8(pending)
|
|
}
|
|
|
|
pub fn flush_decode_stream(&self, pending: &mut Vec<u8>) -> String {
|
|
let text = String::from_utf8_lossy(pending).into_owned();
|
|
pending.clear();
|
|
text
|
|
}
|
|
|
|
pub fn eos_token_id(&self) -> Option<u32> {
|
|
self.eos_token_id
|
|
}
|
|
|
|
/// True if `id` is any end-of-generation token (a model may have several;
|
|
/// gpt-oss/harmony ends on <|return|>, <|call|>, or <|endoftext|>).
|
|
pub fn is_eos(&self, id: u32) -> bool {
|
|
self.eos_token_ids.contains(&id)
|
|
}
|
|
|
|
pub fn vocab_size(&self) -> usize {
|
|
self.decoder.len()
|
|
}
|
|
|
|
pub fn special_token_id(&self, name: &str) -> Option<u32> {
|
|
self.special_tokens.get(name).copied()
|
|
}
|
|
}
|
|
|
|
/// Convert a token string from HF vocab (which uses Unicode replacements for bytes)
|
|
/// back to raw bytes. GPT-2 uses a byte-to-unicode mapping where e.g. byte 0x20 (space)
|
|
/// is represented as 'Ġ' (U+0120).
|
|
fn token_str_to_bytes(s: &str) -> Vec<u8> {
|
|
s.chars().map(|c| unicode_to_byte(c)).collect()
|
|
}
|
|
|
|
fn take_valid_utf8(pending: &mut Vec<u8>) -> String {
|
|
match std::str::from_utf8(pending) {
|
|
Ok(text) => {
|
|
let text = text.to_string();
|
|
pending.clear();
|
|
text
|
|
}
|
|
Err(err) => {
|
|
let valid_up_to = err.valid_up_to();
|
|
if valid_up_to == 0 {
|
|
if let Some(error_len) = err.error_len() {
|
|
let invalid_len = error_len.min(pending.len());
|
|
let text = String::from_utf8_lossy(&pending[..invalid_len]).into_owned();
|
|
pending.drain(..invalid_len);
|
|
return text;
|
|
}
|
|
return String::new();
|
|
}
|
|
let text = String::from_utf8_lossy(&pending[..valid_up_to]).into_owned();
|
|
pending.drain(..valid_up_to);
|
|
text
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Convert a Unicode char back to the byte it represents in GPT-2 encoding.
|
|
fn unicode_to_byte(c: char) -> u8 {
|
|
// Build the inverse map on first use
|
|
use std::sync::OnceLock;
|
|
static INV_MAP: OnceLock<HashMap<u32, u8>> = OnceLock::new();
|
|
|
|
let map = INV_MAP.get_or_init(|| {
|
|
let mut m = HashMap::new();
|
|
// Build GPT-2's bytes_to_unicode forward map, then invert
|
|
let mut n = 0u32;
|
|
for b in 0..=255u16 {
|
|
let byte = b as u8;
|
|
let unicode = match byte {
|
|
0x21..=0x7E | 0xA1..=0xAC | 0xAE..=0xFF => byte as u32,
|
|
_ => {
|
|
let u = 256 + n;
|
|
n += 1;
|
|
u
|
|
}
|
|
};
|
|
m.insert(unicode, byte);
|
|
}
|
|
m
|
|
});
|
|
|
|
*map.get(&(c as u32)).unwrap_or_else(|| {
|
|
panic!("unmapped unicode char U+{:04X} in tokenizer", c as u32)
|
|
})
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::{take_valid_utf8, Tokenizer};
|
|
|
|
#[test]
|
|
fn qwen_added_tokens_are_indivisible_and_im_end_is_eos() {
|
|
let path =
|
|
std::env::temp_dir().join(format!("xserv-tokenizer-test-{}.json", std::process::id()));
|
|
std::fs::write(
|
|
&path,
|
|
r#"{
|
|
"model": {
|
|
"vocab": {},
|
|
"merges": [],
|
|
"byte_fallback": false
|
|
},
|
|
"added_tokens": [
|
|
{"id":151643,"content":"<|endoftext|>","special":true},
|
|
{"id":151644,"content":"<|im_start|>","special":true},
|
|
{"id":151645,"content":"<|im_end|>","special":true},
|
|
{"id":151667,"content":"<think>","special":false},
|
|
{"id":151668,"content":"</think>","special":false}
|
|
]
|
|
}"#,
|
|
)
|
|
.unwrap();
|
|
|
|
let tokenizer = Tokenizer::from_file(&path);
|
|
let _ = std::fs::remove_file(&path);
|
|
|
|
assert_eq!(tokenizer.eos_token_id(), Some(151645));
|
|
assert_eq!(tokenizer.encode("<think>"), vec![151667]);
|
|
assert_eq!(tokenizer.encode("</think>"), vec![151668]);
|
|
assert_eq!(tokenizer.decode(&[151645]), "<|im_end|>");
|
|
}
|
|
|
|
#[test]
|
|
fn stream_decode_buffers_incomplete_utf8() {
|
|
let mut pending = vec![0xF0, 0x9F];
|
|
assert_eq!(take_valid_utf8(&mut pending), "");
|
|
pending.extend_from_slice(&[0x98, 0x8A, b'!']);
|
|
assert_eq!(take_valid_utf8(&mut pending), "😊!");
|
|
assert!(pending.is_empty());
|
|
}
|
|
}
|