Update translation.py

translation.py

@@ -8,6 +8,7 @@ LANGUAGES = {
     "ar": ("العربية", "Arabic"), "ru": ("Русский", "Russian"), "ja": ("日本語", "Japanese")
 }
 
+# Cache resource to load a specific translation model pair
 @st.cache_resource
 def _load_model_pair(source_lang, target_lang):
     try:
@@ -18,6 +19,7 @@ def _load_model_pair(source_lang, target_lang):
     except Exception:
         return None, None
 
+# Cache resource to load all possible model combinations
 @st.cache_resource
 def _load_all_models():
     models = {}
@@ -29,6 +31,12 @@ def _load_all_models():
 
 all_models = _load_all_models()
 
+# Class to handle combined translation through English pivot
+class CombinedModel:
+    def generate(self, **kwargs):
+        return torch.tensor([combined_translate(tokenizer.decode(x, skip_special_tokens=True)) for x in kwargs['input_ids']])
+
+# Function to load appropriate translation model
 def load_model(source_lang, target_lang):
     if source_lang == target_lang:
         return _load_default_model()
@@ -47,12 +55,10 @@ def load_model(source_lang, target_lang):
             en_to_tgt_tokenizer, en_to_tgt_model = all_models.get(("en", target_lang), _load_model_pair("en", target_lang)) or _load_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=500))[0], skip_special_tokens=True)
         return en_text
-    class CombinedModel:
-        def generate(self, **kwargs):
-            return torch.tensor([combined_translate(tokenizer.decode(x, skip_special_tokens=True)) for x in kwargs['input_ids']])
     tokenizer, _ = _load_default_model()
     return tokenizer, CombinedModel()
 
+# Cache resource to load default translation model
 @st.cache_resource
 def _load_default_model():
     model_name = "Helsinki-NLP/opus-mt-en-hi"
@@ -60,6 +66,7 @@ def _load_default_model():
     model = MarianMTModel.from_pretrained(model_name)
     return tokenizer, model
 
+# Function to perform the translation
 def translate(text, source_lang, target_lang):
     if not text:
         return ""