Update translation.py

translation.py

@@ -36,17 +36,33 @@ def combined_translate(text, source_lang, target_lang, default_tokenizer, defaul
     with torch.no_grad():
         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)
+                src_to_inter_tokenizer, src_to_inter_model = None, None
+                # Try multiple intermediates, prefer English first
+                for inter in ["en", "fr", "es", "de", "ru"]:  # Prioritize common languages
+                    pair = all_models.get((source_lang, inter))
+                    if pair and pair[0] and pair[1]:
+                        src_to_inter_tokenizer, src_to_inter_model = pair
+                        break
+                if src_to_inter_tokenizer and src_to_inter_model:
+                    inter_text = src_to_inter_tokenizer.decode(src_to_inter_model.generate(**src_to_inter_tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=500))[0], skip_special_tokens=True)
+                else:
+                    inter_text = text  # Fallback to input if no path found
             else:
-                en_text = text
+                inter_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
+                inter_to_tgt_tokenizer, inter_to_tgt_model = None, None
+                for inter in ["en", "fr", "es", "de", "ru"]:
+                    pair = all_models.get((inter, target_lang))
+                    if pair and pair[0] and pair[1]:
+                        inter_to_tgt_tokenizer, inter_to_tgt_model = pair
+                        break
+                if inter_to_tgt_tokenizer and inter_to_tgt_model:
+                    translated = inter_to_tgt_tokenizer.decode(inter_to_tgt_model.generate(**inter_to_tgt_tokenizer(inter_text, return_tensors="pt", padding=True, truncation=True, max_length=1000))[0], skip_special_tokens=True)
+                    return translated if translated.strip() else text
+            return inter_text
         return text
 
-# Class to handle combined translation through English pivot
+# Class to handle combined translation through multiple intermediates
 class CombinedModel:
     def __init__(self, source_lang, target_lang, default_tokenizer, default_model):
         self.source_lang = source_lang
@@ -78,7 +94,7 @@ def load_model(source_lang, target_lang):
             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
+    # Fallback to default model with CombinedModel
     default_tokenizer, default_model = _load_default_model()
     return default_tokenizer, CombinedModel(source_lang, target_lang, default_tokenizer, default_model)
 
@@ -101,7 +117,7 @@ def translate(text, source_lang, target_lang):
         if inputs['input_ids'].size(0) > 1:
             inputs = {k: v[0].unsqueeze(0) for k, v in inputs.items()}
         with torch.no_grad():
-            translated_ids = model.generate(**inputs, max_length=1000 if target_lang in ["hi", "ja"] else 500, num_beams=4, early_stopping=True)
+            translated_ids = model.generate(**inputs, max_length=1000 if target_lang in ["hi", "zh", "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: