Update translation.py

translation.py

@@ -8,6 +8,51 @@ LANGUAGES = {
     "ar": ("العربية", "Arabic"), "ru": ("Русский", "Russian"), "ja": ("日本語", "Japanese")
 }
 
+@st.cache_resource
+def _load_model_pair(source_lang, target_lang):
+    try:
+        model_name = f"Helsinki-NLP/opus-mt-{source_lang}-{target_lang}"
+        tokenizer = MarianTokenizer.from_pretrained(model_name)
+        model = MarianMTModel.from_pretrained(model_name)
+        return tokenizer, model
+    except Exception:
+        return None, None
+
+@st.cache_resource
+def _load_all_models():
+    models = {}
+    for src in LANGUAGES.keys():
+        for tgt in LANGUAGES.keys():
+            if src != tgt:
+                models[(src, tgt)] = _load_model_pair(src, tgt)
+    return models
+
+all_models = _load_all_models()
+
+def load_model(source_lang, target_lang):
+    if source_lang == target_lang:
+        return _load_default_model()
+    model_key = (source_lang, target_lang)
+    if all_models.get(model_key) and all_models[model_key][0] and all_models[model_key][1]:
+        return all_models[model_key]
+    # Pivot through English
+    def combined_translate(text):
+        en_tokenizer, en_model = all_models.get(("en", "en"), _load_default_model())
+        if source_lang != "en":
+            src_to_en_tokenizer, src_to_en_model = all_models.get((source_lang, "en"), _load_model_pair(source_lang, "en")) or _load_default_model()
+            en_text = src_to_en_tokenizer.decode(src_to_en_model.generate(**src_to_en_tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=500))[0], skip_special_tokens=True)
+        else:
+            en_text = text
+        if target_lang != "en":
+            en_to_tgt_tokenizer, en_to_tgt_model = all_models.get(("en", target_lang), _load_model_pair("en", target_lang)) or _load_default_model()
+            return en_to_tgt_tokenizer.decode(en_to_tgt_model.generate(**en_to_tgt_tokenizer(en_text, return_tensors="pt", padding=True, truncation=True, max_length=500))[0], skip_special_tokens=True)
+        return en_text
+    class CombinedModel:
+        def generate(self, **kwargs):
+            return torch.tensor([combined_translate(tokenizer.decode(x, skip_special_tokens=True)) for x in kwargs['input_ids']])
+    tokenizer, _ = _load_default_model()
+    return tokenizer, CombinedModel()
+
 @st.cache_resource
 def _load_default_model():
     model_name = "Helsinki-NLP/opus-mt-en-hi"
@@ -15,33 +60,6 @@ def _load_default_model():
     model = MarianMTModel.from_pretrained(model_name)
     return tokenizer, model
 
-@st.cache_resource
-def load_model(source_lang, target_lang):
-    try:
-        if source_lang == target_lang:
-            return _load_default_model()
-        # Try direct model
-        model_name = f"Helsinki-NLP/opus-mt-{source_lang}-{target_lang}"
-        try:
-            tokenizer = MarianTokenizer.from_pretrained(model_name)
-            model = MarianMTModel.from_pretrained(model_name)
-            return tokenizer, model
-        except Exception:
-            # Pivot through English for non-English pairs
-            if source_lang != "en" and target_lang != "en":
-                def combined_translate(text):
-                    en_tokenizer, en_model = load_model(source_lang, "en")
-                    en_text = en_tokenizer.decode(en_model.generate(**en_tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=500))[0], skip_special_tokens=True)
-                    target_tokenizer, target_model = load_model("en", target_lang)
-                    return target_tokenizer.decode(target_model.generate(**target_tokenizer(en_text, return_tensors="pt", padding=True, truncation=True, max_length=500))[0], skip_special_tokens=True)
-                class CombinedModel:
-                    def generate(self, **kwargs):
-                        return torch.tensor([combined_translate(tokenizer.decode(x, skip_special_tokens=True)) for x in kwargs['input_ids']])
-                return MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-hi"), CombinedModel()
-            return _load_default_model()
-    except Exception:
-        return _load_default_model()
-
 def translate(text, source_lang, target_lang):
     if not text:
         return ""