Update app.py

app.py

@@ -1,11 +1,13 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from optimum.onnxruntime import ORTModelForSeq2SeqLM
+
 
 def check_grammar(sentence):
-  tokenizer = AutoTokenizer.from_pretrained("leslyarun/grammatical-error-correction")
-  model = AutoModelForSeq2SeqLM.from_pretrained("leslyarun/grammatical-error-correction")
+  tokenizer = AutoTokenizer.from_pretrained("leslyarun/grammatical-error-correction-quantized")
+  model = ORTModelForSeq2SeqLM.from_pretrained("leslyarun/grammatical-error-correction-quantized")
   
-  inputs = tokenizer("grammar: " + sentence, truncation=True, return_tensors='pt')
+  inputs = tokenizer("grammar: " + sentence, 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)