Update app.py

app.py

@@ -1,34 +1,26 @@
 import streamlit as st
 from transformers import MarianMTModel, MarianTokenizer
 
-# Load the MarianMT model and tokenizer
-model_name = "Helsinki-NLP/opus-mt-en-many-to-many-mmt"
-model = MarianMTModel.from_pretrained(model_name)
-tokenizer = MarianTokenizer.from_pretrained(model_name)
-
-# Define target language codes directly
-language_codes = {
-    'French': 'fr',
-    'German': 'de',
-    'Italian': 'it',
-    'Portuguese': 'pt',
-    'Spanish': 'es',
-    'Chinese': 'zh',
-    'Japanese': 'ja',
-    'Korean': 'ko',
-    'Russian': 'ru',
-    'Arabic': 'ar',
-    'Turkish': 'tr',
-    'Hindi': 'hi',
-    'Swahili': 'sw',
-    # Add more languages if needed
+# Define a dictionary to map language names to model identifiers
+models = {
+    'French': 'Helsinki-NLP/opus-mt-en-fr',
+    'Spanish': 'Helsinki-NLP/opus-mt-en-es',
+    'German': 'Helsinki-NLP/opus-mt-en-de',
+    'Italian': 'Helsinki-NLP/opus-mt-en-it',
+    'Portuguese': 'Helsinki-NLP/opus-mt-en-pt',
+    # Add more language models if needed
 }
 
-def translate_text(text, target_lang_code):
-    # Prepare the input and translate
+def load_model(model_name):
+    """Load the model and tokenizer based on the selected model name."""
+    model = MarianMTModel.from_pretrained(model_name)
+    tokenizer = MarianTokenizer.from_pretrained(model_name)
+    return model, tokenizer
+
+def translate_text(text, model, tokenizer):
+    """Translate text using the provided model and tokenizer."""
     inputs = tokenizer.encode(text, return_tensors="pt")
-    # Generate the translation
-    translated = model.generate(inputs, decoder_start_token_id=tokenizer.lang_code_to_id[target_lang_code])
+    translated = model.generate(inputs)
     translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
     return translated_text
 
@@ -39,12 +31,15 @@ def main():
     text_to_translate = st.text_area("Enter text in English:")
 
     # Language options
-    selected_language = st.selectbox("Select target language:", list(language_codes.keys()))
+    selected_language = st.selectbox("Select target language:", list(models.keys()))
 
     if st.button("Translate"):
         if text_to_translate:
-            target_lang_code = language_codes[selected_language]
-            translated_text = translate_text(text_to_translate, target_lang_code)
+            # Load the selected model
+            model_name = models[selected_language]
+            model, tokenizer = load_model(model_name)
+            
+            translated_text = translate_text(text_to_translate, model, tokenizer)
             st.write(f"**Translation in {selected_language}:**")
             st.write(translated_text)
         else: