voirs_g2p/
models.rs

1//! G2P model definitions, training, and loading.
2
3use crate::{G2pError, LanguageCode, Phoneme, Result};
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use std::path::{Path, PathBuf};
7use std::time::{Duration, SystemTime};
8
9/// G2P model types
10#[derive(Debug, Clone, Serialize, Deserialize)]
11pub enum ModelType {
12    /// Rule-based model
13    RuleBased,
14    /// Statistical model
15    Statistical,
16    /// Neural network model
17    Neural,
18    /// Hybrid model combining multiple approaches
19    Hybrid,
20}
21
22/// Model architecture configuration
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct ModelConfig {
25    /// Model type
26    pub model_type: ModelType,
27    /// Model architecture parameters
28    pub architecture: ArchitectureConfig,
29    /// Training configuration
30    pub training: TrainingConfig,
31    /// Model metadata
32    pub metadata: ModelMetadata,
33}
34
35/// Neural network architecture configuration
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct ArchitectureConfig {
38    /// Input vocabulary size
39    pub vocab_size: usize,
40    /// Hidden layer dimensions
41    pub hidden_dims: Vec<usize>,
42    /// Number of layers
43    pub num_layers: usize,
44    /// Dropout rate
45    pub dropout: f32,
46    /// Attention mechanism enabled
47    pub use_attention: bool,
48    /// Bidirectional processing
49    pub bidirectional: bool,
50    /// Activation function
51    pub activation: String,
52}
53
54/// Training configuration
55#[derive(Debug, Clone, Serialize, Deserialize)]
56pub struct TrainingConfig {
57    /// Learning rate
58    pub learning_rate: f32,
59    /// Batch size
60    pub batch_size: usize,
61    /// Number of epochs
62    pub epochs: usize,
63    /// Validation split ratio
64    pub validation_split: f32,
65    /// Early stopping patience
66    pub early_stopping_patience: usize,
67    /// Optimizer type
68    pub optimizer: String,
69    /// Learning rate schedule
70    pub lr_schedule: Option<LearningRateSchedule>,
71    /// Regularization settings
72    pub regularization: RegularizationConfig,
73}
74
75/// Learning rate schedule configuration
76#[derive(Debug, Clone, Serialize, Deserialize)]
77pub struct LearningRateSchedule {
78    /// Schedule type (step, exponential, cosine)
79    pub schedule_type: String,
80    /// Decay rate
81    pub decay_rate: f32,
82    /// Decay steps
83    pub decay_steps: usize,
84    /// Minimum learning rate
85    pub min_lr: f32,
86}
87
88/// Regularization configuration
89#[derive(Debug, Clone, Serialize, Deserialize)]
90pub struct RegularizationConfig {
91    /// L1 regularization strength
92    pub l1: f32,
93    /// L2 regularization strength
94    pub l2: f32,
95    /// Dropout rate
96    pub dropout: f32,
97    /// Gradient clipping threshold
98    pub gradient_clip: Option<f32>,
99}
100
101/// Model metadata
102#[derive(Debug, Clone, Serialize, Deserialize)]
103pub struct ModelMetadata {
104    /// Model name
105    pub name: String,
106    /// Model version
107    pub version: String,
108    /// Description
109    pub description: String,
110    /// Target language
111    pub language: LanguageCode,
112    /// Creation timestamp
113    pub created_at: SystemTime,
114    /// Training duration
115    pub training_duration: Option<Duration>,
116    /// Training dataset info
117    pub dataset_info: Option<DatasetInfo>,
118    /// Model performance metrics
119    pub performance_metrics: HashMap<String, f32>,
120    /// Model size in bytes
121    pub model_size: Option<u64>,
122}
123
124/// Dataset information
125#[derive(Debug, Clone, Serialize, Deserialize)]
126pub struct DatasetInfo {
127    /// Dataset name
128    pub name: String,
129    /// Number of training examples
130    pub train_size: usize,
131    /// Number of validation examples
132    pub validation_size: usize,
133    /// Number of test examples
134    pub test_size: Option<usize>,
135    /// Dataset source
136    pub source: String,
137    /// Dataset version
138    pub version: String,
139}
140
141/// Training progress tracking
142#[derive(Debug, Clone, Serialize, Deserialize)]
143pub struct TrainingProgress {
144    /// Current epoch
145    pub epoch: usize,
146    /// Total epochs
147    pub total_epochs: usize,
148    /// Current step in epoch
149    pub step: usize,
150    /// Total steps in epoch
151    pub total_steps: usize,
152    /// Training loss
153    pub train_loss: f32,
154    /// Validation loss
155    pub val_loss: Option<f32>,
156    /// Training accuracy
157    pub train_accuracy: f32,
158    /// Validation accuracy
159    pub val_accuracy: Option<f32>,
160    /// Learning rate
161    pub learning_rate: f32,
162    /// Elapsed time
163    pub elapsed_time: Duration,
164    /// Estimated time remaining
165    pub eta: Option<Duration>,
166}
167
168/// Model evaluation metrics
169#[derive(Debug, Clone, Serialize, Deserialize)]
170pub struct EvaluationMetrics {
171    /// Phoneme-level accuracy
172    pub phoneme_accuracy: f32,
173    /// Word-level accuracy
174    pub word_accuracy: f32,
175    /// Edit distance (Levenshtein)
176    pub edit_distance: f32,
177    /// BLEU score
178    pub bleu_score: Option<f32>,
179    /// Perplexity
180    pub perplexity: Option<f32>,
181    /// Confidence score distribution
182    pub confidence_stats: ConfidenceStats,
183    /// Per-phoneme accuracy breakdown
184    pub phoneme_breakdown: HashMap<String, f32>,
185}
186
187/// Confidence score statistics
188#[derive(Debug, Clone, Serialize, Deserialize)]
189pub struct ConfidenceStats {
190    /// Mean confidence score
191    pub mean: f32,
192    /// Standard deviation
193    pub std_dev: f32,
194    /// Minimum confidence
195    pub min: f32,
196    /// Maximum confidence
197    pub max: f32,
198    /// Median confidence
199    pub median: f32,
200}
201
202/// G2P model with training capabilities
203#[derive(Debug, Clone, Serialize, Deserialize)]
204pub struct G2pModel {
205    /// Model configuration
206    pub config: ModelConfig,
207    /// Model weights and parameters
208    pub parameters: ModelParameters,
209    /// Training history
210    pub training_history: Vec<TrainingProgress>,
211    /// Evaluation metrics
212    pub evaluation_metrics: Option<EvaluationMetrics>,
213    /// Model file path
214    pub model_path: Option<PathBuf>,
215}
216
217/// Model parameters storage
218#[derive(Debug, Clone, Serialize, Deserialize)]
219pub struct ModelParameters {
220    /// Model weights as byte array
221    pub weights: Vec<u8>,
222    /// Vocabulary mapping
223    pub vocabulary: HashMap<String, usize>,
224    /// Phoneme mapping
225    pub phoneme_mapping: HashMap<String, usize>,
226    /// Additional parameters
227    pub additional_params: HashMap<String, Vec<u8>>,
228}
229
230/// Training dataset representation
231#[derive(Debug, Clone, Serialize, Deserialize)]
232pub struct TrainingDataset {
233    /// Training examples
234    pub examples: Vec<TrainingExample>,
235    /// Dataset metadata
236    pub metadata: DatasetInfo,
237    /// Language code
238    pub language: LanguageCode,
239}
240
241/// Individual training example
242#[derive(Debug, Clone, Serialize, Deserialize)]
243pub struct TrainingExample {
244    /// Input text
245    pub text: String,
246    /// Target phonemes
247    pub phonemes: Vec<Phoneme>,
248    /// Additional context
249    pub context: Option<String>,
250    /// Example weight for training
251    pub weight: f32,
252}
253
254/// Transfer learning configuration
255#[derive(Debug, Clone, Serialize, Deserialize)]
256pub struct TransferLearningConfig {
257    /// Source model path
258    pub source_model_path: PathBuf,
259    /// Layers to freeze during transfer
260    pub freeze_layers: Vec<usize>,
261    /// Fine-tuning learning rate
262    pub fine_tune_lr: f32,
263    /// Target language for transfer
264    pub target_language: LanguageCode,
265    /// Adaptation strategy
266    pub adaptation_strategy: AdaptationStrategy,
267}
268
269/// Model adaptation strategies
270#[derive(Debug, Clone, Serialize, Deserialize)]
271pub enum AdaptationStrategy {
272    /// Full fine-tuning
273    FullFineTuning,
274    /// Feature extraction (freeze base, train head)
275    FeatureExtraction,
276    /// Gradual unfreezing
277    GradualUnfreezing,
278    /// Domain adaptation
279    DomainAdaptation,
280}
281
282/// Few-shot learning configuration
283#[derive(Debug, Clone, Serialize, Deserialize)]
284pub struct FewShotConfig {
285    /// Number of examples per phoneme
286    pub examples_per_phoneme: usize,
287    /// Meta-learning algorithm
288    pub meta_learning_algorithm: String,
289    /// Support set size
290    pub support_set_size: usize,
291    /// Query set size
292    pub query_set_size: usize,
293    /// Number of gradient steps
294    pub gradient_steps: usize,
295    /// Inner learning rate
296    pub inner_lr: f32,
297    /// Outer learning rate
298    pub outer_lr: f32,
299}
300
301/// Pronunciation customization settings
302#[derive(Debug, Clone, Serialize, Deserialize, Default)]
303pub struct PronunciationCustomization {
304    /// Custom pronunciation dictionary
305    pub custom_dict: HashMap<String, Vec<Phoneme>>,
306    /// Phoneme substitution rules
307    pub substitution_rules: Vec<SubstitutionRule>,
308    /// Regional accent modifications
309    pub accent_modifications: HashMap<String, AccentModification>,
310    /// Context-sensitive rules
311    pub context_rules: Vec<ContextRule>,
312}
313
314/// Phoneme substitution rule
315#[derive(Debug, Clone, Serialize, Deserialize)]
316pub struct SubstitutionRule {
317    /// Source phoneme pattern
318    pub source_pattern: String,
319    /// Target phoneme pattern
320    pub target_pattern: String,
321    /// Context conditions
322    pub context_conditions: Vec<String>,
323    /// Rule priority
324    pub priority: usize,
325}
326
327/// Regional accent modification
328#[derive(Debug, Clone, Serialize, Deserialize)]
329pub struct AccentModification {
330    /// Accent name
331    pub accent_name: String,
332    /// Phoneme transformations
333    pub transformations: Vec<SubstitutionRule>,
334    /// Accent strength (0.0-1.0)
335    pub strength: f32,
336}
337
338/// Context-sensitive pronunciation rule
339#[derive(Debug, Clone, Serialize, Deserialize)]
340pub struct ContextRule {
341    /// Word pattern
342    pub word_pattern: String,
343    /// Preceding context
344    pub preceding_context: Option<String>,
345    /// Following context
346    pub following_context: Option<String>,
347    /// Target pronunciation
348    pub target_pronunciation: Vec<Phoneme>,
349    /// Rule confidence
350    pub confidence: f32,
351}
352
353impl G2pModel {
354    /// Create new G2P model with configuration
355    pub fn new(config: ModelConfig) -> Self {
356        Self {
357            config,
358            parameters: ModelParameters {
359                weights: Vec::new(),
360                vocabulary: HashMap::new(),
361                phoneme_mapping: HashMap::new(),
362                additional_params: HashMap::new(),
363            },
364            training_history: Vec::new(),
365            evaluation_metrics: None,
366            model_path: None,
367        }
368    }
369
370    /// Load model from file
371    pub fn load_from_file<P: AsRef<Path>>(path: P) -> Result<Self> {
372        let file_content = std::fs::read(path.as_ref())
373            .map_err(|e| G2pError::ModelError(format!("Failed to read model file: {e}")))?;
374
375        // bincode 2: use serde integration module
376        let (model, _len): (G2pModel, usize) =
377            oxicode::serde::decode_from_slice(&file_content, oxicode::config::standard())
378                .map_err(|e| G2pError::ModelError(format!("Failed to deserialize model: {e}")))?;
379
380        Ok(model)
381    }
382
383    /// Save model to file
384    pub fn save_to_file<P: AsRef<Path>>(&mut self, path: P) -> Result<()> {
385        self.model_path = Some(path.as_ref().to_path_buf());
386
387        let serialized = oxicode::serde::encode_to_vec(self, oxicode::config::standard())
388            .map_err(|e| G2pError::ModelError(format!("Failed to serialize model: {e}")))?;
389
390        std::fs::write(path.as_ref(), serialized)
391            .map_err(|e| G2pError::ModelError(format!("Failed to write model file: {e}")))?;
392
393        Ok(())
394    }
395
396    /// Get model size in bytes
397    pub fn model_size(&self) -> usize {
398        self.parameters.weights.len()
399            + self.parameters.vocabulary.len() * 32 // rough estimate
400            + self.parameters.phoneme_mapping.len() * 32
401            + self.parameters.additional_params.values().map(|v| v.len()).sum::<usize>()
402    }
403
404    /// Add training progress entry
405    pub fn add_training_progress(&mut self, progress: TrainingProgress) {
406        self.training_history.push(progress);
407    }
408
409    /// Get latest training progress
410    pub fn latest_training_progress(&self) -> Option<&TrainingProgress> {
411        self.training_history.last()
412    }
413
414    /// Set evaluation metrics
415    pub fn set_evaluation_metrics(&mut self, metrics: EvaluationMetrics) {
416        self.evaluation_metrics = Some(metrics);
417    }
418
419    /// Check if model is trained
420    pub fn is_trained(&self) -> bool {
421        !self.parameters.weights.is_empty() && !self.training_history.is_empty()
422    }
423
424    /// Get supported language
425    pub fn language(&self) -> LanguageCode {
426        self.config.metadata.language
427    }
428
429    /// Get model performance summary
430    pub fn performance_summary(&self) -> HashMap<String, f32> {
431        let mut summary = HashMap::new();
432
433        if let Some(metrics) = &self.evaluation_metrics {
434            summary.insert("phoneme_accuracy".to_string(), metrics.phoneme_accuracy);
435            summary.insert("word_accuracy".to_string(), metrics.word_accuracy);
436            summary.insert("edit_distance".to_string(), metrics.edit_distance);
437
438            if let Some(bleu) = metrics.bleu_score {
439                summary.insert("bleu_score".to_string(), bleu);
440            }
441
442            if let Some(perplexity) = metrics.perplexity {
443                summary.insert("perplexity".to_string(), perplexity);
444            }
445        }
446
447        summary
448    }
449}
450
451#[cfg(test)]
452mod bincode_migration_tests {
453    use super::*;
454    use std::time::SystemTime;
455
456    #[test]
457    fn serialize_deserialize_roundtrip() {
458        let config = ModelConfig {
459            model_type: ModelType::Neural,
460            architecture: ArchitectureConfig {
461                vocab_size: 100,
462                hidden_dims: vec![64, 64],
463                num_layers: 2,
464                dropout: 0.1,
465                use_attention: true,
466                bidirectional: false,
467                activation: "relu".into(),
468            },
469            training: TrainingConfig {
470                learning_rate: 1e-3,
471                batch_size: 4,
472                epochs: 1,
473                validation_split: 0.1,
474                early_stopping_patience: 2,
475                optimizer: "adam".into(),
476                lr_schedule: None,
477                regularization: RegularizationConfig {
478                    l1: 0.0,
479                    l2: 0.0,
480                    dropout: 0.1,
481                    gradient_clip: Some(1.0),
482                },
483            },
484            metadata: ModelMetadata {
485                name: "test-model".into(),
486                version: "0.1.0".into(),
487                description: "test".into(),
488                language: LanguageCode::EnUs,
489                created_at: SystemTime::now(),
490                training_duration: None,
491                dataset_info: None,
492                performance_metrics: HashMap::new(),
493                model_size: None,
494            },
495        };
496        let mut model = G2pModel::new(config);
497        model.parameters.weights = vec![1, 2, 3];
498        model.parameters.vocabulary.insert("hello".into(), 42);
499
500        let bytes =
501            oxicode::serde::encode_to_vec(&model, oxicode::config::standard()).expect("encode");
502        assert!(!bytes.is_empty());
503
504        let (decoded, _len): (G2pModel, usize) =
505            oxicode::serde::decode_from_slice(&bytes, oxicode::config::standard()).expect("decode");
506        assert_eq!(decoded.parameters.weights.len(), 3);
507        assert_eq!(decoded.parameters.vocabulary.get("hello"), Some(&42));
508    }
509
510    #[test]
511    fn deserialize_failure_on_truncated_bytes() {
512        let config = ModelConfig {
513            model_type: ModelType::Neural,
514            architecture: ArchitectureConfig {
515                vocab_size: 10,
516                hidden_dims: vec![8],
517                num_layers: 1,
518                dropout: 0.0,
519                use_attention: false,
520                bidirectional: false,
521                activation: "relu".into(),
522            },
523            training: TrainingConfig {
524                learning_rate: 1e-3,
525                batch_size: 2,
526                epochs: 1,
527                validation_split: 0.0,
528                early_stopping_patience: 1,
529                optimizer: "adam".into(),
530                lr_schedule: None,
531                regularization: RegularizationConfig {
532                    l1: 0.0,
533                    l2: 0.0,
534                    dropout: 0.0,
535                    gradient_clip: None,
536                },
537            },
538            metadata: ModelMetadata {
539                name: "truncated".into(),
540                version: "0.0.1".into(),
541                description: "truncate test".into(),
542                language: LanguageCode::Ja,
543                created_at: SystemTime::now(),
544                training_duration: None,
545                dataset_info: None,
546                performance_metrics: HashMap::new(),
547                model_size: None,
548            },
549        };
550        let mut model = G2pModel::new(config);
551        model.parameters.weights = vec![7, 8];
552
553        let bytes = oxicode::serde::encode_to_vec(&model, oxicode::config::standard()).unwrap();
554        let truncated = &bytes[0..bytes.len() / 3];
555        let result = oxicode::serde::decode_from_slice::<G2pModel, _>(
556            truncated,
557            oxicode::config::standard(),
558        );
559        assert!(result.is_err(), "Expected error on truncated input");
560    }
561}
562
563impl TrainingDataset {
564    /// Create new training dataset
565    pub fn new(language: LanguageCode, metadata: DatasetInfo) -> Self {
566        Self {
567            examples: Vec::new(),
568            metadata,
569            language,
570        }
571    }
572
573    /// Add training example
574    pub fn add_example(&mut self, example: TrainingExample) {
575        self.examples.push(example);
576    }
577
578    /// Load dataset from file
579    pub fn load_from_file<P: AsRef<Path>>(path: P) -> Result<Self> {
580        let file_content = std::fs::read_to_string(path.as_ref())
581            .map_err(|e| G2pError::ModelError(format!("Failed to read dataset file: {e}")))?;
582
583        let dataset: TrainingDataset = serde_json::from_str(&file_content)
584            .map_err(|e| G2pError::ModelError(format!("Failed to parse dataset: {e}")))?;
585
586        Ok(dataset)
587    }
588
589    /// Save dataset to file
590    pub fn save_to_file<P: AsRef<Path>>(&self, path: P) -> Result<()> {
591        let serialized = serde_json::to_string_pretty(self)
592            .map_err(|e| G2pError::ModelError(format!("Failed to serialize dataset: {e}")))?;
593
594        std::fs::write(path.as_ref(), serialized)
595            .map_err(|e| G2pError::ModelError(format!("Failed to write dataset file: {e}")))?;
596
597        Ok(())
598    }
599
600    /// Split dataset into train/validation/test sets
601    pub fn split(
602        &self,
603        train_ratio: f32,
604        val_ratio: f32,
605    ) -> (TrainingDataset, TrainingDataset, TrainingDataset) {
606        let total_size = self.examples.len();
607        let train_size = (total_size as f32 * train_ratio) as usize;
608        let val_size = (total_size as f32 * val_ratio) as usize;
609
610        let mut train_examples = Vec::new();
611        let mut val_examples = Vec::new();
612        let mut test_examples = Vec::new();
613
614        for (i, example) in self.examples.iter().enumerate() {
615            if i < train_size {
616                train_examples.push(example.clone());
617            } else if i < train_size + val_size {
618                val_examples.push(example.clone());
619            } else {
620                test_examples.push(example.clone());
621            }
622        }
623
624        let train_dataset = TrainingDataset {
625            examples: train_examples,
626            metadata: DatasetInfo {
627                name: format!("{}_train", self.metadata.name),
628                train_size,
629                validation_size: 0,
630                test_size: None,
631                source: self.metadata.source.clone(),
632                version: self.metadata.version.clone(),
633            },
634            language: self.language,
635        };
636
637        let val_dataset = TrainingDataset {
638            examples: val_examples,
639            metadata: DatasetInfo {
640                name: format!("{}_val", self.metadata.name),
641                train_size: 0,
642                validation_size: val_size,
643                test_size: None,
644                source: self.metadata.source.clone(),
645                version: self.metadata.version.clone(),
646            },
647            language: self.language,
648        };
649
650        let test_dataset = TrainingDataset {
651            examples: test_examples,
652            metadata: DatasetInfo {
653                name: format!("{}_test", self.metadata.name),
654                train_size: 0,
655                validation_size: 0,
656                test_size: Some(total_size - train_size - val_size),
657                source: self.metadata.source.clone(),
658                version: self.metadata.version.clone(),
659            },
660            language: self.language,
661        };
662
663        (train_dataset, val_dataset, test_dataset)
664    }
665
666    /// Get dataset statistics
667    pub fn statistics(&self) -> DatasetStatistics {
668        let total_examples = self.examples.len();
669        let total_phonemes: usize = self.examples.iter().map(|e| e.phonemes.len()).sum();
670        let avg_phonemes_per_example = if total_examples > 0 {
671            total_phonemes as f32 / total_examples as f32
672        } else {
673            0.0
674        };
675
676        let mut phoneme_counts = HashMap::new();
677        for example in &self.examples {
678            for phoneme in &example.phonemes {
679                *phoneme_counts.entry(phoneme.symbol.clone()).or_insert(0) += 1;
680            }
681        }
682
683        DatasetStatistics {
684            total_examples,
685            total_phonemes,
686            avg_phonemes_per_example,
687            phoneme_distribution: phoneme_counts,
688            language: self.language,
689        }
690    }
691}
692
693/// Dataset statistics
694#[derive(Debug, Clone, Serialize, Deserialize)]
695pub struct DatasetStatistics {
696    /// Total number of examples
697    pub total_examples: usize,
698    /// Total number of phonemes
699    pub total_phonemes: usize,
700    /// Average phonemes per example
701    pub avg_phonemes_per_example: f32,
702    /// Phoneme distribution
703    pub phoneme_distribution: HashMap<String, usize>,
704    /// Dataset language
705    pub language: LanguageCode,
706}
707
708/// Default implementations
709impl Default for ArchitectureConfig {
710    fn default() -> Self {
711        Self {
712            vocab_size: 10000,
713            hidden_dims: vec![256, 128],
714            num_layers: 2,
715            dropout: 0.1,
716            use_attention: true,
717            bidirectional: true,
718            activation: "relu".to_string(),
719        }
720    }
721}
722
723impl Default for TrainingConfig {
724    fn default() -> Self {
725        Self {
726            learning_rate: 0.001,
727            batch_size: 32,
728            epochs: 100,
729            validation_split: 0.2,
730            early_stopping_patience: 10,
731            optimizer: "adam".to_string(),
732            lr_schedule: None,
733            regularization: RegularizationConfig::default(),
734        }
735    }
736}
737
738impl Default for RegularizationConfig {
739    fn default() -> Self {
740        Self {
741            l1: 0.0,
742            l2: 0.0001,
743            dropout: 0.1,
744            gradient_clip: Some(1.0),
745        }
746    }
747}
748
749impl Default for FewShotConfig {
750    fn default() -> Self {
751        Self {
752            examples_per_phoneme: 5,
753            meta_learning_algorithm: "maml".to_string(),
754            support_set_size: 10,
755            query_set_size: 5,
756            gradient_steps: 1,
757            inner_lr: 0.01,
758            outer_lr: 0.001,
759        }
760    }
761}
voirs_g2p/models.rs

voirs_g2p/
models.rs
