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 =["Look if ther is fire on the top",
           "Where os you're car?",
           "Iu is going to rain",
           "Feel free to raach out to me",
           "Will return it to yu once it is donr",
           "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, 
    num_return_sequences=2
  )
  
  st.subheader("Suggested sentences: ")
  
  i = 0
  for x in outputs:
      out_text = ttokenizer.decode(x, skip_special_tokens=True)
      i = i + 1
      st.success(str(i) + '. ' + out_text)