|
import re |
|
import gradio as gr |
|
from tqdm.auto import tqdm |
|
from transformers import pipeline |
|
|
|
|
|
checker = pipeline("text-classification", "textattack/roberta-base-CoLA") |
|
corrector = pipeline("text2text-generation", "pszemraj/flan-t5-large-grammar-synthesis") |
|
|
|
def split_text(text: str) -> list: |
|
|
|
sentences = re.split(r"(?<=[^A-Z].[.?]) +(?=[A-Z])", text) |
|
|
|
|
|
sentence_batches = [] |
|
|
|
|
|
temp_batch = [] |
|
|
|
|
|
for sentence in sentences: |
|
|
|
temp_batch.append(sentence) |
|
|
|
|
|
if len(temp_batch) >= 2 and len(temp_batch) <= 3 or sentence == sentences[-1]: |
|
sentence_batches.append(temp_batch) |
|
temp_batch = [] |
|
|
|
return sentence_batches |
|
|
|
def correct_text(text: str, checker, corrector, separator: str = " ") -> str: |
|
|
|
sentence_batches = split_text(text) |
|
|
|
|
|
corrected_text = [] |
|
|
|
|
|
for batch in tqdm( |
|
sentence_batches, total=len(sentence_batches), desc="correcting text.." |
|
): |
|
|
|
raw_text = " ".join(batch) |
|
|
|
|
|
results = checker(raw_text) |
|
|
|
|
|
if results[0]["label"] != "LABEL_1" or ( |
|
results[0]["label"] == "LABEL_1" and results[0]["score"] < 0.9 |
|
): |
|
|
|
corrected_batch = corrector(raw_text) |
|
corrected_text.append(corrected_batch[0]["generated_text"]) |
|
else: |
|
corrected_text.append(raw_text) |
|
|
|
|
|
corrected_text = separator.join(corrected_text) |
|
|
|
return corrected_text |
|
|
|
def update(text: str, checker, corrector): |
|
text = text[:4000] |
|
return correct_text(text, checker, corrector) |
|
|
|
with gr.Blocks() as demo: |
|
gr.Markdown("<center># Robust Grammar Correction with FLAN-T5</center>") |
|
gr.Markdown("Enter the text you want to correct in the textbox below. The text will be truncated to 4000 characters. Then click Run to see the corrected text.") |
|
with gr.Row(): |
|
inp = gr.Textbox(placeholder="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() |
|
btn = gr.Button("Process Text") |
|
btn.click(fn=update, inputs=inp, outputs=out) |
|
|
|
demo.launch() |