Update translation.py

translation.py

@@ -36,14 +36,15 @@ 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
+# Function to load appropriate translation model with optimized pivot
 def load_model(source_lang, target_lang):
     if source_lang == target_lang:
         return _load_default_model()
     model_key = (source_lang, target_lang)
-    if all_models.get(model_key) and all_models[model_key][0] and all_models[model_key][1]:
-        return all_models[model_key]
-    # Pivot through English
+    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
+    # Optimize pivot through English using preloaded models
     def combined_translate(text):
         en_tokenizer, en_model = all_models.get(("en", "en"), _load_default_model())
         if source_lang != "en":
@@ -55,8 +56,8 @@ 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
-    tokenizer, _ = _load_default_model()
-    return tokenizer, CombinedModel()
+    default_tokenizer, _ = _load_default_model()
+    return default_tokenizer, CombinedModel()
 
 # Cache resource to load default translation model
 @st.cache_resource