inferencelayer 0.2.8

//! Tokenizer BPE « HF json » sans dépendance — encode/décode depuis `tokenizer.json`, pour
//! les cibles où le crate `tokenizers` ne compile pas (wasm32 : la page charge le LLM sur
//! WebGPU et doit tokenizer sur place). Couvre les BPE convertis de sentencepiece (Gemma) :
//! normalisation espace→▁, merges par rang, byte-fallback `<0xXX>`, added_tokens découpés
//! AVANT le BPE (plus long match), décodage ▁→espace + fusion des octets.
//!
//! La fidélité est gardée par un test NATIF d'équivalence STRICTE contre `tokenizers`
//! (feature `cli`) sur un corpus français/anglais, tokenizer.json d'E2B.

use std::collections::HashMap;

use anyhow::{Context, Result};

pub struct HfBpe {
    vocab: HashMap<String, u32>,
    id_to_tok: Vec<String>,
    /// (gauche, droite) → (rang, id fusionné).
    ranks: HashMap<(u32, u32), (u32, u32)>,
    /// added_tokens, triés du plus long au plus court (le scan préfère le plus long match).
    specials: Vec<(String, u32)>,
}

impl HfBpe {
    pub fn from_tokenizer_json(bytes: &[u8]) -> Result<Self> {
        let v: serde_json::Value = serde_json::from_slice(bytes)?;
        let model = v.get("model").context("model")?;
        let vocab_obj = model
            .get("vocab")
            .and_then(|x| x.as_object())
            .context("vocab")?;
        let mut vocab = HashMap::with_capacity(vocab_obj.len());
        let mut max_id = 0usize;
        for (tok, id) in vocab_obj {
            let id = id.as_u64().context("id")? as u32;
            max_id = max_id.max(id as usize);
            vocab.insert(tok.clone(), id);
        }
        let mut id_to_tok = vec![String::new(); max_id + 1];
        for (tok, &id) in &vocab {
            id_to_tok[id as usize] = tok.clone();
        }
        let merges = model
            .get("merges")
            .and_then(|x| x.as_array())
            .context("merges")?;
        let mut ranks = HashMap::with_capacity(merges.len());
        for (rank, m) in merges.iter().enumerate() {
            // Deux formats existent dans la nature : "gauche droite" ou ["gauche", "droite"].
            let (l, r) = match m {
                serde_json::Value::String(s) => {
                    let (l, r) = s.split_once(' ').context("merge")?;
                    (l.to_string(), r.to_string())
                }
                serde_json::Value::Array(a) => (
                    a[0].as_str().context("merge l")?.to_string(),
                    a[1].as_str().context("merge r")?.to_string(),
                ),
                _ => anyhow::bail!("format de merge inconnu"),
            };
            let (Some(&li), Some(&ri)) = (vocab.get(&l), vocab.get(&r)) else {
                continue;
            };
            let merged = format!("{l}{r}");
            let Some(&mi) = vocab.get(&merged) else {
                continue;
            };
            ranks.entry((li, ri)).or_insert((rank as u32, mi));
        }
        let mut specials: Vec<(String, u32)> = v
            .get("added_tokens")
            .and_then(|x| x.as_array())
            .map(|a| {
                a.iter()
                    .filter_map(|t| {
                        Some((
                            t.get("content")?.as_str()?.to_string(),
                            t.get("id")?.as_u64()? as u32,
                        ))
                    })
                    .collect()
            })
            .unwrap_or_default();
        specials.sort_by_key(|(s, _)| std::cmp::Reverse(s.len()));
        Ok(Self {
            vocab,
            id_to_tok,
            ranks,
            specials,
        })
    }

    /// BPE d'un segment SANS added_tokens (déjà découpés par [`Self::encode`]).
    fn bpe(&self, text: &str, out: &mut Vec<u32>) {
        // Normalisation sentencepiece : espace → ▁ ; unités de départ = caractères, avec
        // byte-fallback <0xXX> pour un caractère absent du vocab.
        let norm = text.replace(' ', "\u{2581}");
        let mut ids: Vec<u32> = Vec::with_capacity(norm.len());
        for ch in norm.chars() {
            let s = ch.to_string();
            match self.vocab.get(&s) {
                Some(&id) => ids.push(id),
                None => {
                    let mut buf = [0u8; 4];
                    for b in ch.encode_utf8(&mut buf).as_bytes() {
                        ids.push(self.vocab[&format!("<0x{b:02X}>")]);
                    }
                }
            }
        }
        // Boucle de merges : fusionner la paire de PLUS PETIT rang, répéter. O(n²) — les
        // prompts d'une app de dictée font des centaines de tokens, pas des millions.
        loop {
            let mut best: Option<(u32, usize, u32)> = None; // (rang, position, id fusionné)
            for i in 0..ids.len().saturating_sub(1) {
                if let Some(&(rank, merged)) = self.ranks.get(&(ids[i], ids[i + 1]))
                    && best.map_or(true, |(br, _, _)| rank < br)
                {
                    best = Some((rank, i, merged));
                }
            }
            match best {
                Some((_, i, merged)) => {
                    ids[i] = merged;
                    ids.remove(i + 1);
                }
                None => break,
            }
        }
        out.extend_from_slice(&ids);
    }

    /// Encode `text`, added_tokens découpés d'abord (plus long match, occurrences littérales).
    pub fn encode(&self, text: &str) -> Vec<u32> {
        let mut out = Vec::new();
        let mut rest = text;
        'outer: while !rest.is_empty() {
            // Première occurrence du special le plus prioritaire (le plus long d'abord).
            let mut first: Option<(usize, &str, u32)> = None;
            for (s, id) in &self.specials {
                if let Some(pos) = rest.find(s.as_str())
                    && first.map_or(true, |(p, f, _)| pos < p || (pos == p && s.len() > f.len()))
                {
                    first = Some((pos, s, *id));
                }
            }
            match first {
                Some((pos, s, id)) => {
                    if pos > 0 {
                        self.bpe(&rest[..pos], &mut out);
                    }
                    out.push(id);
                    rest = &rest[pos + s.len()..];
                }
                None => {
                    self.bpe(rest, &mut out);
                    break 'outer;
                }
            }
        }
        out
    }

    /// Décode : concatène les tokens, ▁→espace, fusion des `<0xXX>` en octets, specials omis.
    pub fn decode(&self, ids: &[u32]) -> String {
        let mut bytes: Vec<u8> = Vec::new();
        let special: std::collections::HashSet<&str> =
            self.specials.iter().map(|(s, _)| s.as_str()).collect();
        for &id in ids {
            let Some(tok) = self.id_to_tok.get(id as usize) else {
                continue;
            };
            if special.contains(tok.as_str()) {
                continue;
            }
            if tok.len() == 6 && tok.starts_with("<0x") && tok.ends_with('>') {
                if let Ok(b) = u8::from_str_radix(&tok[3..5], 16) {
                    bytes.push(b);
                    continue;
                }
            }
            bytes.extend_from_slice(tok.replace('\u{2581}', " ").as_bytes());
        }
        String::from_utf8_lossy(&bytes).into_owned()
    }
}

#[cfg(all(test, not(target_arch = "wasm32"), feature = "cli"))]
mod tests {
    use super::*;

    /// Équivalence STRICTE avec le crate `tokenizers` sur le tokenizer.json d'E2B — la seule
    /// définition de « correct » qui compte (skip si le checkpoint est absent).
    #[test]
    fn should_match_the_hf_tokenizer_exactly() {
        let dir = std::path::PathBuf::from(std::env::var("HOME").unwrap())
            .join(".cache/inferencelayer/gemma-4-e2b-it");
        if !dir.join("tokenizer.json").exists() {
            eprintln!("SKIP: {} absent", dir.display());
            return;
        }
        let bytes = std::fs::read(dir.join("tokenizer.json")).unwrap();
        let ours = HfBpe::from_tokenizer_json(&bytes).expect("parse");
        let hf = tokenizers::Tokenizer::from_file(dir.join("tokenizer.json")).expect("hf");
        let corpus = [
            "La capitale de la France est",
            "Les audits internes ont été réalisés conformément au plan annuel.",
            "<bos><|turn>user\nCorrige : c'est comme meme chelou<turn|>\n<|turn>model\n",
            "Hello, world! 123 — éàüç œ 😀",
            "  espaces  multiples  et\ntabulations\t.",
        ];
        for text in corpus {
            let a = ours.encode(text);
            let b: Vec<u32> = hf.encode(text, false).expect("hf enc").get_ids().to_vec();
            assert_eq!(a, b, "encode divergent sur {text:?}");
            let d = ours.decode(&a);
            let hd = hf.decode(&b, true).expect("hf dec");
            assert_eq!(d, hd, "decode divergent sur {text:?}");
        }
    }
}