Update translation.py

translation.py

@@ -45,13 +45,11 @@ 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
-    # Use direct English pivot only if direct model unavailable
+    # Use direct English pivot only if necessary
     if source_lang != "en" and target_lang != "en":
-        en_pivot_pair = all_models.get(("en", target_lang)) or _load_model_pair("en", target_lang)
+        en_pivot_pair = all_models.get((source_lang, "en")) or _load_model_pair(source_lang, "en")
         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
+            return en_pivot_pair
     default_tokenizer, _ = _load_default_model()
     def combined_translate(text):
         with torch.no_grad():
@@ -83,8 +81,9 @@ def translate(text, source_lang, target_lang):
         tokenizer, model = load_model(source_lang, target_lang)
         inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=500)
         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)
+            translated = model.generate(**inputs, max_length=500, num_beams=4, early_stopping=True)  # Increased beams for better accuracy
+        result = tokenizer.decode(translated[0], skip_special_tokens=True)
+        return result if result.strip() else text
     except Exception as e:
         st.error(f"Translation error: {e}")
         return text