Update translation.py

translation.py

@@ -20,18 +20,24 @@ def load_model(source_lang, target_lang):
     try:
         if source_lang == target_lang:
             return _load_default_model()
-        # Try direct model first
+        # 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
+            # Pivot through English
             if source_lang != "en" and target_lang != "en":
-                en_to_target = load_model("en", target_lang)
-                source_to_en = load_model(source_lang, "en")
-                return source_to_en if source_lang == "en" else en_to_target
+                en_to_target_tokenizer, en_to_target_model = load_model("en", target_lang)
+                source_to_en_tokenizer, source_to_en_model = load_model(source_lang, "en")
+                def combined_translate(text):
+                    en_text = source_to_en_tokenizer.decode(source_to_en_model.generate(**source_to_en_tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=500))[0], skip_special_tokens=True)
+                    return en_to_target_tokenizer.decode(en_to_target_model.generate(**en_to_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()
@@ -45,5 +51,6 @@ def translate(text, source_lang, target_lang):
         with torch.no_grad():
             translated = model.generate(**inputs, max_length=500, num_beams=2, early_stopping=True)
         return tokenizer.decode(translated[0], skip_special_tokens=True)
-    except Exception:
+    except Exception as e:
+        st.error(f"Translation error: {e}")
         return text