1use std::{fmt, str::FromStr};
4
5use serde::{Deserialize, Serialize};
6
7#[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 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 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#[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#[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 pub const fn source(&self) -> CuratedLanguage {
102 self.source
103 }
104
105 pub const fn target(&self) -> CuratedLanguage {
107 self.target
108 }
109
110 pub fn model_id(&self) -> &str {
112 &self.model_id
113 }
114}
115
116#[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 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 pub fn from_language_codes(source: &str, target: &str) -> Result<Self, TranslationPlanError> {
175 Self::new(source.parse()?, target.parse()?)
176 }
177
178 pub const fn source(&self) -> CuratedLanguage {
180 self.source
181 }
182
183 pub const fn target(&self) -> CuratedLanguage {
185 self.target
186 }
187
188 pub const fn provenance(&self) -> TranslationPlanProvenance {
190 self.provenance
191 }
192
193 pub fn legs(&self) -> &[TranslationLeg] {
195 &self.legs
196 }
197}
198
199fn 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#[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}