aoiandroid/IndexTTS-Rust

1//! Text tokenization for TTS2//!3//! Uses SentencePiece BPE tokenization for converting text to tokens4 5use crate::{Error, Result};6use std::collections::HashMap;7use std::path::Path;8 9/// Tokenizer configuration10#[derive(Debug, Clone)]11pub struct TokenizerConfig {12    /// Path to BPE model13    pub model_path: String,14    /// Vocabulary size15    pub vocab_size: usize,16    /// Start of text token ID17    pub bos_id: i64,18    /// End of text token ID19    pub eos_id: i64,20    /// Unknown token ID21    pub unk_id: i64,22    /// Padding token ID23    pub pad_id: i64,24}25 26impl Default for TokenizerConfig {27    fn default() -> Self {28        Self {29            model_path: "models/bpe.model".to_string(),30            vocab_size: 6681,31            bos_id: 1,32            eos_id: 2,33            unk_id: 0,34            pad_id: 3,35        }36    }37}38 39/// Text tokenizer using BPE (Byte Pair Encoding)40#[derive(Debug)]41pub struct TextTokenizer {42    /// Configuration43    config: TokenizerConfig,44    /// Token to ID mapping45    token_to_id: HashMap<String, i64>,46    /// ID to token mapping47    id_to_token: HashMap<i64, String>,48    /// Character-level fallback vocabulary49    char_vocab: HashMap<char, i64>,50}51 52impl TextTokenizer {53    /// Create new tokenizer with default vocabulary54    pub fn new(config: TokenizerConfig) -> Result<Self> {55        let mut token_to_id = HashMap::new();56        let mut id_to_token = HashMap::new();57        let mut char_vocab = HashMap::new();58 59        // Add special tokens60        token_to_id.insert("<unk>".to_string(), config.unk_id);61        token_to_id.insert("<s>".to_string(), config.bos_id);62        token_to_id.insert("</s>".to_string(), config.eos_id);63        token_to_id.insert("<pad>".to_string(), config.pad_id);64 65        id_to_token.insert(config.unk_id, "<unk>".to_string());66        id_to_token.insert(config.bos_id, "<s>".to_string());67        id_to_token.insert(config.eos_id, "</s>".to_string());68        id_to_token.insert(config.pad_id, "<pad>".to_string());69 70        // Add basic ASCII characters71        let mut next_id = 4i64;72        for c in ' '..='~' {73            char_vocab.insert(c, next_id);74            token_to_id.insert(c.to_string(), next_id);75            id_to_token.insert(next_id, c.to_string());76            next_id += 1;77        }78 79        // Add Chinese character range (simplified approach)80        // In production, this would load from the actual BPE model81        for code_point in 0x4E00u32..=0x9FFF {82            if let Some(c) = char::from_u32(code_point) {83                char_vocab.insert(c, next_id);84                token_to_id.insert(c.to_string(), next_id);85                id_to_token.insert(next_id, c.to_string());86                next_id += 1;87 88                if next_id >= config.vocab_size as i64 {89                    break;90                }91            }92        }93 94        Ok(Self {95            config,96            token_to_id,97            id_to_token,98            char_vocab,99        })100    }101 102    /// Load tokenizer from model file103    pub fn load<P: AsRef<Path>>(path: P) -> Result<Self> {104        let path = path.as_ref();105        if !path.exists() {106            return Err(Error::FileNotFound(path.display().to_string()));107        }108 109        // In production, this would load the actual SentencePiece model110        // For now, create a character-level tokenizer111        let config = TokenizerConfig {112            model_path: path.display().to_string(),113            ..Default::default()114        };115 116        Self::new(config)117    }118 119    /// Encode text to token IDs120    pub fn encode(&self, text: &str) -> Result<Vec<i64>> {121        let mut tokens = Vec::new();122 123        // Add BOS token124        tokens.push(self.config.bos_id);125 126        // Tokenize character by character (simplified)127        // In production, this would use BPE merging128        for ch in text.chars() {129            if let Some(&id) = self.char_vocab.get(&ch) {130                tokens.push(id);131            } else if let Some(&id) = self.token_to_id.get(&ch.to_string()) {132                tokens.push(id);133            } else {134                // Unknown token135                tokens.push(self.config.unk_id);136            }137        }138 139        // Add EOS token140        tokens.push(self.config.eos_id);141 142        Ok(tokens)143    }144 145    /// Encode text without special tokens146    pub fn encode_without_special(&self, text: &str) -> Result<Vec<i64>> {147        let mut tokens = Vec::new();148 149        for ch in text.chars() {150            if let Some(&id) = self.char_vocab.get(&ch) {151                tokens.push(id);152            } else if let Some(&id) = self.token_to_id.get(&ch.to_string()) {153                tokens.push(id);154            } else {155                tokens.push(self.config.unk_id);156            }157        }158 159        Ok(tokens)160    }161 162    /// Decode token IDs to text163    pub fn decode(&self, tokens: &[i64]) -> Result<String> {164        let mut text = String::new();165 166        for &token_id in tokens {167            // Skip special tokens168            if token_id == self.config.bos_id169                || token_id == self.config.eos_id170                || token_id == self.config.pad_id171            {172                continue;173            }174 175            if let Some(token) = self.id_to_token.get(&token_id) {176                text.push_str(token);177            } else {178                // Unknown token placeholder179                text.push('?');180            }181        }182 183        Ok(text)184    }185 186    /// Get vocabulary size187    pub fn vocab_size(&self) -> usize {188        self.config.vocab_size189    }190 191    /// Get BOS token ID192    pub fn bos_id(&self) -> i64 {193        self.config.bos_id194    }195 196    /// Get EOS token ID197    pub fn eos_id(&self) -> i64 {198        self.config.eos_id199    }200 201    /// Get UNK token ID202    pub fn unk_id(&self) -> i64 {203        self.config.unk_id204    }205 206    /// Get PAD token ID207    pub fn pad_id(&self) -> i64 {208        self.config.pad_id209    }210 211    /// Pad sequences to same length212    pub fn pad_sequences(&self, sequences: &[Vec<i64>], max_len: Option<usize>) -> Vec<Vec<i64>> {213        let max_length = max_len.unwrap_or_else(|| sequences.iter().map(|s| s.len()).max().unwrap_or(0));214 215        sequences216            .iter()217            .map(|seq| {218                let mut padded = seq.clone();219                while padded.len() < max_length {220                    padded.push(self.config.pad_id);221                }222                padded.truncate(max_length);223                padded224            })225            .collect()226    }227 228    /// Create attention mask (1 for real tokens, 0 for padding)229    pub fn create_attention_mask(&self, tokens: &[i64]) -> Vec<i64> {230        tokens231            .iter()232            .map(|&t| if t == self.config.pad_id { 0 } else { 1 })233            .collect()234    }235 236    /// Batch encode multiple texts237    pub fn batch_encode(&self, texts: &[&str]) -> Result<Vec<Vec<i64>>> {238        texts.iter().map(|text| self.encode(text)).collect()239    }240 241    /// Batch encode and pad242    pub fn batch_encode_padded(243        &self,244        texts: &[&str],245        max_len: Option<usize>,246    ) -> Result<Vec<Vec<i64>>> {247        let encoded: Vec<Vec<i64>> = self.batch_encode(texts)?;248        Ok(self.pad_sequences(&encoded, max_len))249    }250}251 252#[cfg(test)]253mod tests {254    use super::*;255 256    #[test]257    fn test_tokenizer_creation() {258        let config = TokenizerConfig::default();259        let tokenizer = TextTokenizer::new(config).unwrap();260        assert!(tokenizer.vocab_size() > 0);261    }262 263    #[test]264    fn test_encode_decode() {265        let config = TokenizerConfig::default();266        let tokenizer = TextTokenizer::new(config).unwrap();267 268        let text = "Hello world";269        let tokens = tokenizer.encode(text).unwrap();270 271        // Should start with BOS and end with EOS272        assert_eq!(tokens[0], tokenizer.bos_id());273        assert_eq!(*tokens.last().unwrap(), tokenizer.eos_id());274 275        let decoded = tokenizer.decode(&tokens).unwrap();276        assert_eq!(decoded, text);277    }278 279    #[test]280    fn test_encode_chinese() {281        let config = TokenizerConfig::default();282        let tokenizer = TextTokenizer::new(config).unwrap();283 284        let text = "你好";285        let tokens = tokenizer.encode(text).unwrap();286 287        // Should have BOS + 2 chars + EOS = 4 tokens288        assert_eq!(tokens.len(), 4);289    }290 291    #[test]292    fn test_pad_sequences() {293        let config = TokenizerConfig::default();294        let tokenizer = TextTokenizer::new(config).unwrap();295 296        let seq1 = vec![1, 2, 3];297        let seq2 = vec![1, 2, 3, 4, 5];298 299        let padded = tokenizer.pad_sequences(&[seq1, seq2], None);300 301        assert_eq!(padded[0].len(), 5);302        assert_eq!(padded[1].len(), 5);303        assert_eq!(padded[0][3], tokenizer.pad_id());304    }305 306    #[test]307    fn test_attention_mask() {308        let config = TokenizerConfig::default();309        let tokenizer = TextTokenizer::new(config).unwrap();310 311        let tokens = vec![1, 2, tokenizer.pad_id(), tokenizer.pad_id()];312        let mask = tokenizer.create_attention_mask(&tokens);313 314        assert_eq!(mask, vec![1, 1, 0, 0]);315    }316}317