Update app.py

app.py

@@ -1,11 +1,17 @@
 import streamlit as st
 from transformers import MarianMTModel, MarianTokenizer, pipeline
 
-# Function to load models dynamically
-def load_translation_model(model_name):
-    tokenizer = MarianTokenizer.from_pretrained(model_name)
-    model = MarianMTModel.from_pretrained(model_name)
-    return pipeline("translation", model=model, tokenizer=tokenizer)
+# Preload translation models
+models = {
+    "en-ur": pipeline("translation", model=MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-en-ur"),
+                      tokenizer=MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ur")),
+    "ur-en": pipeline("translation", model=MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-ur-en"),
+                      tokenizer=MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ur-en")),
+    "en-fr": pipeline("translation_en_to_fr"),
+    "fr-en": pipeline("translation_fr_to_en"),
+    "en-de": pipeline("translation_en_to_de"),
+    "de-en": pipeline("translation_de_to_en")
+}
 
 # Streamlit UI
 st.title("Language Translator")
@@ -16,25 +22,14 @@ input_lang = st.sidebar.selectbox("Select Input Language", ["English", "Urdu", "
 output_lang = st.sidebar.selectbox("Select Output Language", ["Urdu", "English", "French", "Spanish", "German"])
 
 # Text input
-input_text = st.text_area("Enter Text to Translate", placeholder =  "Type here...")
+input_text = st.text_area("Enter Text to Translate",placeholder =  "Type here...")
 
 # Translation logic
 translator = None
 if st.button("Translate"):
-    if input_lang == "English" and output_lang == "Urdu":
-        translator = load_translation_model("Helsinki-NLP/opus-mt-en-ur")
-    elif input_lang == "Urdu" and output_lang == "English":
-        translator = load_translation_model("Helsinki-NLP/opus-mt-ur-en")
-    elif input_lang == "English" and output_lang == "French":
-        translator = pipeline("translation_en_to_fr")
-    elif input_lang == "French" and output_lang == "English":
-        translator = pipeline("translation_fr_to_en")
-    elif input_lang == "English" and output_lang == "German":
-        translator = pipeline("translation_en_to_de")
-    elif input_lang == "German" and output_lang == "English":
-        translator = pipeline("translation_de_to_en")
-
-    if translator:
+    model_key = f"{input_lang.lower()[0:2]}-{output_lang.lower()[0:2]}"
+    if model_key in models:
+        translator = models[model_key]
         translation = translator(input_text)
         st.write("Translated Text:")
         st.write(translation[0]['translation_text'])