import transformers
import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline


st.title("English-Vietnamese Text Translator")
st.write("A simple interface to translate from English to Vietnamese, and vice versa.")


tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
@st.cache
def load_model(model_name):
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
    return model

model = load_model("facebook/nllb-200-distilled-600M")

src_lang_selection = st.radio(
    "Select Your Source Language:",
    ('English', 'Vietnamese'))
if src_lang_selection == "English":
    src_lang = "eng_Latn"
    tgt_lang = "vie_Latn"
else:
    src_lang = "vie_Latn"
    tgt_lang = "eng_Latn"
# default_value = "UN Chief says there is no military solution in Syria"
sent = st.text_area("Input Your Text Here", height = 275)

if st.button("Run"):
    with st.spinner("Working Hard..."):
        translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=src_lang, tgt_lang=tgt_lang)
        trans_text = translator(sent)[0]["translation_text"]
        st.write(trans_text)
    st.success("Done!")

st.write("For feedback/requests, write to mel.nguyen273@gmail.com.")