Peter commited on
Commit
84b8f8b
1 Parent(s): 151824c

add onnx support

Browse files

Signed-off-by: Peter <74869040+pszemraj@users.noreply.github.com>

Files changed (2) hide show
  1. app.py +39 -8
  2. requirements.txt +3 -1
app.py CHANGED
@@ -1,13 +1,31 @@
1
  import re
 
2
 
3
  from cleantext import clean
4
  import gradio as gr
5
  from tqdm.auto import tqdm
6
  from transformers import pipeline
7
 
 
 
 
 
8
  # pipelines
9
- checker = pipeline("text-classification", "textattack/roberta-base-CoLA")
10
- corrector = pipeline("text2text-generation", "pszemraj/flan-t5-large-grammar-synthesis")
 
 
 
 
 
 
 
 
 
 
 
 
 
11
 
12
  def split_text(text: str) -> list:
13
  # Split the text into sentences using regex
@@ -31,6 +49,7 @@ def split_text(text: str) -> list:
31
 
32
  return sentence_batches
33
 
 
34
  def correct_text(text: str, checker, corrector, separator: str = " ") -> str:
35
  # Split the text into sentence batches
36
  sentence_batches = split_text(text)
@@ -63,23 +82,35 @@ def correct_text(text: str, checker, corrector, separator: str = " ") -> str:
63
 
64
  return corrected_text
65
 
 
66
  def update(text: str):
67
  text = clean(text[:4000], lower=False)
68
  return correct_text(text, checker, corrector)
69
 
 
70
  with gr.Blocks() as demo:
71
  gr.Markdown("# <center>Robust Grammar Correction with FLAN-T5</center>")
72
- gr.Markdown("**Instructions:** Enter the text you want to correct in the textbox below (_text will be truncated to 4000 characters_). Click 'Process' to run.")
73
- gr.Markdown("""Models:
 
 
 
74
  - `textattack/roberta-base-CoLA` for grammar quality detection
75
  - `pszemraj/flan-t5-large-grammar-synthesis` for grammar correction
76
- """)
 
77
  with gr.Row():
78
- inp = gr.Textbox(label="input", placeholder="PUT TEXT TO CHECK & CORRECT BROSKI", value="I wen to the store yesturday to bye some food. I needd milk, bread, and a few otter things. The store was really crowed and I had a hard time finding everyting I needed. I finaly made it to the check out line and payed for my stuff.")
 
 
 
 
79
  out = gr.Textbox(label="output", interactive=False)
80
  btn = gr.Button("Process")
81
  btn.click(fn=update, inputs=inp, outputs=out)
82
  gr.Markdown("---")
83
- gr.Markdown("- see the [model card](https://huggingface.co/pszemraj/flan-t5-large-grammar-synthesis) for more info")
 
 
84
  gr.Markdown("- if experiencing long wait times, feel free to duplicate the space!")
85
- demo.launch()
 
1
  import re
2
+ import os
3
 
4
  from cleantext import clean
5
  import gradio as gr
6
  from tqdm.auto import tqdm
7
  from transformers import pipeline
8
 
9
+
10
+ checker_model_name = "textattack/roberta-base-CoLA"
11
+ corrector_model_name = "pszemraj/flan-t5-large-grammar-synthesis"
12
+
13
  # pipelines
14
+ checker = pipeline(
15
+ "text-classification",
16
+ checker_model_name,
17
+ )
18
+
19
+ if os.environ.get("HF_DEMO_NO_USE_ONNX") is None:
20
+ # load onnx runtime unless HF_DEMO_NO_USE_ONNX is set
21
+ from optimum.pipelines import pipeline
22
+
23
+ corrector = pipeline(
24
+ "text2text-generation", model=corrector_model_name, accelerator="ort"
25
+ )
26
+ else:
27
+ corrector = pipeline("text2text-generation", corrector_model_name)
28
+
29
 
30
  def split_text(text: str) -> list:
31
  # Split the text into sentences using regex
 
49
 
50
  return sentence_batches
51
 
52
+
53
  def correct_text(text: str, checker, corrector, separator: str = " ") -> str:
54
  # Split the text into sentence batches
55
  sentence_batches = split_text(text)
 
82
 
83
  return corrected_text
84
 
85
+
86
  def update(text: str):
87
  text = clean(text[:4000], lower=False)
88
  return correct_text(text, checker, corrector)
89
 
90
+
91
  with gr.Blocks() as demo:
92
  gr.Markdown("# <center>Robust Grammar Correction with FLAN-T5</center>")
93
+ gr.Markdown(
94
+ "**Instructions:** Enter the text you want to correct in the textbox below (_text will be truncated to 4000 characters_). Click 'Process' to run."
95
+ )
96
+ gr.Markdown(
97
+ """Models:
98
  - `textattack/roberta-base-CoLA` for grammar quality detection
99
  - `pszemraj/flan-t5-large-grammar-synthesis` for grammar correction
100
+ """
101
+ )
102
  with gr.Row():
103
+ inp = gr.Textbox(
104
+ label="input",
105
+ placeholder="PUT TEXT TO CHECK & CORRECT BROSKI",
106
+ value="I wen to the store yesturday to bye some food. I needd milk, bread, and a few otter things. The store was really crowed and I had a hard time finding everyting I needed. I finaly made it to the check out line and payed for my stuff.",
107
+ )
108
  out = gr.Textbox(label="output", interactive=False)
109
  btn = gr.Button("Process")
110
  btn.click(fn=update, inputs=inp, outputs=out)
111
  gr.Markdown("---")
112
+ gr.Markdown(
113
+ "- see the [model card](https://huggingface.co/pszemraj/flan-t5-large-grammar-synthesis) for more info"
114
+ )
115
  gr.Markdown("- if experiencing long wait times, feel free to duplicate the space!")
116
+ demo.launch()
requirements.txt CHANGED
@@ -2,4 +2,6 @@ transformers
2
  gradio
3
  tqdm
4
  torch
5
- clean-text
 
 
 
2
  gradio
3
  tqdm
4
  torch
5
+ clean-text
6
+ accelerate
7
+ optimum[onnxruntime]