inferencelayer 0.2.4

//! `token id → raw bytes` table, so the FSM can be walked one TOKEN at a time (the constrained
//! decoder masks logits over token ids, then advances the grammar by the chosen token's bytes).
//!
//! Byte-level BPE (GPT-2 / Qwen / LFM2) is the primary path: it reuses the engine's existing
//! [`crate::vocab::build_vocab_blob`] (the same wasm-friendly, `onig`-free decoder the fleet ships to
//! donors) and re-parses its documented blob into `id → bytes`. Tokenizers whose vocab is not
//! byte-level fall back to an SPM/Unigram piece builder (`▁` U+2581 → space, `<0xNN>` byte-fallback
//! tokens → the raw byte). If a tokenizer is neither, construction REFUSES loudly rather than
//! silently producing a wrong byte table — a wrong table would mis-mask, defeating the whole point.

use anyhow::{Result, anyhow};
use serde_json::Value;

/// Maps every token id in a vocabulary to the raw bytes it contributes to the decoded output.
/// Special/control tokens map to empty (they contribute no output bytes).
#[derive(Clone, Debug)]
pub struct TokenByteTable {
    pieces: Vec<Vec<u8>>,
}

impl TokenByteTable {
    /// Build from a `tokenizer.json`, choosing the byte-level path when possible and the SPM
    /// fallback otherwise; refuses loudly if neither applies.
    pub fn from_tokenizer_json(tokenizer_json: &[u8]) -> Result<Self> {
        if let Some(blob) = crate::vocab::build_vocab_blob(tokenizer_json) {
            return Self::from_vocab_blob(&blob);
        }
        Self::from_spm_tokenizer_json(tokenizer_json)
    }

    /// Parse the engine's `id→bytes` vocab blob (`u32 count`, then per id `u32 len` + bytes) into a
    /// byte table. This is the byte-level BPE path.
    pub fn from_vocab_blob(blob: &[u8]) -> Result<Self> {
        let count = u32::from_le_bytes(
            blob.get(0..4)
                .ok_or_else(|| anyhow!("vocab blob truncated (no count)"))?
                .try_into()
                .unwrap(),
        ) as usize;
        let mut pieces = Vec::with_capacity(count);
        let mut o = 4;
        for id in 0..count {
            let len = u32::from_le_bytes(
                blob.get(o..o + 4)
                    .ok_or_else(|| anyhow!("vocab blob truncated at id {id} length"))?
                    .try_into()
                    .unwrap(),
            ) as usize;
            o += 4;
            let bytes = blob
                .get(o..o + len)
                .ok_or_else(|| anyhow!("vocab blob truncated at id {id} bytes"))?;
            pieces.push(bytes.to_vec());
            o += len;
        }
        Ok(Self { pieces })
    }

    /// SPM/Unigram fallback: `model.vocab` is an array of `[piece, score]` (index = id). Pieces are
    /// literal text with `▁` (U+2581) denoting a leading space and `<0xNN>` denoting a raw byte.
    fn from_spm_tokenizer_json(tokenizer_json: &[u8]) -> Result<Self> {
        let v: Value = serde_json::from_slice(tokenizer_json)
            .map_err(|e| anyhow!("tokenizer.json is not valid JSON: {e}"))?;
        let vocab = v.get("model").and_then(|m| m.get("vocab")).ok_or_else(|| {
            anyhow!("tokenizer is neither byte-level BPE nor SPM/Unigram: no `model.vocab`")
        })?;
        let arr = vocab.as_array().ok_or_else(|| {
            anyhow!(
                "tokenizer is neither byte-level BPE nor SPM/Unigram: `model.vocab` is not a \
                 Unigram piece array"
            )
        })?;
        let mut pieces = vec![Vec::new(); arr.len()];
        for (id, entry) in arr.iter().enumerate() {
            let piece = entry
                .get(0)
                .and_then(|p| p.as_str())
                .ok_or_else(|| anyhow!("Unigram vocab entry {id} is malformed"))?;
            pieces[id] = spm_piece_to_bytes(piece);
        }
        // Special tokens render as empty; explicit non-special added tokens keep their literal text.
        if let Some(added) = v.get("added_tokens").and_then(|a| a.as_array()) {
            for tok in added {
                let Some(id) = tok.get("id").and_then(|i| i.as_u64()) else {
                    continue;
                };
                let Some(slot) = pieces.get_mut(id as usize) else {
                    continue;
                };
                let special = tok.get("special").and_then(|s| s.as_bool()).unwrap_or(true);
                *slot = if special {
                    Vec::new()
                } else {
                    tok.get("content")
                        .and_then(|c| c.as_str())
                        .map(|s| s.as_bytes().to_vec())
                        .unwrap_or_default()
                };
            }
        }
        Ok(Self { pieces })
    }

    /// The raw bytes for a token id, or `None` if the id is out of range.
    pub fn bytes(&self, id: u32) -> Option<&[u8]> {
        self.pieces.get(id as usize).map(Vec::as_slice)
    }

    /// The number of token ids the table covers.
    pub fn len(&self) -> usize {
        self.pieces.len()
    }

    pub fn is_empty(&self) -> bool {
        self.pieces.is_empty()
    }
}

/// Convert one SPM piece to its output bytes: `<0xNN>` → the raw byte, otherwise the literal text
/// with every `▁` (U+2581) replaced by a space.
fn spm_piece_to_bytes(piece: &str) -> Vec<u8> {
    if let Some(hex) = piece.strip_prefix("<0x").and_then(|s| s.strip_suffix('>'))
        && hex.len() == 2
        && let Ok(byte) = u8::from_str_radix(hex, 16)
    {
        return vec![byte];
    }
    piece.replace('\u{2581}', " ").into_bytes()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn should_build_byte_table_from_byte_level_blob() {
        // "ĠParis" → " Paris" in the GPT-2 byte alphabet (Ġ = U+0120 = byte 0x20).
        let tok = r#"{ "model": { "vocab": { "ĠParis": 0, "Hello": 1 } } }"#;
        let table = TokenByteTable::from_tokenizer_json(tok.as_bytes()).expect("byte table");
        assert_eq!(table.bytes(1), Some(&b"Hello"[..]));
        assert_eq!(table.bytes(0), Some(&b" Paris"[..]));
    }

    #[test]
    fn should_fall_back_to_spm_piece_builder() {
        // Unigram vocab is an ARRAY of [piece, score]; build_vocab_blob returns None on it.
        let tok = r#"{
            "model": { "type": "Unigram", "vocab": [ ["<unk>", 0.0], ["▁", -1.0], ["▁hi", -2.0], ["<0x0A>", -3.0] ] }
        }"#;
        let table = TokenByteTable::from_tokenizer_json(tok.as_bytes()).expect("spm table");
        assert_eq!(table.bytes(1), Some(&b" "[..]), "\u{2581} → space");
        assert_eq!(table.bytes(2), Some(&b" hi"[..]), "\u{2581}hi → ' hi'");
        assert_eq!(
            table.bytes(3),
            Some(&[0x0A][..]),
            "<0x0A> → raw newline byte"
        );
    }

    #[test]
    fn should_refuse_unknown_tokenizer_loudly() {
        // No byte-level object vocab and no Unigram piece array ⇒ refuse rather than mis-mask.
        let tok = r#"{ "model": { "type": "WordPiece" } }"#;
        let err = TokenByteTable::from_tokenizer_json(tok.as_bytes())
            .expect_err("must refuse an unrecognized tokenizer");
        assert!(
            err.to_string().contains("neither byte-level BPE nor SPM"),
            "refusal must explain why: {err}"
        );
    }
}