Update translation.py

translation.py

@@ -31,10 +31,29 @@ def _load_all_models():
 
 all_models = _load_all_models()
 
+# Define combined_translate outside load_model
+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
+
 # Class to handle combined translation through English pivot
 class CombinedModel:
+    def __init__(self, source_lang, target_lang, default_tokenizer, default_model):
+        self.source_lang = source_lang
+        self.target_lang = target_lang
+        self.default_tokenizer = default_tokenizer
+        self.default_model = default_model
+
     def generate(self, **kwargs):
-        return torch.tensor([combined_translate(tokenizer.decode(x, skip_special_tokens=True)) for x in kwargs['input_ids']])
+        return torch.tensor([combined_translate(tokenizer.decode(x, skip_special_tokens=True), self.source_lang, self.target_lang, self.default_tokenizer, self.default_model) for x in kwargs['input_ids']])
 
 # Function to load appropriate translation model with optimized caching
 @st.cache_resource
@@ -45,20 +64,9 @@ 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
-    # Simplified pivot through English
+    # Use simplified pivot through English with CombinedModel
     default_tokenizer, default_model = _load_default_model()
-    def combined_translate(text):
-        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
-    return default_tokenizer, CombinedModel()
+    return default_tokenizer, CombinedModel(source_lang, target_lang, default_tokenizer, default_model)
 
 # Cache resource to load default translation model
 @st.cache_resource