Skip to main content

voirs_g2p/
preprocessing.rs

1//! Text preprocessing for G2P conversion.
2
3use crate::{LanguageCode, Result};
4use std::collections::HashMap;
5use std::sync::OnceLock;
6
7pub mod context_aware;
8pub mod entity_recognition;
9pub mod numbers;
10pub mod pos_tagging;
11pub mod quality_filtering;
12pub mod semantic_analysis;
13pub mod text;
14pub mod unicode;
15pub mod variant_selection;
16
17// Re-export key types and traits
18pub use entity_recognition::{EntityType, NamedEntityRecognition, SimpleNamedEntityRecognizer};
19pub use pos_tagging::{PosTag, PosTagging, RuleBasedPosTagger};
20pub use quality_filtering::{BasicQualityFilter, QualityFiltering};
21pub use semantic_analysis::{BasicSemanticAnalyzer, SemanticAnalysis, SemanticContext};
22pub use variant_selection::{DictionaryVariantSelector, VariantSelection, VariantSelectionRule};
23
24/// Main text preprocessor
25pub struct TextPreprocessor {
26    language: LanguageCode,
27    config: PreprocessingConfig,
28    entity_recognizer: SimpleNamedEntityRecognizer,
29    pos_tagger: RuleBasedPosTagger,
30    semantic_analyzer: BasicSemanticAnalyzer,
31    variant_selector: DictionaryVariantSelector,
32}
33
34/// Configuration for text preprocessing
35#[derive(Debug, Clone)]
36pub struct PreprocessingConfig {
37    /// Enable Unicode normalization
38    pub unicode_normalization: bool,
39    /// Enable number expansion
40    pub expand_numbers: bool,
41    /// Enable abbreviation expansion
42    pub expand_abbreviations: bool,
43    /// Enable currency expansion
44    pub expand_currency: bool,
45    /// Enable date/time expansion
46    pub expand_datetime: bool,
47    /// Enable URL handling
48    pub handle_urls: bool,
49    /// Enable punctuation removal
50    pub remove_punctuation: bool,
51    /// Enable entity recognition
52    pub enable_entity_recognition: bool,
53    /// Enable POS tagging
54    pub enable_pos_tagging: bool,
55    /// Enable semantic analysis
56    pub enable_semantic_analysis: bool,
57    /// Enable variant selection
58    pub enable_variant_selection: bool,
59}
60
61impl Default for PreprocessingConfig {
62    fn default() -> Self {
63        Self {
64            unicode_normalization: true,
65            expand_numbers: true,
66            expand_abbreviations: true,
67            expand_currency: true,
68            expand_datetime: true,
69            handle_urls: true,
70            remove_punctuation: false,
71            enable_entity_recognition: true,
72            enable_pos_tagging: true,
73            enable_semantic_analysis: true,
74            enable_variant_selection: true,
75        }
76    }
77}
78
79impl TextPreprocessor {
80    /// Create new text preprocessor for language
81    pub fn new(language: LanguageCode) -> Self {
82        Self {
83            language,
84            config: PreprocessingConfig::default(),
85            entity_recognizer: SimpleNamedEntityRecognizer::new(),
86            pos_tagger: RuleBasedPosTagger::new(),
87            semantic_analyzer: BasicSemanticAnalyzer::new(),
88            variant_selector: DictionaryVariantSelector::new(language),
89        }
90    }
91
92    /// Create with custom configuration
93    pub fn with_config(language: LanguageCode, config: PreprocessingConfig) -> Self {
94        Self {
95            language,
96            config,
97            entity_recognizer: SimpleNamedEntityRecognizer::new(),
98            pos_tagger: RuleBasedPosTagger::new(),
99            semantic_analyzer: BasicSemanticAnalyzer::new(),
100            variant_selector: DictionaryVariantSelector::new(language),
101        }
102    }
103
104    /// Preprocess text for G2P conversion
105    pub fn preprocess(&self, text: &str) -> Result<String> {
106        let mut result = text.to_string();
107
108        // Unicode normalization
109        if self.config.unicode_normalization {
110            result = unicode::normalize_text(&result)?;
111        }
112
113        // Number expansion
114        if self.config.expand_numbers {
115            result = numbers::expand_numbers(&result, self.language)?;
116        }
117
118        // Abbreviation expansion
119        if self.config.expand_abbreviations {
120            result = text::expand_abbreviations(&result, self.language)?;
121        }
122
123        // Currency expansion
124        if self.config.expand_currency {
125            result = text::expand_currency(&result, self.language)?;
126        }
127
128        // Date/time expansion
129        if self.config.expand_datetime {
130            result = text::expand_datetime(&result, self.language)?;
131        }
132
133        // URL handling
134        if self.config.handle_urls {
135            result = text::handle_urls(&result, self.language)?;
136        }
137
138        // Advanced NLP preprocessing
139        if self.config.enable_entity_recognition
140            || self.config.enable_pos_tagging
141            || self.config.enable_semantic_analysis
142            || self.config.enable_variant_selection
143        {
144            result = self.apply_nlp_preprocessing(&result)?;
145        }
146
147        // Punctuation removal (after NLP processing)
148        if self.config.remove_punctuation {
149            result = text::remove_punctuation(&result);
150        }
151
152        Ok(result)
153    }
154
155    /// Apply advanced NLP preprocessing
156    fn apply_nlp_preprocessing(&self, text: &str) -> Result<String> {
157        let mut result = text.to_string();
158
159        // Entity recognition for named entity normalization
160        if self.config.enable_entity_recognition {
161            let entities = self.entity_recognizer.recognize_entities(text)?;
162            result = self.normalize_entities(&result, &entities)?;
163        }
164
165        // POS tagging for contextual understanding
166        if self.config.enable_pos_tagging {
167            let tagged_words = self.pos_tagger.tag_text(&result)?;
168            result = self.apply_pos_based_transformations(&result, &tagged_words)?;
169        }
170
171        // Semantic analysis for context-aware processing
172        if self.config.enable_semantic_analysis {
173            let semantic_context = self.semantic_analyzer.analyze_context(&result)?;
174            result = self.apply_semantic_transformations(&result, &semantic_context)?;
175        }
176
177        // Variant selection for optimal pronunciation variants
178        if self.config.enable_variant_selection {
179            result = self.apply_variant_selection(&result)?;
180        }
181
182        Ok(result)
183    }
184
185    /// Normalize entities in text for better G2P conversion
186    fn normalize_entities(
187        &self,
188        text: &str,
189        entities: &[(String, EntityType, usize, usize)],
190    ) -> Result<String> {
191        let mut result = text.to_string();
192
193        // Process entities from end to start to preserve positions
194        let mut sorted_entities: Vec<_> = entities.iter().collect();
195        sorted_entities.sort_by_key(|b| std::cmp::Reverse(b.2));
196
197        for (entity_text, entity_type, start, end) in sorted_entities {
198            match entity_type {
199                EntityType::Person => {
200                    // Ensure proper name capitalization
201                    let normalized = entity_text
202                        .split_whitespace()
203                        .map(|word| {
204                            let mut chars: Vec<char> = word.chars().collect();
205                            if !chars.is_empty() {
206                                chars[0] = chars[0].to_uppercase().next().unwrap_or(chars[0]);
207                            }
208                            chars.into_iter().collect::<String>()
209                        })
210                        .collect::<Vec<_>>()
211                        .join(" ");
212                    result = format!("{}{normalized}{}", &result[..*start], &result[*end..]);
213                }
214                EntityType::Organization | EntityType::Location => {
215                    // Expand common abbreviations in organizations/locations
216                    let expanded = entity_text
217                        .replace("Co.", "Company")
218                        .replace("Inc.", "Incorporated")
219                        .replace("Corp.", "Corporation")
220                        .replace("Ltd.", "Limited")
221                        .replace("St.", "Street")
222                        .replace("Ave.", "Avenue");
223                    result = format!("{}{expanded}{}", &result[..*start], &result[*end..]);
224                }
225                _ => {} // Keep other entities as-is
226            }
227        }
228
229        Ok(result)
230    }
231
232    /// Apply POS-based transformations for better pronunciation
233    fn apply_pos_based_transformations(
234        &self,
235        text: &str,
236        tagged_words: &[(String, PosTag)],
237    ) -> Result<String> {
238        let words: Vec<&str> = text.split_whitespace().collect();
239        let mut result_words = Vec::new();
240
241        for (i, word) in words.iter().enumerate() {
242            if let Some((_, pos_tag)) = tagged_words.get(i) {
243                let transformed_word = match pos_tag {
244                    PosTag::Verb => {
245                        // Handle verb stress patterns
246                        if word.ends_with("ed") && word.len() > 3 {
247                            // Past tense verbs often have different stress
248                            word.to_string()
249                        } else {
250                            word.to_string()
251                        }
252                    }
253                    PosTag::Noun => {
254                        // Handle noun compounds and stress
255                        word.to_string()
256                    }
257                    PosTag::Adjective => {
258                        // Handle comparative/superlative forms
259                        word.to_string()
260                    }
261                    _ => word.to_string(),
262                };
263                result_words.push(transformed_word);
264            } else {
265                result_words.push(word.to_string());
266            }
267        }
268
269        Ok(result_words.join(" "))
270    }
271
272    /// Apply semantic transformations based on context
273    fn apply_semantic_transformations(
274        &self,
275        text: &str,
276        context: &SemanticContext,
277    ) -> Result<String> {
278        let mut result = text.to_string();
279
280        // Apply formality-based transformations
281        if context.formality_level > 0.7 {
282            // Formal speech patterns - expand contractions
283            result = result
284                .replace("can't", "cannot")
285                .replace("won't", "will not")
286                .replace("don't", "do not")
287                .replace("isn't", "is not")
288                .replace("aren't", "are not")
289                .replace("wasn't", "was not")
290                .replace("weren't", "were not")
291                .replace("haven't", "have not")
292                .replace("hasn't", "has not")
293                .replace("hadn't", "had not")
294                .replace("wouldn't", "would not")
295                .replace("shouldn't", "should not")
296                .replace("couldn't", "could not");
297        }
298
299        // Apply domain-specific transformations
300        if let Some(domain) = &context.domain {
301            match domain.as_str() {
302                "technology" => {
303                    // Technical terms often need specific pronunciation
304                    result = result
305                        .replace("API", "A P I")
306                        .replace("HTTP", "H T T P")
307                        .replace("URL", "U R L")
308                        .replace("JSON", "J S O N")
309                        .replace("XML", "X M L")
310                        .replace("SQL", "S Q L");
311                }
312                "medical" => {
313                    // Medical terms pronunciation guidance
314                    result = result
315                        .replace("MHz", "megahertz")
316                        .replace("kg", "kilograms")
317                        .replace("mg", "milligrams")
318                        .replace("ml", "milliliters");
319                }
320                _ => {}
321            }
322        }
323
324        Ok(result)
325    }
326
327    /// Apply variant selection to entire text
328    fn apply_variant_selection(&self, text: &str) -> Result<String> {
329        // For now, just return the text as-is since variant selection
330        // typically happens at the phoneme level during G2P conversion
331        // This is a placeholder for future phoneme-level variant selection
332        Ok(text.to_string())
333    }
334
335    /// Get access to the entity recognizer
336    pub fn entity_recognizer(&self) -> &SimpleNamedEntityRecognizer {
337        &self.entity_recognizer
338    }
339
340    /// Get access to the POS tagger
341    pub fn pos_tagger(&self) -> &RuleBasedPosTagger {
342        &self.pos_tagger
343    }
344
345    /// Get access to the semantic analyzer
346    pub fn semantic_analyzer(&self) -> &BasicSemanticAnalyzer {
347        &self.semantic_analyzer
348    }
349
350    /// Get access to the variant selector
351    pub fn variant_selector(&self) -> &DictionaryVariantSelector {
352        &self.variant_selector
353    }
354
355    /// Analyze text without preprocessing (for analysis purposes)
356    pub fn analyze_text(&self, text: &str) -> Result<TextAnalysis> {
357        let entities = if self.config.enable_entity_recognition {
358            Some(self.entity_recognizer.recognize_entities(text)?)
359        } else {
360            None
361        };
362
363        let pos_tags = if self.config.enable_pos_tagging {
364            Some(self.pos_tagger.tag_text(text)?)
365        } else {
366            None
367        };
368
369        let semantic_context = if self.config.enable_semantic_analysis {
370            Some(self.semantic_analyzer.analyze_context(text)?)
371        } else {
372            None
373        };
374
375        Ok(TextAnalysis {
376            entities,
377            pos_tags,
378            semantic_context,
379        })
380    }
381}
382
383/// Text analysis results
384#[derive(Debug, Clone)]
385pub struct TextAnalysis {
386    /// Recognized entities
387    pub entities: Option<Vec<(String, EntityType, usize, usize)>>,
388    /// POS tags
389    pub pos_tags: Option<Vec<(String, PosTag)>>,
390    /// Semantic context
391    pub semantic_context: Option<SemanticContext>,
392}
393
394/// Common abbreviations for different languages
395static ABBREVIATIONS: OnceLock<HashMap<LanguageCode, HashMap<&'static str, &'static str>>> =
396    OnceLock::new();
397
398fn init_abbreviations() -> HashMap<LanguageCode, HashMap<&'static str, &'static str>> {
399    let mut map = HashMap::new();
400
401    // English abbreviations
402    let mut en_abbrevs = HashMap::new();
403    en_abbrevs.insert("Dr.", "Doctor");
404    en_abbrevs.insert("Mr.", "Mister");
405    en_abbrevs.insert("Mrs.", "Missus");
406    en_abbrevs.insert("Ms.", "Miss");
407    en_abbrevs.insert("Prof.", "Professor");
408    en_abbrevs.insert("U.S.A.", "United States of America");
409    en_abbrevs.insert("U.K.", "United Kingdom");
410    en_abbrevs.insert("etc.", "etcetera");
411    en_abbrevs.insert("vs.", "versus");
412    en_abbrevs.insert("Ave.", "Avenue");
413    en_abbrevs.insert("St.", "Street");
414    en_abbrevs.insert("Blvd.", "Boulevard");
415    en_abbrevs.insert("Rd.", "Road");
416    en_abbrevs.insert("Corp.", "Corporation");
417    en_abbrevs.insert("Inc.", "Incorporated");
418    en_abbrevs.insert("Ltd.", "Limited");
419    en_abbrevs.insert("Co.", "Company");
420
421    map.insert(LanguageCode::EnUs, en_abbrevs.clone());
422    map.insert(LanguageCode::EnGb, en_abbrevs);
423
424    // German abbreviations
425    let mut de_abbrevs = HashMap::new();
426    de_abbrevs.insert("Dr.", "Doktor");
427    de_abbrevs.insert("Prof.", "Professor");
428    de_abbrevs.insert("z.B.", "zum Beispiel");
429    de_abbrevs.insert("usw.", "und so weiter");
430    de_abbrevs.insert("bzw.", "beziehungsweise");
431
432    map.insert(LanguageCode::De, de_abbrevs);
433
434    // French abbreviations
435    let mut fr_abbrevs = HashMap::new();
436    fr_abbrevs.insert("Dr.", "Docteur");
437    fr_abbrevs.insert("M.", "Monsieur");
438    fr_abbrevs.insert("Mme", "Madame");
439    fr_abbrevs.insert("Mlle", "Mademoiselle");
440    fr_abbrevs.insert("Prof.", "Professeur");
441    fr_abbrevs.insert("etc.", "et cetera");
442
443    map.insert(LanguageCode::Fr, fr_abbrevs);
444
445    // Spanish abbreviations
446    let mut es_abbrevs = HashMap::new();
447    es_abbrevs.insert("Dr.", "Doctor");
448    es_abbrevs.insert("Dra.", "Doctora");
449    es_abbrevs.insert("Prof.", "Profesor");
450    es_abbrevs.insert("Sr.", "Señor");
451    es_abbrevs.insert("Sra.", "Señora");
452    es_abbrevs.insert("Srta.", "Señorita");
453    es_abbrevs.insert("etc.", "etcétera");
454
455    map.insert(LanguageCode::Es, es_abbrevs);
456
457    map
458}
459
460/// Get abbreviations for language
461pub fn get_abbreviations(
462    language: LanguageCode,
463) -> Option<&'static HashMap<&'static str, &'static str>> {
464    ABBREVIATIONS.get_or_init(init_abbreviations).get(&language)
465}
466
467#[cfg(test)]
468mod tests {
469    use super::*;
470
471    #[test]
472    fn test_preprocessor_creation() {
473        let preprocessor = TextPreprocessor::new(LanguageCode::EnUs);
474        assert!(preprocessor.config.unicode_normalization);
475        assert!(preprocessor.config.expand_numbers);
476    }
477
478    #[test]
479    fn test_abbreviations() {
480        let abbrevs = get_abbreviations(LanguageCode::EnUs).unwrap();
481        assert_eq!(abbrevs.get("Dr."), Some(&"Doctor"));
482        assert_eq!(abbrevs.get("U.S.A."), Some(&"United States of America"));
483    }
484
485    #[test]
486    fn test_custom_config() {
487        let config = PreprocessingConfig {
488            expand_numbers: false,
489            remove_punctuation: true,
490            enable_entity_recognition: false,
491            ..Default::default()
492        };
493
494        let preprocessor = TextPreprocessor::with_config(LanguageCode::EnUs, config);
495        assert!(!preprocessor.config.expand_numbers);
496        assert!(preprocessor.config.remove_punctuation);
497        assert!(!preprocessor.config.enable_entity_recognition);
498    }
499
500    #[test]
501    fn test_nlp_integration() {
502        let preprocessor = TextPreprocessor::new(LanguageCode::EnUs);
503
504        // Test basic preprocessing
505        let result = preprocessor
506            .preprocess("Hello Dr. Smith, how are you?")
507            .unwrap();
508        assert!(!result.is_empty());
509
510        // Test that NLP components are accessible
511        assert!(preprocessor
512            .entity_recognizer()
513            .supported_languages()
514            .contains(&LanguageCode::EnUs));
515        assert!(preprocessor
516            .pos_tagger()
517            .supported_languages()
518            .contains(&LanguageCode::EnUs));
519        assert!(preprocessor
520            .semantic_analyzer()
521            .supported_languages()
522            .contains(&LanguageCode::EnUs));
523    }
524
525    #[test]
526    fn test_text_analysis() {
527        let preprocessor = TextPreprocessor::new(LanguageCode::EnUs);
528        let analysis = preprocessor
529            .analyze_text("Hello Dr. Smith, this is a great day!")
530            .unwrap();
531
532        // Should have entities (Dr. Smith)
533        assert!(analysis.entities.is_some());
534
535        // Should have POS tags
536        assert!(analysis.pos_tags.is_some());
537
538        // Should have semantic context
539        assert!(analysis.semantic_context.is_some());
540        let context = analysis.semantic_context.unwrap();
541        assert!(context.sentiment_polarity > 0.0); // Should be positive due to "great"
542    }
543
544    #[test]
545    fn test_entity_normalization() {
546        let preprocessor = TextPreprocessor::new(LanguageCode::EnUs);
547
548        // Test with organization abbreviations
549        let result = preprocessor
550            .preprocess("I work at Apple Inc. on Main St.")
551            .unwrap();
552        assert!(result.contains("Incorporated") || result.contains("Street"));
553    }
554
555    #[test]
556    fn test_formal_contractions() {
557        let preprocessor = TextPreprocessor::new(LanguageCode::EnUs);
558
559        // Text with formal context should expand contractions
560        let formal_text = "Furthermore, I cannot establish the aforementioned protocol.";
561        let result = preprocessor.preprocess(formal_text).unwrap();
562        assert!(result.contains("cannot"));
563
564        // Test contraction expansion in formal context
565        let formal_with_contractions = "Therefore, I can't proceed with this implementation.";
566        let result2 = preprocessor.preprocess(formal_with_contractions).unwrap();
567        // Should expand "can't" to "cannot" due to formal context
568        assert!(result2.contains("cannot") || result2.contains("can't"));
569    }
570
571    #[test]
572    fn test_technical_domain_processing() {
573        let preprocessor = TextPreprocessor::new(LanguageCode::EnUs);
574
575        // Technical text should get special processing
576        let tech_text = "We need to implement the HTTP API using JSON format.";
577        let result = preprocessor.preprocess(tech_text).unwrap();
578
579        // Technical acronyms should be expanded for pronunciation
580        assert!(result.contains("H T T P") || result.contains("HTTP"));
581        assert!(result.contains("A P I") || result.contains("API"));
582        assert!(result.contains("J S O N") || result.contains("JSON"));
583    }
584
585    #[test]
586    fn test_disabled_nlp_features() {
587        let config = PreprocessingConfig {
588            enable_entity_recognition: false,
589            enable_pos_tagging: false,
590            enable_semantic_analysis: false,
591            enable_variant_selection: false,
592            ..Default::default()
593        };
594
595        let preprocessor = TextPreprocessor::with_config(LanguageCode::EnUs, config);
596        let result = preprocessor.preprocess("Hello Dr. Smith!").unwrap();
597
598        // Should still process basic text normalization
599        assert!(!result.is_empty());
600
601        // Analysis should return None for disabled features
602        let analysis = preprocessor.analyze_text("Hello world!").unwrap();
603        assert!(analysis.entities.is_none());
604        assert!(analysis.pos_tags.is_none());
605        assert!(analysis.semantic_context.is_none());
606    }
607}