rustytorch/crates/models/rtx-tts/src/text/mod.rs

//! Text processing for TTS
//!
//! This module provides comprehensive text processing capabilities for TTS including:
//! - Text normalization (case, punctuation, abbreviations, numbers)
//! - Grapheme-to-phoneme conversion
//! - Prosody prediction (pauses, emphasis, intonation)

pub mod normalizer;
pub mod phonemizer;
pub mod prosody;

// Re-export commonly used types
pub use normalizer::{TextNormalizer, TextNormalizerConfig};
pub use phonemizer::{Phonemizer, PhonemizerConfig, PhonemizerBackend, Phoneme};
pub use prosody::{Prosody, ProsodyConfig, ProsodyMarker, ProsodyAnnotation};

use crate::error::Result;

/// Common trait for text processors
pub trait TextProcessor {
    /// Process input text
    ///
    /// # Errors
    ///
    /// Returns error if text processing fails
    fn process(&self, text: &str) -> Result<String>;
}

/// Text processing pipeline
///
/// Combines normalization, phonemization, and prosody prediction
#[derive(Debug)]
pub struct TextPipeline {
    normalizer: TextNormalizer,
    phonemizer: Phonemizer,
    prosody: Prosody,
}

impl TextPipeline {
    /// Create a new text processing pipeline
    #[must_use]
    pub fn new(
        normalizer: TextNormalizer,
        phonemizer: Phonemizer,
        prosody: Prosody,
    ) -> Self {
        Self {
            normalizer,
            phonemizer,
            prosody,
        }
    }

    /// Create a pipeline with default configuration
    #[must_use]
    pub fn default() -> Self {
        Self {
            normalizer: TextNormalizer::default(),
            phonemizer: Phonemizer::default(),
            prosody: Prosody::default(),
        }
    }

    /// Process text through the full pipeline
    ///
    /// # Errors
    ///
    /// Returns error if any stage of processing fails
    pub fn process_full(&self, text: &str) -> Result<ProcessedText> {
        // 1. Normalize text
        let normalized = self.normalizer.normalize(text)?;

        // 2. Convert to phonemes
        let phonemes = self.phonemizer.g2p(&normalized)?;

        // 3. Predict prosody
        let prosody_annotations = self.prosody.predict(&normalized)?;

        Ok(ProcessedText {
            original: text.to_string(),
            normalized,
            phonemes,
            prosody_annotations,
        })
    }

    /// Get reference to the normalizer
    #[must_use]
    pub fn normalizer(&self) -> &TextNormalizer {
        &self.normalizer
    }

    /// Get reference to the phonemizer
    #[must_use]
    pub fn phonemizer(&self) -> &Phonemizer {
        &self.phonemizer
    }

    /// Get reference to the prosody predictor
    #[must_use]
    pub fn prosody(&self) -> &Prosody {
        &self.prosody
    }
}

/// Processed text output containing all stages of processing
#[derive(Debug, Clone)]
pub struct ProcessedText {
    /// Original input text
    pub original: String,
    /// Normalized text
    pub normalized: String,
    /// Phoneme sequences for each word
    pub phonemes: Vec<Vec<Phoneme>>,
    /// Prosody annotations
    pub prosody_annotations: Vec<ProsodyAnnotation>,
}

impl ProcessedText {
    /// Get the phoneme string representation
    #[must_use]
    pub fn phoneme_string(&self, include_stress: bool) -> String {
        self.phonemes
            .iter()
            .map(|word_phonemes| {
                word_phonemes
                    .iter()
                    .map(|p| {
                        if include_stress {
                            p.to_string_with_stress()
                        } else {
                            p.symbol.clone()
                        }
                    })
                    .collect::<Vec<_>>()
                    .join(" ")
            })
            .collect::<Vec<_>>()
            .join(" | ")
    }
}

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

    #[test]
    fn test_pipeline_creation() {
        let pipeline = TextPipeline::default();
        assert!(pipeline.normalizer().config().lowercase);
    }

    #[test]
    fn test_pipeline_full_processing() {
        let pipeline = TextPipeline::default();
        let result = pipeline.process_full("Hello, world! I have 3 apples.").unwrap();

        assert!(!result.original.is_empty());
        assert!(!result.normalized.is_empty());
        assert!(!result.phonemes.is_empty());
        assert!(!result.prosody_annotations.is_empty());
    }

    #[test]
    fn test_processed_text_phoneme_string() {
        let pipeline = TextPipeline::default();
        let result = pipeline.process_full("hello").unwrap();

        let phoneme_str_with_stress = result.phoneme_string(true);
        let phoneme_str_without_stress = result.phoneme_string(false);

        assert!(!phoneme_str_with_stress.is_empty());
        assert!(!phoneme_str_without_stress.is_empty());
    }

    #[test]
    fn test_pipeline_accessors() {
        let pipeline = TextPipeline::default();

        let normalizer = pipeline.normalizer();
        assert!(normalizer.config().lowercase);

        let phonemizer = pipeline.phonemizer();
        assert_eq!(phonemizer.config().language, "en-us");

        let prosody = pipeline.prosody();
        assert_eq!(prosody.config().default_speaking_rate, 1.0);
    }

    #[test]
    fn test_pipeline_with_numbers() {
        let pipeline = TextPipeline::default();
        let result = pipeline.process_full("I have 42 apples.").unwrap();

        assert!(result.normalized.contains("forty two"));
        assert!(!result.phonemes.is_empty());
    }

    #[test]
    fn test_pipeline_with_abbreviations() {
        let pipeline = TextPipeline::default();
        let result = pipeline.process_full("Dr. Smith is here.").unwrap();

        assert!(result.normalized.contains("doctor"));
        assert!(!result.phonemes.is_empty());
    }

    #[test]
    fn test_pipeline_with_questions() {
        let pipeline = TextPipeline::default();
        let result = pipeline.process_full("How are you?").unwrap();

        // Should detect rising intonation for questions
        let has_rising = result.prosody_annotations.iter().any(|a| {
            a.markers.iter().any(|m| matches!(m, ProsodyMarker::Rising))
        });
        assert!(has_rising);
    }

    #[test]
    fn test_pipeline_preserves_original() {
        let pipeline = TextPipeline::default();
        let original = "Hello, World!";
        let result = pipeline.process_full(original).unwrap();

        assert_eq!(result.original, original);
    }

    #[test]
    fn test_pipeline_empty_text() {
        let pipeline = TextPipeline::default();
        let result = pipeline.process_full("").unwrap();

        assert_eq!(result.original, "");
        assert!(result.normalized.is_empty());
        assert!(result.phonemes.is_empty());
        assert!(result.prosody_annotations.is_empty());
    }
}