import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

def check_grammar(sentence):
  tokenizer = AutoTokenizer.from_pretrained("leslyarun/grammatical-error-correction")
  model = AutoModelForSeq2SeqLM.from_pretrained("leslyarun/grammatical-error-correction")
  
  inputs = tokenizer("grammar: " + sentence, truncation=True, return_tensors='pt')
  
  output = model.generate(inputs['input_ids'], num_beams=5, max_length=512, early_stopping=True)
  correction = tokenizer.batch_decode(output, skip_special_tokens=True)
  return "".join(correction)

  
demo = gr.Interface(check_grammar, inputs=['text'],
                    outputs="text",
                    title = "English Grammar Corrector")
                    
if __name__ == "__main__":
    demo.launch(debug=True)