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,
    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}"));

        let byte_fallback = tj.model.byte_fallback;

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

        // Special tokens
        let mut special_tokens = HashMap::new();
        let mut special_token_ids = HashMap::new();
        let mut eos_token_id = None;
        for at in &tj.added_tokens {
            if at.special {
                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();
                if at.content == "<|endoftext|>" || at.content == "<|end_of_text|>" {
                    eos_token_id = Some(at.id);
                }
            }
        }

        // 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().map(|&b| {
                *self.encoder.get(&vec![b]).unwrap_or_else(|| {
                    panic!("byte {b} (0x{b:02X}) not in vocab")
                })
            }).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 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()
}

/// 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)
    })
}