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


if 'textbox' not in st.session_state:
    st.session_state['textbox'] = "Vi bruker ikke tegnsetting eller store bokstaver når vi prater. Vi slår også sammen ord, og i praksis er dermed heller ikke mellomrom meningsbærende. Prøv å fjerne tegnsetting, store bokstaver og mellomrom fra dette avsnittet. Se om den nye North-T5-modellen greier å sette sammen til et nytt meningsbærende avsnitt."
 
@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)
    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 uncase():
    st.session_state['textbox'] = st.session_state['textbox'].lower()

def unpunct():
    text = st.session_state['textbox']
    trans_chars = "'\",.:;-_*?/\n"
    trans_table = text.maketrans("", "", trans_chars)
    st.session_state['textbox'] = text.translate(trans_table)
    
def unspace():
    st.session_state['textbox'] = st.session_state['textbox'].replace(" ","")
    
def sidebar_callback():
    st.session_state['textbox'] = st.session_state['prefilled']

st.title("DeUnCaser")
st.sidebar.write("This web app adds spaces, punctation and capitalisation back into the text.")
st.sidebar.write("You can use the examples below, but too really test the effect of the model: Write or copy text from the Internet, and then manually remove spaces, puctation, cases etc. Try to restore the text.")
option = st.sidebar.selectbox(
     "Examples:",
     ("Vi bruker ikke tegnsetting eller store bokstaver når vi prater. Vi slår også sammen ord, og i praksis er dermed heller ikke mellomrom meningsbærende. Prøv å fjerne tegnsetting, store bokstaver og mellomrom fra dette avsnittet. Se om den nye North-T5-modellen greier å sette sammen til et nytt meningsbærende avsnitt.","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", "sentpå60talletvardetfaktisknoensomkalteungensinperegilkummervoldidagerdetikkelengersåvanligåbrukedobbeltnavninorgehvasynesduomdet"),key='prefilled',on_change=sidebar_callback)

st.button('Uncase', on_click=uncase)
st.button('Remove punctation', on_click=unpunct)
st.button('Remove spaces', on_click=unspace)

text = st.text_area(f"",max_chars=1000,height=120,key="textbox")

run = st.button('Run DeUnCaser')

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