aoiandroid/IndexTTS-Rust

1//! Text processing module for IndexTTS2//!3//! Provides text normalization, tokenization, and phoneme conversion.4 5mod normalizer;6mod phoneme;7mod tokenizer;8 9pub use normalizer::{Language, TextNormalizer};10pub use phoneme::{g2p_english, pinyin_to_phones};11pub use tokenizer::{TextTokenizer, TokenizerConfig};12 13use crate::Result;14 15/// Process text through the complete frontend pipeline16pub fn process_text(text: &str, tokenizer: &TextTokenizer) -> Result<Vec<i64>> {17    // Normalize text18    let normalizer = TextNormalizer::new();19    let normalized = normalizer.normalize(text)?;20 21    // Tokenize22    let tokens = tokenizer.encode(&normalized)?;23 24    Ok(tokens)25}26 27/// Detect language of text28pub fn detect_language(text: &str) -> Language {29    let mut chinese_count = 0;30    let mut english_count = 0;31 32    for ch in text.chars() {33        if is_chinese_char(ch) {34            chinese_count += 1;35        } else if ch.is_ascii_alphabetic() {36            english_count += 1;37        }38    }39 40    if chinese_count > 0 && english_count == 0 {41        Language::Chinese42    } else if english_count > 0 && chinese_count == 0 {43        Language::English44    } else if chinese_count > 0 && english_count > 0 {45        Language::Mixed46    } else {47        // Default to English for pure punctuation or empty48        Language::English49    }50}51 52/// Check if character is Chinese53pub fn is_chinese_char(ch: char) -> bool {54    matches!(ch as u32,55        0x4E00..=0x9FFF |     // CJK Unified Ideographs56        0x3400..=0x4DBF |     // CJK Unified Ideographs Extension A57        0x20000..=0x2A6DF |   // CJK Unified Ideographs Extension B58        0x2A700..=0x2B73F |   // CJK Unified Ideographs Extension C59        0x2B740..=0x2B81F |   // CJK Unified Ideographs Extension D60        0xF900..=0xFAFF |     // CJK Compatibility Ideographs61        0x2F800..=0x2FA1F     // CJK Compatibility Ideographs Supplement62    )63}64 65/// Check if text contains Chinese characters66pub fn contains_chinese(text: &str) -> bool {67    text.chars().any(is_chinese_char)68}69 70/// Check if text contains only ASCII71pub fn is_ascii_only(text: &str) -> bool {72    text.is_ascii()73}74 75/// Split text into segments by language76pub fn split_by_language(text: &str) -> Vec<(String, Language)> {77    let mut segments = Vec::new();78    let mut current_segment = String::new();79    let mut current_lang = None;80 81    for ch in text.chars() {82        let char_lang = if is_chinese_char(ch) {83            Some(Language::Chinese)84        } else if ch.is_ascii_alphabetic() {85            Some(Language::English)86        } else {87            None // Punctuation or other88        };89 90        match (current_lang, char_lang) {91            (None, Some(lang)) => {92                current_lang = Some(lang);93                current_segment.push(ch);94            }95            (Some(curr), Some(lang)) if curr == lang => {96                current_segment.push(ch);97            }98            (Some(curr), Some(lang)) if curr != lang => {99                if !current_segment.trim().is_empty() {100                    segments.push((current_segment.clone(), curr));101                }102                current_segment = ch.to_string();103                current_lang = Some(lang);104            }105            (Some(_), None) => {106                // Punctuation - add to current segment107                current_segment.push(ch);108            }109            (None, None) => {110                // Pure punctuation111                if !current_segment.is_empty() {112                    current_segment.push(ch);113                }114            }115            _ => {}116        }117    }118 119    if !current_segment.trim().is_empty() {120        if let Some(lang) = current_lang {121            segments.push((current_segment, lang));122        }123    }124 125    segments126}127 128#[cfg(test)]129mod tests {130    use super::*;131 132    #[test]133    fn test_is_chinese_char() {134        assert!(is_chinese_char('中'));135        assert!(is_chinese_char('文'));136        assert!(!is_chinese_char('a'));137        assert!(!is_chinese_char('1'));138    }139 140    #[test]141    fn test_detect_language() {142        assert_eq!(detect_language("Hello world"), Language::English);143        assert_eq!(detect_language("你好世界"), Language::Chinese);144        assert_eq!(detect_language("Hello 世界"), Language::Mixed);145    }146 147    #[test]148    fn test_contains_chinese() {149        assert!(contains_chinese("Hello 世界"));150        assert!(contains_chinese("你好"));151        assert!(!contains_chinese("Hello world"));152    }153}154