Update app.py

app.py

@@ -1,34 +1,25 @@
 
 import streamlit as st
-from transformers import pipeline
+from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
 
-def main():
-    st.title("Multilingual Translator")
+# Load the multilingual translation model and tokenizer
+model_name = "facebook/mbart-large-50"  # Choose a suitable model
+tokenizer = MBart50TokenizerFast.from_pretrained(model_name)
+model = MBartForConditionalGeneration.from_pretrained(model_name)
 
-    # Get user input
-    input_text = st.text_area("Enter text to translate:")
+# Create the Streamlit app interface
+st.title("Multilingual Translator")
 
-    # Select source and target languages
-    source_lang = st.selectbox("Select source language:", get_languages())
-    target_lang = st.selectbox("Select target language:", get_languages())
+source_text = st.text_area("Enter text to translate")
+target_language = st.selectbox("Choose target language", tokenizer.lang_codes.keys())
 
-    # Translate text
-    if st.button("Translate"):
-        if input_text:
-            translation = translate_text(input_text, source_lang, target_lang)
-            st.success(f"Translated text: {translation}")
-        else:
-            st.warning("Please enter text to translate.")
+if st.button("Translate"):
+    translated_text = translate_text(model, tokenizer, source_text, target_language)
+    st.write("Translated text:", translated_text)
 
-def get_languages():
-    # You can customize this list based on the languages you want to support
-    return ["en", "es", "fr", "de", "it", "ja", "ko", "zh-CN"]
-
-def translate_text(text, source_lang, target_lang):
-    translator = pipeline(task="translation", model=f"{source_lang}-to-{target_lang}")
-    translation = translator(text, max_length=500)[0]['translation_text']
-    return translation
-
-if __name__ == "__main__":
-    main()
-streamlit run app.py
+# Define the translation function
+def translate_text(model, tokenizer, source_text, target_language):
+    inputs = tokenizer(source_text, return_tensors="pt")
+    outputs = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id[target_language])
+    translated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+    return translated_text