rustytorch/crates/production/rtx-wasm-inference/src/tokenizer.rs

//! WASM Tokenizer
//!
//! Tokenization for WebAssembly inference.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use wasm_bindgen::prelude::*;

/// WASM-compatible tokenizer
#[wasm_bindgen]
#[derive(Debug, Clone)]
pub struct WasmTokenizer {
    /// Token to ID mapping
    vocab: HashMap<String, u32>,
    /// ID to token mapping
    id_to_token: HashMap<u32, String>,
    /// Special tokens
    special_tokens: SpecialTokens,
    /// Vocabulary size
    vocab_size: usize,
}

/// Special token IDs
#[derive(Debug, Clone, Serialize, Deserialize)]
struct SpecialTokens {
    bos_token_id: u32,
    eos_token_id: u32,
    pad_token_id: u32,
    unk_token_id: u32,
}

impl Default for SpecialTokens {
    fn default() -> Self {
        Self {
            bos_token_id: 1,
            eos_token_id: 2,
            pad_token_id: 0,
            unk_token_id: 3,
        }
    }
}

#[wasm_bindgen]
impl WasmTokenizer {
    /// Create a new tokenizer with basic vocabulary
    #[wasm_bindgen(constructor)]
    pub fn new() -> Self {
        let mut vocab = HashMap::new();
        let mut id_to_token = HashMap::new();

        // Add special tokens
        vocab.insert("<pad>".to_string(), 0);
        vocab.insert("<s>".to_string(), 1);
        vocab.insert("</s>".to_string(), 2);
        vocab.insert("<unk>".to_string(), 3);

        id_to_token.insert(0, "<pad>".to_string());
        id_to_token.insert(1, "<s>".to_string());
        id_to_token.insert(2, "</s>".to_string());
        id_to_token.insert(3, "<unk>".to_string());

        // Add basic character vocabulary
        for (i, c) in
            "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 .,!?'\"-:;()[]{}"
                .chars()
                .enumerate()
        {
            let id = (i + 4) as u32;
            vocab.insert(c.to_string(), id);
            id_to_token.insert(id, c.to_string());
        }

        let vocab_size = vocab.len();

        Self {
            vocab,
            id_to_token,
            special_tokens: SpecialTokens::default(),
            vocab_size,
        }
    }

    /// Load tokenizer from JSON
    pub fn from_json(json: &str) -> Result<Self, JsError> {
        // Parse tokenizer JSON (simplified)
        let parsed: serde_json::Value = serde_json::from_str(json)
            .map_err(|e| JsError::new(&format!("Failed to parse tokenizer JSON: {}", e)))?;

        // Extract vocabulary
        let mut vocab = HashMap::new();
        let mut id_to_token = HashMap::new();

        if let Some(model) = parsed.get("model") {
            if let Some(vocab_obj) = model.get("vocab") {
                if let Some(vocab_map) = vocab_obj.as_object() {
                    for (token, id) in vocab_map {
                        if let Some(id_num) = id.as_u64() {
                            vocab.insert(token.clone(), id_num as u32);
                            id_to_token.insert(id_num as u32, token.clone());
                        }
                    }
                }
            }
        }

        // If vocabulary is empty, create basic one
        if vocab.is_empty() {
            let tokenizer = Self::new();
            return Ok(tokenizer);
        }

        // Extract special tokens
        let special_tokens = SpecialTokens {
            bos_token_id: parsed
                .get("bos_token_id")
                .and_then(serde_json::Value::as_u64)
                .unwrap_or(1) as u32,
            eos_token_id: parsed
                .get("eos_token_id")
                .and_then(serde_json::Value::as_u64)
                .unwrap_or(2) as u32,
            pad_token_id: parsed
                .get("pad_token_id")
                .and_then(serde_json::Value::as_u64)
                .unwrap_or(0) as u32,
            unk_token_id: parsed
                .get("unk_token_id")
                .and_then(serde_json::Value::as_u64)
                .unwrap_or(3) as u32,
        };

        let vocab_size = vocab.len();

        Ok(Self {
            vocab,
            id_to_token,
            special_tokens,
            vocab_size,
        })
    }

    /// Encode text to token IDs
    pub fn encode(&self, text: &str) -> Result<Vec<u32>, JsError> {
        let mut tokens = Vec::new();

        // Add BOS token
        tokens.push(self.special_tokens.bos_token_id);

        // Simple character-level tokenization
        for c in text.chars() {
            let token_id = self
                .vocab
                .get(&c.to_string())
                .copied()
                .unwrap_or(self.special_tokens.unk_token_id);
            tokens.push(token_id);
        }

        Ok(tokens)
    }

    /// Decode token IDs to text
    pub fn decode(&self, ids: &[u32]) -> Result<String, JsError> {
        let mut text = String::new();

        for &id in ids {
            // Skip special tokens in output
            if id == self.special_tokens.bos_token_id
                || id == self.special_tokens.eos_token_id
                || id == self.special_tokens.pad_token_id
            {
                continue;
            }

            if let Some(token) = self.id_to_token.get(&id) {
                text.push_str(token);
            } else {
                text.push_str("<unk>");
            }
        }

        Ok(text)
    }

    /// Get BOS token ID
    #[wasm_bindgen]
    pub fn bos_token_id(&self) -> u32 {
        self.special_tokens.bos_token_id
    }

    /// Get EOS token ID
    #[wasm_bindgen]
    pub fn eos_token_id(&self) -> u32 {
        self.special_tokens.eos_token_id
    }

    /// Get PAD token ID
    #[wasm_bindgen]
    pub fn pad_token_id(&self) -> u32 {
        self.special_tokens.pad_token_id
    }

    /// Get UNK token ID
    #[wasm_bindgen]
    pub fn unk_token_id(&self) -> u32 {
        self.special_tokens.unk_token_id
    }

    /// Get vocabulary size
    #[wasm_bindgen]
    pub fn vocab_size(&self) -> usize {
        self.vocab_size
    }

    /// Check if token ID is special
    #[wasm_bindgen]
    pub fn is_special_token(&self, id: u32) -> bool {
        id == self.special_tokens.bos_token_id
            || id == self.special_tokens.eos_token_id
            || id == self.special_tokens.pad_token_id
            || id == self.special_tokens.unk_token_id
    }

    /// Get token for ID
    #[wasm_bindgen]
    pub fn id_to_token(&self, id: u32) -> Option<String> {
        self.id_to_token.get(&id).cloned()
    }

    /// Get ID for token
    #[wasm_bindgen]
    pub fn token_to_id(&self, token: &str) -> Option<u32> {
        self.vocab.get(token).copied()
    }

    /// Memory usage in bytes
    #[wasm_bindgen]
    pub fn memory_usage(&self) -> usize {
        // Rough estimate: each entry is ~50 bytes average
        self.vocab_size * 50
    }
}

impl Default for WasmTokenizer {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn test_tokenizer_basic() {
        let tokenizer = WasmTokenizer::new();
        assert!(tokenizer.vocab_size() > 0);
    }

    #[test]
    fn test_encode_decode() {
        let tokenizer = WasmTokenizer::new();

        let text = "Hello";
        let ids = tokenizer.encode(text).unwrap();

        assert!(!ids.is_empty());
        assert_eq!(ids[0], tokenizer.bos_token_id());

        let decoded = tokenizer.decode(&ids[1..]).unwrap();
        assert_eq!(decoded, text);
    }

    #[test]
    fn test_special_tokens() {
        let tokenizer = WasmTokenizer::new();

        assert!(tokenizer.is_special_token(tokenizer.bos_token_id()));
        assert!(tokenizer.is_special_token(tokenizer.eos_token_id()));
        assert!(!tokenizer.is_special_token(100));
    }
}