import gradio as gr
from transformers import AutoTokenizer
from optimum.onnxruntime import ORTModelForSeq2SeqLM
from optimum.pipelines import pipeline


def check_grammar(sentence):
  tokenizer = AutoTokenizer.from_pretrained("leslyarun/grammatical-error-correction-quantized")
  model = ORTModelForSeq2SeqLM.from_pretrained("leslyarun/grammatical-error-correction-quantized",
                                                   encoder_file_name="encoder_model_quantized.onnx",
                                                   decoder_file_name="decoder_model_quantized.onnx",
                                                   decoder_with_past_file_name="decoder_with_past_model_quantized.onnx")
  
  text2text_generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
  output = text2text_generator("grammar: " + sentence)
  return output[0]["generated_text"]

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