import streamlit as st
from transformers import T5ForConditionalGeneration, AutoTokenizer

st.title("SpellCorrectorT5")
st.markdown('SpellCorrectorT5 is a fine-tuned version of **pre-trained t5-small model** modelled on randomly selected 50000 sentences modified by [imputing random noises/errors](./random_noiser.py) and trained using transformers. It not only looks for _spelling errors but also looks for the semantics_ in the sentence and suggest other possible words for the incorrect word.')
ttokenizer = AutoTokenizer.from_pretrained("./")
tmodel = T5ForConditionalGeneration.from_pretrained('./')

form = st.form("T5-form")

examples = ["I will return it to yu once it is donr",
            "Iu is going to rain",
            "Feel free to raach out to me",
            "Wheir do you live?",
            "It wis great mieting with you all"]

input_text = form.selectbox(label="Choose an example",
        options=examples)        
form.write("(or)")
input_text = form.text_input(label='Enter your own sentence', value=input_text)
submit = form.form_submit_button("Submit")

if submit:
  input_ids = ttokenizer.encode('seq: '+ input_text, return_tensors='pt')
  
  # generate text until the output length (which includes the context length) reaches 50
  outputs = tmodel.generate(
    input_ids,
    do_sample=True, 
    max_length=50,
    top_p=0.99, 
    top_k=50,
    num_return_sequences=3
  )
  
  for x in outputs:
    st.subheader("Most probable: ")
    
    out_text = ttokenizer.decode(x, skip_special_tokens=True)
    st.success(out_text.capitalize())
    
    st.markdown("### Edited sentences:")
    c_text = ""
    for x in out_text.lower().split(" "):
      if x in input_text.lower().split(" "):
        c_text = c_text + x + " "
      else:
        c_text = c_text + '<span style="font-weight:bold; color:rgb(150,255,100);">' + x + '</span>' + " "
        
    ct = c_text.capitalize()
    st.markdown(str(ct), unsafe_allow_html=True)