Update translation.py

translation.py

@@ -36,7 +36,8 @@ 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 with optimized pivot
+# Function to load appropriate translation model with optimized caching
+@st.cache_resource
 def load_model(source_lang, target_lang):
     if source_lang == target_lang:
         return _load_default_model()
@@ -44,17 +45,19 @@ 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
-    # Optimize pivot through English using preloaded models
+    # Optimized 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":
             src_to_en_tokenizer, src_to_en_model = all_models.get((source_lang, "en"), _load_model_pair(source_lang, "en")) or _load_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)
+            with torch.no_grad():
+                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), _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)
+            with torch.no_grad():
+                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
     default_tokenizer, _ = _load_default_model()
     return default_tokenizer, CombinedModel()
@@ -67,7 +70,8 @@ def _load_default_model():
     model = MarianMTModel.from_pretrained(model_name)
     return tokenizer, model
 
-# Function to perform the translation
+# Cache translation results to improve speed
+@st.cache_data
 def translate(text, source_lang, target_lang):
     if not text:
         return ""