Update translation.py

translation.py

@@ -34,15 +34,17 @@ all_models = _load_all_models()
 # Define combined_translate outside load_model with explicit parameters
 def combined_translate(text, source_lang, target_lang, default_tokenizer, default_model):
     with torch.no_grad():
-        if source_lang != "en":
-            src_to_en_tokenizer, src_to_en_model = all_models.get((source_lang, "en"), (default_tokenizer, 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), (default_tokenizer, 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=1000))[0], skip_special_tokens=True)
-        return en_text
+        if source_lang != target_lang:  # Only translate if languages differ
+            if source_lang != "en":
+                src_to_en_tokenizer, src_to_en_model = all_models.get((source_lang, "en"), (default_tokenizer, 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), (default_tokenizer, default_model))
+                translated = 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=1000))[0], skip_special_tokens=True)
+                return translated if translated.strip() else text
+        return text
 
 # Class to handle combined translation through English pivot
 class CombinedModel:
@@ -54,9 +56,11 @@ class CombinedModel:
 
     def generate(self, **kwargs):
         input_ids = kwargs.get('input_ids')
-        if not input_ids:
+        if not input_ids or input_ids.size(0) == 0:
             return torch.tensor([])
-        return torch.tensor([combined_translate(self.default_tokenizer.decode(x, skip_special_tokens=True), self.source_lang, self.target_lang, self.default_tokenizer, self.default_model) for x in input_ids])
+        inputs = self.default_tokenizer.batch_decode(input_ids, skip_special_tokens=True)
+        translated = [combined_translate(text, self.source_lang, self.target_lang, self.default_tokenizer, self.default_model) for text in inputs]
+        return torch.tensor([self.default_tokenizer.encode(t, return_tensors="pt", padding=True, truncation=True, max_length=500)[0] for t in translated])
 
 # Function to load appropriate translation model with optimized caching
 @st.cache_resource
@@ -67,13 +71,14 @@ 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
-    # Prefer direct model if available, then pivot
-    for src in [source_lang, "en"]:
-        for tgt in [target_lang, "en"]:
-            if src != tgt:
-                pair = all_models.get((src, tgt))
-                if pair and pair[0] and pair[1]:
-                    return pair
+    # Try to find the best path through any intermediate language
+    for inter in LANGUAGES.keys():
+        if inter != source_lang and inter != target_lang:
+            pair1 = all_models.get((source_lang, inter))
+            pair2 = all_models.get((inter, target_lang))
+            if pair1 and pair1[0] and pair1[1] and pair2 and pair2[0] and pair2[1]:
+                return pair1
+    # Fallback to pivot through English
     default_tokenizer, default_model = _load_default_model()
     return default_tokenizer, CombinedModel(source_lang, target_lang, default_tokenizer, default_model)
 
@@ -93,11 +98,11 @@ def translate(text, source_lang, target_lang):
     try:
         tokenizer, model = load_model(source_lang, target_lang)
         inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=500)
-        if inputs['input_ids'].size(0) > 1:  # Ensure single sequence
+        if inputs['input_ids'].size(0) > 1:
             inputs = {k: v[0].unsqueeze(0) for k, v in inputs.items()}
         with torch.no_grad():
-            translated = model.generate(**inputs, max_length=1000 if target_lang == "ja" else 500, num_beams=4, early_stopping=True)
-        result = tokenizer.decode(translated[0], skip_special_tokens=True)
+            translated_ids = model.generate(**inputs, max_length=1000 if target_lang in ["hi", "ja"] else 500, num_beams=4, early_stopping=True)
+        result = tokenizer.decode(translated_ids[0], skip_special_tokens=True)
         return result if result.strip() else text
     except Exception as e:
         st.error(f"Translation error: {e}")