import streamlit as st
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast

# 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)

# Create the Streamlit app interface
st.title("Multilingual Translator")

source_text = st.text_area("Enter text to translate")
target_language = st.selectbox("Choose target language", tokenizer.lang_codes.keys())

if st.button("Translate"):
    translated_text = translate_text(model, tokenizer, source_text, target_language)
    st.write("Translated text:", translated_text)

# 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