import streamlit as st
from transformers import T5ForConditionalGeneration, T5TokenizerFast, T5Config

@st.cache(allow_output_mutation=True, suppress_st_warning=True)
def load_model():
    model_name = "north/demo-nynorsk-base"
    config = T5Config.from_pretrained(model_name)
    #Debug
    #st.text(config)
    #st.text("north/demo-nynorsk-base")
    model = T5ForConditionalGeneration.from_pretrained(model_name,config=config)
    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.")

option = st.sidebar.selectbox(
     'Examples:',
     ('Loven har også som formål å sikre et arbeidsmiljø som gir grunnlag for en helsefremmende og meningsfylt arbeidssituasjon, og bidra til et inkluderende arbeidsliv.','En av de vanskeligste oppgavene når man oversetter fra bokmål til nynorsk, er å passe på at man bruker riktige pronomen. Man kan for eksempel si at man eier en bil og at den er rød.', 'Alle søknader behandles konfidensielt.', 'Kommunens nettsider henviser til kommunens vedtak.'))

text = st.text_area(f"Bokmål text: ",max_chars=1000, value=option)

run = st.button('Run Nynorsk Translation')



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