aoiandroid/IndexTTS-Rust

1//! Text normalization for TTS2 3use crate::Result;4use lazy_static::lazy_static;5use regex::Regex;6use std::collections::HashMap;7 8#[derive(Debug, Clone, Copy, PartialEq, Eq)]9pub enum Language {10    Chinese,11    English,12    Mixed,13}14 15#[derive(Debug)]16pub struct TextNormalizer {17    punct_map: HashMap<char, char>,18    number_words: HashMap<u64, &'static str>,19}20 21lazy_static! {22    static ref NUMBER_REGEX: Regex = Regex::new(r"\d+").unwrap();23    static ref WHITESPACE_REGEX: Regex = Regex::new(r"\s+").unwrap();24}25 26impl TextNormalizer {27    pub fn new() -> Self {28        let mut punct_map = HashMap::new();29        punct_map.insert('\u{FF0C}', ',');30        punct_map.insert('\u{3002}', '.');31        punct_map.insert('\u{FF01}', '!');32        punct_map.insert('\u{FF1F}', '?');33        punct_map.insert('\u{FF1B}', ';');34        punct_map.insert('\u{FF1A}', ':');35        punct_map.insert('\u{201C}', '\u{0022}');36        punct_map.insert('\u{201D}', '\u{0022}');37        punct_map.insert('\u{2018}', '\'');38        punct_map.insert('\u{2019}', '\'');39 40        let mut number_words = HashMap::new();41        number_words.insert(0, "zero");42        number_words.insert(1, "one");43        number_words.insert(2, "two");44        number_words.insert(3, "three");45        number_words.insert(4, "four");46        number_words.insert(5, "five");47        number_words.insert(6, "six");48        number_words.insert(7, "seven");49        number_words.insert(8, "eight");50        number_words.insert(9, "nine");51        number_words.insert(10, "ten");52        number_words.insert(20, "twenty");53        number_words.insert(30, "thirty");54 55        Self { punct_map, number_words }56    }57 58    pub fn normalize(&self, text: &str) -> Result<String> {59        let mut result = self.normalize_punctuation(text);60        result = self.normalize_whitespace(&result);61        Ok(result)62    }63 64    pub fn normalize_punctuation(&self, text: &str) -> String {65        text.chars()66            .map(|c| *self.punct_map.get(&c).unwrap_or(&c))67            .collect()68    }69 70    pub fn normalize_whitespace(&self, text: &str) -> String {71        WHITESPACE_REGEX.replace_all(text, " ").trim().to_string()72    }73 74    pub fn split_sentences(&self, text: &str) -> Vec<String> {75        let mut sentences = Vec::new();76        let mut current = String::new();77 78        for ch in text.chars() {79            current.push(ch);80            if ch == '.' || ch == '!' || ch == '?' {81                let trimmed = current.trim().to_string();82                if !trimmed.is_empty() {83                    sentences.push(trimmed);84                }85                current.clear();86            }87        }88 89        let trimmed = current.trim().to_string();90        if !trimmed.is_empty() {91            sentences.push(trimmed);92        }93 94        sentences95    }96}97 98impl Default for TextNormalizer {99    fn default() -> Self {100        Self::new()101    }102}103 104#[cfg(test)]105mod tests {106    use super::*;107 108    #[test]109    fn test_normalizer() {110        let n = TextNormalizer::new();111        let r = n.normalize_whitespace("  a  b  ");112        assert_eq!(r.len(), 3);113    }114}115