add bettertransformer for CoLA model

app.py

@@ -1,21 +1,38 @@
 import re
 import os
+import gc
 
 from cleantext import clean
 import gradio as gr
 from tqdm.auto import tqdm
 from transformers import pipeline
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
 
 checker_model_name = "textattack/roberta-base-CoLA"
 corrector_model_name = "pszemraj/flan-t5-large-grammar-synthesis"
 
 # pipelines
-checker = pipeline(
-    "text-classification",
-    checker_model_name,
-)
 
+
+if os.environ.get("HF_DEMO_NO_USE_ONNX") is None:
+    from optimum.bettertransformer import BetterTransformer
+
+    model_hf = AutoModelForSequenceClassification.from_pretrained(checker_model_name)
+    tokenizer = AutoTokenizer.from_pretrained(checker_model_name)
+    model = BetterTransformer.transform(model_hf, keep_original_model=False)
+
+    checker = pipeline(
+        "text-classification",
+        model=model,
+        tokenizer=tokenizer,
+    )
+else:
+    checker = pipeline(
+        "text-classification",
+        checker_model_name,
+    )
+gc.collect()
 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