Update app.py

app.py

@@ -1,42 +1,50 @@
 import streamlit as st
-from transformers import pipeline
-
-# Function to load the translation pipeline based on the target language
-@st.cache_resource
-def load_translation_pipeline(target_language):
-    if target_language == 'French':
-        model_name = 'Helsinki-NLP/opus-mt-en-fr'
-    elif target_language == 'Spanish':
-        model_name = 'Helsinki-NLP/opus-mt-en-es'
-    elif target_language == 'German':
-        model_name = 'Helsinki-NLP/opus-mt-en-de'
-    else:
-        st.error('Target language not supported!')
-        return None
-    return pipeline('translation', model=model_name)
-
-# Streamlit app layout
-st.title('Language Translator')
-
-# Input text to translate
-text = st.text_area('Enter text in English to translate:')
-
-# Select target language
-target_language = st.selectbox(
-    'Select target language:',
-    ['French', 'Spanish', 'German']  # Add more languages if needed
-)
-
-# Translate button
-if st.button('Translate'):
-    if text:
-        # Load the translation pipeline based on selected language
-        translation_pipeline = load_translation_pipeline(target_language)
-        if translation_pipeline:
-            # Perform translation
-            translation = translation_pipeline(text)
-            translated_text = translation[0]['translation_text']
-            st.write(f'**Translated text in {target_language}:**')
+from transformers import MarianMTModel, MarianTokenizer
+
+# 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',
+    'Urdu': 'Helsinki-NLP/opus-mt-en-ur',
+    'Arabic': 'Helsinki-NLP/opus-mt-en-ar', 
+    # Add more language models if needed
+}
+
+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")
+    translated = model.generate(inputs)
+    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
+    return translated_text
+
+def main():
+    st.title("Multilingual Translator")
+
+    # User input for English text
+    text_to_translate = st.text_area("Enter text in English:")
+
+    # Language options
+    selected_language = st.selectbox("Select target language:", list(models.keys()))
+
+    if st.button("Translate"):
+        if text_to_translate:
+            # 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:
-        st.error('Please enter text to translate.')
+        else:
+            st.warning("Please enter text to translate.")
+
+if __name__ == "__main__":
+    main()