Skip to main content

native_whisperx/translation/
planning.rs

1//! Deterministic planning for the curated native translation surface.
2
3use std::{fmt, str::FromStr};
4
5use serde::{Deserialize, Serialize};
6
7/// A language supported by the curated native translation registry.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
9pub enum CuratedLanguage {
10    #[serde(rename = "en")]
11    English,
12    #[serde(rename = "de")]
13    German,
14    #[serde(rename = "fr")]
15    French,
16    #[serde(rename = "es")]
17    Spanish,
18    #[serde(rename = "it")]
19    Italian,
20    #[serde(rename = "pt")]
21    Portuguese,
22    #[serde(rename = "nl")]
23    Dutch,
24    #[serde(rename = "pl")]
25    Polish,
26}
27
28impl CuratedLanguage {
29    /// The complete curated language set in stable registry order.
30    pub const ALL: [Self; 8] = [
31        Self::English,
32        Self::German,
33        Self::French,
34        Self::Spanish,
35        Self::Italian,
36        Self::Portuguese,
37        Self::Dutch,
38        Self::Polish,
39    ];
40
41    /// Returns the ISO 639-1 language code used by OPUS-MT plans.
42    pub const fn code(self) -> &'static str {
43        match self {
44            Self::English => "en",
45            Self::German => "de",
46            Self::French => "fr",
47            Self::Spanish => "es",
48            Self::Italian => "it",
49            Self::Portuguese => "pt",
50            Self::Dutch => "nl",
51            Self::Polish => "pl",
52        }
53    }
54}
55
56impl fmt::Display for CuratedLanguage {
57    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
58        formatter.write_str(self.code())
59    }
60}
61
62impl FromStr for CuratedLanguage {
63    type Err = TranslationPlanError;
64
65    fn from_str(value: &str) -> Result<Self, Self::Err> {
66        match value.trim().to_ascii_lowercase().as_str() {
67            "en" => Ok(Self::English),
68            "de" => Ok(Self::German),
69            "fr" => Ok(Self::French),
70            "es" => Ok(Self::Spanish),
71            "it" => Ok(Self::Italian),
72            "pt" => Ok(Self::Portuguese),
73            "nl" => Ok(Self::Dutch),
74            "pl" => Ok(Self::Polish),
75            _ => Err(TranslationPlanError::UnsupportedLanguage {
76                language: value.to_string(),
77            }),
78        }
79    }
80}
81
82/// Whether a plan uses one direct model or composes two legs through English.
83#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
84#[serde(rename_all = "camelCase", tag = "kind")]
85pub enum TranslationPlanProvenance {
86    Direct,
87    PivotTranslation { pivot: CuratedLanguage },
88}
89
90/// One ordered OPUS-MT model invocation in a translation plan.
91#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
92#[serde(rename_all = "camelCase")]
93pub struct TranslationLeg {
94    source: CuratedLanguage,
95    target: CuratedLanguage,
96    model_id: String,
97}
98
99impl TranslationLeg {
100    /// Returns the source language consumed by this leg.
101    pub const fn source(&self) -> CuratedLanguage {
102        self.source
103    }
104
105    /// Returns the target language produced by this leg.
106    pub const fn target(&self) -> CuratedLanguage {
107        self.target
108    }
109
110    /// Returns the canonical Hugging Face OPUS-MT model ID for this leg.
111    pub fn model_id(&self) -> &str {
112        &self.model_id
113    }
114}
115
116/// A deterministic direct or Pivot Translation plan.
117#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
118#[serde(rename_all = "camelCase")]
119pub struct TranslationPlan {
120    source: CuratedLanguage,
121    target: CuratedLanguage,
122    provenance: TranslationPlanProvenance,
123    legs: Vec<TranslationLeg>,
124}
125
126impl TranslationPlan {
127    /// Selects the stable validated plan for a pair of curated languages.
128    pub fn new(
129        source: CuratedLanguage,
130        target: CuratedLanguage,
131    ) -> Result<Self, TranslationPlanError> {
132        if source == target {
133            return Err(TranslationPlanError::SameLanguage { language: source });
134        }
135
136        if let Some(model_id) = validated_direct_model_id(source, target) {
137            return Ok(Self {
138                source,
139                target,
140                provenance: TranslationPlanProvenance::Direct,
141                legs: vec![TranslationLeg {
142                    source,
143                    target,
144                    model_id: model_id.to_string(),
145                }],
146            });
147        }
148
149        let pivot = CuratedLanguage::English;
150        let source_model = validated_direct_model_id(source, pivot)
151            .expect("every curated source language has a validated English translation model");
152        let target_model = validated_direct_model_id(pivot, target)
153            .expect("every curated target language has a validated English translation model");
154        Ok(Self {
155            source,
156            target,
157            provenance: TranslationPlanProvenance::PivotTranslation { pivot },
158            legs: vec![
159                TranslationLeg {
160                    source,
161                    target: pivot,
162                    model_id: source_model.to_string(),
163                },
164                TranslationLeg {
165                    source: pivot,
166                    target,
167                    model_id: target_model.to_string(),
168                },
169            ],
170        })
171    }
172
173    /// Parses ISO 639-1 codes and selects their stable validated plan.
174    pub fn from_language_codes(source: &str, target: &str) -> Result<Self, TranslationPlanError> {
175        Self::new(source.parse()?, target.parse()?)
176    }
177
178    /// Returns the requested source language.
179    pub const fn source(&self) -> CuratedLanguage {
180        self.source
181    }
182
183    /// Returns the requested target language.
184    pub const fn target(&self) -> CuratedLanguage {
185        self.target
186    }
187
188    /// Returns whether the plan is direct or a Pivot Translation.
189    pub const fn provenance(&self) -> TranslationPlanProvenance {
190        self.provenance
191    }
192
193    /// Returns the one or two model legs in execution order.
194    pub fn legs(&self) -> &[TranslationLeg] {
195        &self.legs
196    }
197}
198
199// This conservative registry contains dedicated pair repositories whose Marian
200// files were validated for the native resolver. Portuguese and Polish use
201// publisher-documented multilingual models only for their required English
202// bridge directions. Missing non-English pairs deliberately pivot via English.
203fn validated_direct_model_id(
204    source: CuratedLanguage,
205    target: CuratedLanguage,
206) -> Option<&'static str> {
207    use CuratedLanguage::{Dutch, English, French, German, Italian, Polish, Portuguese, Spanish};
208
209    match (source, target) {
210        (German, English) => Some("Helsinki-NLP/opus-mt-de-en"),
211        (German, Spanish) => Some("Helsinki-NLP/opus-mt-de-es"),
212        (German, French) => Some("Helsinki-NLP/opus-mt-de-fr"),
213        (German, Italian) => Some("Helsinki-NLP/opus-mt-de-it"),
214        (German, Dutch) => Some("Helsinki-NLP/opus-mt-de-nl"),
215        (German, Polish) => Some("Helsinki-NLP/opus-mt-de-pl"),
216        (English, German) => Some("Helsinki-NLP/opus-mt-en-de"),
217        (English, Spanish) => Some("Helsinki-NLP/opus-mt-en-es"),
218        (English, French) => Some("Helsinki-NLP/opus-mt-en-fr"),
219        (English, Italian) => Some("Helsinki-NLP/opus-mt-en-it"),
220        (English, Portuguese) => Some("Helsinki-NLP/opus-mt-tc-big-en-pt"),
221        (English, Dutch) => Some("Helsinki-NLP/opus-mt-en-nl"),
222        (English, Polish) => Some("Helsinki-NLP/opus-mt-en-zlw"),
223        (Spanish, German) => Some("Helsinki-NLP/opus-mt-es-de"),
224        (Spanish, English) => Some("Helsinki-NLP/opus-mt-es-en"),
225        (Spanish, French) => Some("Helsinki-NLP/opus-mt-es-fr"),
226        (Spanish, Italian) => Some("Helsinki-NLP/opus-mt-es-it"),
227        (Spanish, Dutch) => Some("Helsinki-NLP/opus-mt-es-nl"),
228        (Spanish, Polish) => Some("Helsinki-NLP/opus-mt-es-pl"),
229        (French, German) => Some("Helsinki-NLP/opus-mt-fr-de"),
230        (French, English) => Some("Helsinki-NLP/opus-mt-fr-en"),
231        (French, Spanish) => Some("Helsinki-NLP/opus-mt-fr-es"),
232        (French, Polish) => Some("Helsinki-NLP/opus-mt-fr-pl"),
233        (Italian, German) => Some("Helsinki-NLP/opus-mt-it-de"),
234        (Italian, English) => Some("Helsinki-NLP/opus-mt-it-en"),
235        (Italian, Spanish) => Some("Helsinki-NLP/opus-mt-it-es"),
236        (Italian, French) => Some("Helsinki-NLP/opus-mt-it-fr"),
237        (Portuguese, English) => Some("Helsinki-NLP/opus-mt-ROMANCE-en"),
238        (Dutch, English) => Some("Helsinki-NLP/opus-mt-nl-en"),
239        (Dutch, Spanish) => Some("Helsinki-NLP/opus-mt-nl-es"),
240        (Dutch, French) => Some("Helsinki-NLP/opus-mt-nl-fr"),
241        (Polish, German) => Some("Helsinki-NLP/opus-mt-pl-de"),
242        (Polish, English) => Some("Helsinki-NLP/opus-mt-pl-en"),
243        (Polish, Spanish) => Some("Helsinki-NLP/opus-mt-pl-es"),
244        (Polish, French) => Some("Helsinki-NLP/opus-mt-pl-fr"),
245        _ => None,
246    }
247}
248
249/// Validation failures returned before any translation model is resolved.
250#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
251pub enum TranslationPlanError {
252    #[error("unsupported translation language `{language}`")]
253    UnsupportedLanguage { language: String },
254    #[error("source and target translation language are both `{language}`")]
255    SameLanguage { language: CuratedLanguage },
256}