xserv/crates/xserv-tokenizer/src/bpe.rs

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>,
    byte_fallback: bool,
}

#[derive(Deserialize)]
struct TokenizerJson {
    model: ModelSection,
    #[serde(default)]
    added_tokens: Vec<AddedToken>,
}

#[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();
        }
        let eos_token_id = special_tokens
            .get("<|im_end|>")
            .or_else(|| special_tokens.get("<|end_of_text|>"))
            .or_else(|| special_tokens.get("<|endoftext|>"))
            .copied();

        // Pre-tokenization regex
        let pre_tokenize_re = if byte_fallback {
            // Qwen-style: split on whitespace boundaries, keep Unicode words/numbers
            Regex::new(r"[\p{L}\p{N}]+|[^\s\p{L}\p{N}]|\s+").unwrap()
        } else {
            // GPT-2 style
            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,
            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
    }

    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());
    }
}