import streamlit as st
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast

model_name = "facebook/mbart-large-50-one-to-many-mmt"
model = MBartForConditionalGeneration.from_pretrained(model_name)
tokenizer = MBart50TokenizerFast.from_pretrained(model_name, src_lang="en_XX")

st.title("Tradutor de Inglês")
st.write(
    f"Faça a tradução do seu texto em inglês para "
    f"{len(tokenizer.lang_code_to_id.keys())} linguagens!"
)

language = st.selectbox(
   "Lingua de saída: ",
   tokenizer.lang_code_to_id.keys(),
   index=0,
   placeholder="Selecione uma linguagem...",
)

input_text = st.text_area("Texto de entrada:")

model_inputs = tokenizer(input_text, return_tensors="pt")

generated_tokens = model.generate(
    **model_inputs,
    forced_bos_token_id=tokenizer.lang_code_to_id[language]
)

translation = tokenizer.batch_decode(
    generated_tokens,
    skip_special_tokens=True
)

st.write(f'Resultado: {translation}')