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)

data = load_data()

def translate_to_nynorsk(model, tokenizer, text):
        tokenizer.src_lang = src_lang
        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("Translate To Nynorsk")


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