import streamlit as st
from transformers import T5ForConditionalGeneration, AutoTokenizer

st.title("SimpleGrammarlyT5")
st.markdown('SimpleGrammarlyT5 is a fine-tuned version of **pre-trained t5-small model** modelled on randomly selected 50000 sentences modified by imputing random noises/errors 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")
input_text = form.text_input(label='Enter a random sentence')
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.98, 
    num_return_sequences=3
  )
  
  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)