Adding smoldocling OCR

app.py

@@ -6,6 +6,14 @@ from docling.datamodel.base_models import InputFormat
 from paddleocr import PPStructureV3
 from pdf2image import convert_from_path
 import numpy as np
+import torch
+from docling_core.types.doc import DoclingDocument
+from docling_core.types.doc.document import DocTagsDocument
+from transformers import AutoProcessor, AutoModelForVision2Seq
+from transformers.image_utils import load_image
+from pathlib import Path
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 pipeline_options = PdfPipelineOptions(enable_remote_services=True)
 converter = DocumentConverter(
@@ -14,6 +22,15 @@ converter = DocumentConverter(
     }
 )
 
+pipeline = PPStructureV3()
+
+processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview")
+model = AutoModelForVision2Seq.from_pretrained(
+    "ds4sd/SmolDocling-256M-preview",
+    torch_dtype=torch.bfloat16,
+    _attn_implementation="flash_attention_2" if DEVICE == "cuda" else "eager",
+).to(DEVICE)
+
 def get_pdf_page_count(pdf_path):
     reader = PdfReader(pdf_path)
     return len(reader.pages)
@@ -24,7 +41,6 @@ def get_docling_ocr(pdf_path, page_num):
     return markdown_text_docling
 
 def get_paddle_ocr(page_image):
-    pipeline = PPStructureV3()
     output = pipeline.predict(input=np.array(page_image))
 
     markdown_list = []
@@ -36,13 +52,49 @@ def get_paddle_ocr(page_image):
     markdown_text_paddleOCR = pipeline.concatenate_markdown_pages(markdown_list)
     return markdown_text_paddleOCR
 
+def get_smoldocling_ocr(page_image):
+    image = load_image(page_image)
+
+    # Create input messages
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image"},
+                {"type": "text", "text": "Convert this page to docling."}
+            ]
+        },
+    ]
+
+    # Prepare inputs
+    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
+    inputs = processor(text=prompt, images=[image], return_tensors="pt")
+    inputs = inputs.to(DEVICE)
+
+    generated_ids = model.generate(**inputs, max_new_tokens=8192)
+    prompt_length = inputs.input_ids.shape[1]
+    trimmed_generated_ids = generated_ids[:, prompt_length:]
+    doctags = processor.batch_decode(
+        trimmed_generated_ids,
+        skip_special_tokens=False,
+    )[0].lstrip()
+
+    # Populate document
+    doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([doctags], [image])
+    doc = DoclingDocument.load_from_doctags(doctags_doc, document_name="Document")
+
+    markdown_text_smoldocling = doc.export_to_markdown()
+    return markdown_text_smoldocling
+    
+
 def inference(pdf_path, page_num):
     docling_ocr = get_docling_ocr(pdf_path, page_num)
     # Extract the first page as an image
     images = convert_from_path(pdf_path, first_page=page_num, last_page=page_num)
     page_image = images[0]
     paddle_ocr = get_paddle_ocr(page_image)
-    return docling_ocr, paddle_ocr
+    smoldocling_ocr = get_smoldocling_ocr(page_image)
+    return docling_ocr, paddle_ocr, smoldocling_ocr
 
 title = "OCR Arena"
 description = "A simple Gradio interface to extract text from PDFs and compare OCR models"
@@ -66,11 +118,12 @@ with gr.Blocks(theme=gr.themes.Glass()) as demo:
                         clear_btn = gr.ClearButton(components=[pdf, page_num])
                         submit_btn = gr.Button("Submit", variant='primary')
 
-                    submit_btn.click(inference, inputs=[pdf, page_num], outputs=[docling_ocr_out, paddle_ocr_out])
+                    submit_btn.click(inference, inputs=[pdf, page_num], outputs=[docling_ocr_out, paddle_ocr_out, smoldocling_ocr_out])
 
         with gr.Column():
             docling_ocr_out = gr.Textbox(label="Docling OCR Output", type="text")
             paddle_ocr_out = gr.Textbox(label="Paddle OCR Output", type="text")
+            smoldocling_ocr_out = gr.Textbox(label="SmolDocling OCR Output", type="text")
 
     examples_obj = gr.Examples(examples=examples, inputs=[pdf])