Update translation.py

translation.py

@@ -45,21 +45,25 @@ def load_model(source_lang, target_lang):
     tokenizer_model_pair = all_models.get(model_key)
     if tokenizer_model_pair and tokenizer_model_pair[0] and tokenizer_model_pair[1]:
         return tokenizer_model_pair
-    # Optimized pivot through English using preloaded models
+    # Use direct English pivot only if direct model unavailable
+    if source_lang != "en" and target_lang != "en":
+        en_pivot_pair = all_models.get(("en", target_lang)) or _load_model_pair("en", target_lang)
+        if en_pivot_pair[0] and en_pivot_pair[1]:
+            src_to_en_pair = all_models.get((source_lang, "en")) or _load_model_pair(source_lang, "en")
+            if src_to_en_pair[0] and src_to_en_pair[1]:
+                return en_pivot_pair  # Prefer direct pivot chain
+    default_tokenizer, _ = _load_default_model()
     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()
-            with torch.no_grad():
+        with torch.no_grad():
+            if source_lang != "en":
+                src_to_en_tokenizer, src_to_en_model = all_models.get((source_lang, "en"), _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()
-            with torch.no_grad():
+            else:
+                en_text = text
+            if target_lang != "en":
+                en_to_tgt_tokenizer, en_to_tgt_model = all_models.get(("en", target_lang), _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
-    default_tokenizer, _ = _load_default_model()
+            return en_text
     return default_tokenizer, CombinedModel()
 
 # Cache resource to load default translation model