pass ocr output through LLM for spell and grammar correction

app.py

@@ -2,16 +2,15 @@ import gradio as gr
 
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 from huggingface_hub import hf_hub_download
-from transformers import AutoModel
+from transformers import pipeline
 from ultralytics import YOLO
 from PIL import Image
-import torch
 
 def process(path, progress = gr.Progress()):
     progress(0, desc="Starting")
     LINE_MODEL_PATH = "Kansallisarkisto/multicentury-textline-detection"
-    #OCR_MODEL_PATH = "Kansallisarkisto/multicentury-htr-model"
     OCR_MODEL_PATH = "microsoft/trocr-large-handwritten"
+    CORRECTOR_PATH = "oliverguhr/spelling-correction-english-base"
 
     # Load the model and processor
     processor = TrOCRProcessor.from_pretrained(OCR_MODEL_PATH)
@@ -43,7 +42,11 @@ def process(path, progress = gr.Progress()):
         generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
         full_text += generated_text
 
-    return full_text
+    fix_spelling = pipeline("text2text-generation",model=CORRECTOR_PATH)
+    fixed_text = fix_spelling(full_text, max_new_tokens=len(full_text)+100)
+    fixed_text = fixed_text[0]['generated_text']
+
+    return fixed_text
 
 if __name__ == "__main__":
     demo = gr.Interface(fn=process, inputs=gr.Image(type="filepath"), outputs="text")