add onnx support

Signed-off-by: Peter <[email protected]>

app.py

@@ -1,13 +1,31 @@
 import re
+import os
 
 from cleantext import clean
 import gradio as gr
 from tqdm.auto import tqdm
 from transformers import pipeline
 
+
+checker_model_name = "textattack/roberta-base-CoLA"
+corrector_model_name = "pszemraj/flan-t5-large-grammar-synthesis"
+
 # pipelines
-checker = pipeline("text-classification", "textattack/roberta-base-CoLA")
-corrector = pipeline("text2text-generation", "pszemraj/flan-t5-large-grammar-synthesis")
+checker = pipeline(
+    "text-classification",
+    checker_model_name,
+)
+
+if os.environ.get("HF_DEMO_NO_USE_ONNX") is None:
+    # load onnx runtime unless HF_DEMO_NO_USE_ONNX is set
+    from optimum.pipelines import pipeline
+
+    corrector = pipeline(
+        "text2text-generation", model=corrector_model_name, accelerator="ort"
+    )
+else:
+    corrector = pipeline("text2text-generation", corrector_model_name)
+
 
 def split_text(text: str) -> list:
     # Split the text into sentences using regex
@@ -31,6 +49,7 @@ def split_text(text: str) -> list:
 
     return sentence_batches
 
+
 def correct_text(text: str, checker, corrector, separator: str = " ") -> str:
     # Split the text into sentence batches
     sentence_batches = split_text(text)
@@ -63,23 +82,35 @@ def correct_text(text: str, checker, corrector, separator: str = " ") -> str:
 
     return corrected_text
 
+
 def update(text: str):
     text = clean(text[:4000], lower=False)
     return correct_text(text, checker, corrector)
 
+
 with gr.Blocks() as demo:
     gr.Markdown("# <center>Robust Grammar Correction with FLAN-T5</center>")
-    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.")
-    gr.Markdown("""Models:
+    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."
+    )
+    gr.Markdown(
+        """Models:
     - `textattack/roberta-base-CoLA` for grammar quality detection
     - `pszemraj/flan-t5-large-grammar-synthesis` for grammar correction
-    """)
+    """
+    )
     with gr.Row():
-        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.")
+        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.",
+        )
         out = gr.Textbox(label="output", interactive=False)
     btn = gr.Button("Process")
     btn.click(fn=update, inputs=inp, outputs=out)
     gr.Markdown("---")
-    gr.Markdown("- see the [model card](https://huggingface.co/pszemraj/flan-t5-large-grammar-synthesis) for more info")
+    gr.Markdown(
+        "- see the [model card](https://huggingface.co/pszemraj/flan-t5-large-grammar-synthesis) for more info"
+    )
     gr.Markdown("- if experiencing long wait times, feel free to duplicate the space!")
-demo.launch()
+demo.launch()

requirements.txt

@@ -2,4 +2,6 @@ transformers
 gradio
 tqdm
 torch
-clean-text
+clean-text
+accelerate
+optimum[onnxruntime]