import streamlit as st
from transformers import T5ForConditionalGeneration, T5TokenizerFast

@st.cache(allow_output_mutation=True, suppress_st_warning=True)
def load_model():
    model_name = "north/demo-nynorsk-base"
    model = T5ForConditionalGeneration.from_pretrained(model_name)
    tokenizer = T5TokenizerFast.from_pretrained(model_name)
    return (model, tokenizer)

def translate_to_nynorsk(model, tokenizer, text):
        encoded_txt = tokenizer(text, return_tensors="pt")
        generated_tokens = model.generate(
            **encoded_txt
        )
        return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)

st.title("Bokmål To Nynorsk")

expander = st.sidebar.expander("About")
expander.write("This web app allows you to translate a text in Norwegian Bokmål to Norwegian Nynorsk.")

st.sidebar.header("Choose or write your own.")

option = st.selectbox(
     'How would you like to be contacted?',
     ('Email', 'Home phone', 'Mobile phone'))

st.write('You selected:', option)

text = st.text_input(f"Bokmål text: ")

#st.text("What you wrote: ")
#st.write(text)

st.text("Nynorsk Translation: ")

if text:
    model, tokenizer = load_model()
    translated_text = translate_to_nynorsk(model, tokenizer, text)
    st.write(translated_text[0] if translated_text else "Unknown Error Translating Text")