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-deuncaser-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 deuncase(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)

def spaces():
    st.ta.value= "About"
    return None

st.title("DeUnCaser")

expander = st.sidebar.expander("About")
expander.write("This web app adds spaces, punctation and capitalisation back into the text.")

optionx = st.sidebar.selectbox(
     "Examples:",
     ("tirsdag var travel for ukrainas president volodymyr zelenskyj på morgenen tok han imot polens statsminister mateusz morawiecki","tirsdagvartravelforukrainaspresidentvolodymyrzelenskyjpåkveldentokhanimotpolensstatsministermateuszmorawiecki","deterikkelettåholderedepåstoreogsmåbokstavermanmåforeksempelhuskestorforbokstavnårmanskriveromkrimhalvøyamenkunbrukelitenforbokstavnårmanhenvisertilenkrimroman","detteerenlitendemosomerlagetavperegilkummervoldhanerenforskersomtidligerejobbetvednasjonalbiblioteketimoirana"))

st.sidebar.write("You can use the examples above, but for best effect: Copy text from the Internet, and remove spaces, puctation, cases etc. Try to restore the text.")

#col1, col2, col3 = st.columns([1,1,1])
#with col1:
#    st.button('Remove Casing')
#with col2:
#    st.button('Remove Punctation')
#with col3:
#    st.button('Remove Spaces')

text = st.text_area(f"Corrupted text: ",key="ta" ,max_chars=1000, value= "No text in here…")

run = st.button("Run DeUnCaser")

col1, col2, col3 = st.columns([1,1,1])

with col1:
    st.button('1')
with col2:
    st.button('2')
with col3:
    st.button('3')


col4, col5, col6 = st.columns(3)

with col4:
    st.header("A cat")
    st.image("https://static.streamlit.io/examples/cat.jpg")

with col5:
    st.header("A dog")
    st.image("https://static.streamlit.io/examples/dog.jpg")

with col6:
    st.header("An owl")
    st.image("https://static.streamlit.io/examples/owl.jpg")

#st.text("Fixed text: ")

#form = st.form(key='my_form')
#form.text_input(value="test",key="test",label='Enter some text')
#submit_button = form.form_submit_button(label='Submit')
#https://blog.streamlit.io/introducing-submit-button-and-forms/
    

#More stuff to try


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