leslyarun commited on
Commit
cf1f5b8
·
1 Parent(s): e453d48

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -6
app.py CHANGED
@@ -1,17 +1,16 @@
1
  import gradio as gr
2
- from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
3
  from optimum.onnxruntime import ORTModelForSeq2SeqLM
 
4
 
5
 
6
  def check_grammar(sentence):
7
  tokenizer = AutoTokenizer.from_pretrained("leslyarun/grammatical-error-correction-quantized")
8
  model = ORTModelForSeq2SeqLM.from_pretrained("leslyarun/grammatical-error-correction-quantized")
9
 
10
- inputs = tokenizer("grammar: " + sentence, return_tensors='pt')
11
-
12
- output = model.generate(inputs['input_ids'], num_beams=5, max_length=512, early_stopping=True)
13
- correction = tokenizer.batch_decode(output, skip_special_tokens=True)
14
- return "".join(correction)
15
 
16
 
17
  demo = gr.Interface(check_grammar, inputs=['text'],
 
1
  import gradio as gr
2
+ from transformers import AutoTokenizer
3
  from optimum.onnxruntime import ORTModelForSeq2SeqLM
4
+ from optimum.pipelines import pipeline
5
 
6
 
7
  def check_grammar(sentence):
8
  tokenizer = AutoTokenizer.from_pretrained("leslyarun/grammatical-error-correction-quantized")
9
  model = ORTModelForSeq2SeqLM.from_pretrained("leslyarun/grammatical-error-correction-quantized")
10
 
11
+ text2text_generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
12
+ output = text2text_generator("grammar: " + sentence)
13
+ return output
 
 
14
 
15
 
16
  demo = gr.Interface(check_grammar, inputs=['text'],