Update app.py

app.py

@@ -1,42 +1,48 @@
 import streamlit as st
 from transformers import MarianMTModel, MarianTokenizer
 
-# Define language options and corresponding language codes
+# Load pre-trained MarianMT model and tokenizer
+model_name = "Helsinki-NLP/opus-mt-en-xx"  # Multilingual model supporting various languages
+model = MarianMTModel.from_pretrained(model_name)
+tokenizer = MarianTokenizer.from_pretrained(model_name)
+
+# Supported languages and their target language codes
 LANGUAGES = {
-    "Spanish": "es",
-    "French": "fr",
-    "German": "de",
-    "Italian": "it",
-    "Portuguese": "pt",
+    'French': 'fr',
+    'Spanish': 'es',
+    'German': 'de',
+    'Chinese': 'zh',
+    'Russian': 'ru',
+    'Japanese': 'ja',
+    'Arabic': 'ar',
+    'Urdu': 'ur',
+    'Hindi': 'hi',
+    'Bengali': 'bn',
 }
 
-# Load translation model and tokenizer
-model_name = "Helsinki-NLP/opus-mt-en-ROMANCE"
-tokenizer = MarianTokenizer.from_pretrained(model_name)
-model = MarianMTModel.from_pretrained(model_name)
-
-def translate_text(text, target_language):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-    try:
-        # Generate translations
-        translated = model.generate(**inputs)
-        # Decode the output
-        translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
-        return translated_text
-    except Exception as e:
-        return f"Error: {str(e)}"
+def translate_text(text, target_lang):
+    # Encode the text and prepare it for translation
+    encoded_text = tokenizer(text, return_tensors="pt")
+    # Translate text
+    translated = model.generate(**encoded_text, forced_bos_token_id=tokenizer.get_lang_id(target_lang))
+    # Decode the translated text
+    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
+    return translated_text
 
 # Streamlit app
-st.title("Language Translator")
+st.title('Language Translator')
 
-# User input
-text_to_translate = st.text_area("Enter text in English:", "")
-selected_language = st.selectbox("Select target language:", list(LANGUAGES.keys()))
+# Input text and language selection
+text_to_translate = st.text_area("Enter text to translate")
+target_language = st.selectbox("Select target language", list(LANGUAGES.keys()))
 
-if st.button("Translate"):
+# Translate button
+if st.button('Translate'):
     if text_to_translate:
-        translated_text = translate_text(text_to_translate, LANGUAGES[selected_language])
-        st.write(f"Translated text ({selected_language}):")
-        st.write(translated_text)
+        if target_language in LANGUAGES:
+            translated_text = translate_text(text_to_translate, LANGUAGES[target_language])
+            st.write(f"Translated text ({target_language}): {translated_text}")
+        else:
+            st.error("Target language not supported.")
     else:
-        st.warning("Please enter text to translate.")
+        st.error("Please enter text to translate.")