Add support for SmolDocling

- Add support for SmolDocling
- Update benchmark and refactor code
- Add script to plot model comparison

Author: dilithjay

README.md

@@ -141,4 +141,5 @@ _Note:_ Benchmarks are currently done in the zero-shot setting.
 | 19 | google/gemma-3-27b-it | 0.624 (±0.357) | 0.750 (±0.327) | 24.51 | 0.00020 |
 | 20 | gpt-4.1 | 0.622 (±0.314) | 0.782 (±0.191) | 34.66 | 0.01461 |
 | 21 | meta-llama/Llama-3.2-90B-Vision-Instruct-Turbo | 0.559 (±0.233) | 0.822 (±0.119) | 27.74 | 0.01102 |
-| 22 | qwen/qwen-2.5-vl-7b-instruct | 0.469 (±0.364) | 0.617 (±0.441) | 13.23 | 0.00060 |
+| 22 | ds4sd/SmolDocling-256M-preview | 0.486 (±0.378) | 0.583 (±0.355) | 108.91 | 0.00000 |
+| 23 | qwen/qwen-2.5-vl-7b-instruct | 0.469 (±0.364) | 0.617 (±0.441) | 13.23 | 0.00060 |

docs/benchmark.rst

@@ -215,6 +215,12 @@ Here are the detailed parsing performance results for various models, sorted by
      - 27.74
      - 0.01102
    * - 22
+     - ds4sd/SmolDocling-256M-preview
+     - 0.486 (±0.378)
+     - 0.583 (±0.355)
+     - 108.91
+     - 0.00000
+   * - 23
      - qwen/qwen-2.5-vl-7b-instruct
      - 0.469 (±0.364)
      - 0.617 (±0.441)

examples/example_notebook.ipynb

@@ -40,6 +40,44 @@
     "# os.environ[\"OPENAI_API_KEY\"] = \"<>\""
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2025-05-25 20:06:10.839\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mlexoid.api\u001b[0m:\u001b[36mparse_chunk\u001b[0m:\u001b[36m66\u001b[0m - \u001b[34m\u001b[1mUsing LLM parser\u001b[0m\n",
+      "Truncation was not explicitly activated but `max_length` is provided a specific value, please use `truncation=True` to explicitly truncate examples to max length. Defaulting to 'longest_first' truncation strategy. If you encode pairs of sequences (GLUE-style) with the tokenizer you can select this strategy more precisely by providing a specific strategy to `truncation`.\n"
+     ]
+    }
+   ],
+   "source": [
+    "result = parse(\"inputs/cvs_coupon.jpg\", parser_type=\"LLM_PARSE\", model=\"ds4sd/SmolDocling-256M-preview\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'<doctag><section_header_level_1><loc_148><loc_102><loc_423><loc_149>$2.00 off</section_header_level_1>\\n<text><loc_148><loc_166><loc_423><loc_233>$2 off CVS HEALTH Topical Pain products</text>\\n<text><loc_148><loc_231><loc_423><loc_244>Expires 11/02/2024 (up to $2.00 value)</text>\\n<text><loc_182><loc_286><loc_345><loc_298>7168 4009 6700 2002</text>\\n<text><loc_109><loc_298><loc_423><loc_326>ExtraCare card required. See www.cvs.com/couponpolicy or policy at register for details.</text>\\n<text><loc_109><loc_326><loc_278><loc_340>ExtraCare Card</text>\\n<text><loc_278><loc_326><loc_293><loc_340>1</text>\\n<text><loc_293><loc_326><loc_298><loc_340>7</text>\\n<text><loc_298><loc_326><loc_303><loc_340>1</text>\\n<text><loc_259><loc_365><loc_264><loc_379>n</text>\\n</doctag>'"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "result[\"raw\"]"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 3,

lexoid/core/parse_type/llm_parser.py

@@ -1,3 +1,4 @@
+from PIL import Image
 import base64
 import io
 import mimetypes
@@ -14,6 +15,8 @@
 from openai import OpenAI
 from requests.exceptions import HTTPError
 from together import Together
+import torch
+from transformers import AutoProcessor, AutoModelForImageTextToText
 
 from lexoid.core.prompt_templates import (
     INSTRUCTIONS_ADD_PG_BREAK,
@@ -65,13 +68,18 @@ def get_api_provider_for_model(model: str) -> str:
         return "fireworks"
     if model.startswith("claude"):
         return "anthropic"
+    if model.startswith("ds4sd/SmolDocling"):
+        return "local"
     raise ValueError(f"Unsupported model: {model}")
 
 
 @retry_on_http_error
 def parse_llm_doc(path: str, **kwargs) -> List[Dict] | str:
-    if "api_provider" in kwargs and kwargs["api_provider"]:
-        return parse_with_api(path, api=kwargs["api_provider"], **kwargs)
+    if "api_provider" in kwargs:
+        if kwargs["api_provider"] == "local":
+            return parse_with_local_model(path, **kwargs)
+        elif kwargs["api_provider"]:
+            return parse_with_api(path, api=kwargs["api_provider"], **kwargs)
 
     model = kwargs.get("model", "gemini-2.0-flash")
     kwargs["model"] = model
@@ -80,8 +88,70 @@ def parse_llm_doc(path: str, **kwargs) -> List[Dict] | str:
 
     if api_provider == "gemini":
         return parse_with_gemini(path, **kwargs)
-    else:
-        return parse_with_api(path, api=api_provider, **kwargs)
+    elif api_provider == "local":
+        return parse_with_local_model(path, **kwargs)
+    return parse_with_api(path, api=api_provider, **kwargs)
+
+
+def parse_with_local_model(path: str, **kwargs) -> List[Dict] | str:
+    # Reference: https://huggingface.co/spaces/aabdullah27/SmolDocling-OCR-App/blob/main/app.py
+    model_name = kwargs.get("model", "ds4sd/SmolDocling-256M-preview")
+    if not model_name.startswith("ds4sd/SmolDocling"):
+        raise ValueError(
+            f"Local model parsing is only supported for 'ds4sd/SmolDocling*', got {model_name}"
+        )
+    processor = AutoProcessor.from_pretrained(model_name)
+    model = AutoModelForImageTextToText.from_pretrained(model_name)
+    images = convert_path_to_images(path)
+    pil_images = [
+        Image.open(io.BytesIO(base64.b64decode(image.split(",")[1])))
+        for _, image in images
+    ]
+    intruction = kwargs.get("docling_command", "Convert this page to docling.")
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image"},
+                {"type": "text", "text": intruction},
+            ],
+        },
+        {
+            "role": "assistant",
+            "content": [
+                {"type": "text", "text": "<output>\n"},
+            ],
+        },
+    ]
+    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
+    inputs = processor(text=prompt, images=pil_images, return_tensors="pt")
+    with torch.no_grad():
+        generated_ids = model.generate(
+            **inputs, max_new_tokens=1500, do_sample=False, num_beams=1, temperature=1.0
+        )
+
+    prompt_length = inputs.input_ids.shape[1]
+    trimmed_generated_ids = generated_ids[:, prompt_length:]
+    output = (
+        processor.batch_decode(
+            trimmed_generated_ids,
+            skip_special_tokens=False,
+        )[0]
+        .strip()
+        .replace("<end_of_utterance>", "")
+    )
+    return {
+        "raw": output,
+        "segments": [
+            {
+                "metadata": {"page": kwargs.get("start", 0) + 1},
+                "content": output,
+            }
+        ],
+        "title": kwargs["title"],
+        "url": kwargs.get("url", ""),
+        "parent_title": kwargs.get("parent_title", ""),
+    }
 
 
 def parse_with_gemini(path: str, **kwargs) -> List[Dict] | str:
@@ -181,6 +251,27 @@ def parse_image_with_gemini(
     }
 
 
+def convert_path_to_images(path):
+    mime_type, _ = mimetypes.guess_type(path)
+    if mime_type and mime_type.startswith("image"):
+        # Single image processing
+        with open(path, "rb") as img_file:
+            image_base64 = base64.b64encode(img_file.read()).decode("utf-8")
+            return [(0, f"data:{mime_type};base64,{image_base64}")]
+    elif mime_type and mime_type.startswith("application/pdf"):
+        # PDF processing
+        pdf_document = pdfium.PdfDocument(path)
+        return [
+            (
+                page_num,
+                f"data:image/png;base64,{convert_pdf_page_to_base64(pdf_document, page_num)}",
+            )
+            for page_num in range(len(pdf_document))
+        ]
+    else:
+        raise ValueError(f"Unsupported file type: {mime_type}")
+
+
 def convert_pdf_page_to_base64(
     pdf_document: pdfium.PdfDocument, page_number: int
 ) -> str:
@@ -376,6 +467,7 @@ def parse_with_api(path: str, api: str, **kwargs) -> List[Dict] | str:
             )
             for page_num in range(len(pdf_document))
         ]
+    images = convert_path_to_images(path)
 
     # Process each page/image
     all_results = []

lexoid/core/utils.py

@@ -4,7 +4,6 @@
 import os
 import re
 import sys
-from difflib import SequenceMatcher
 from hashlib import md5
 from typing import Dict, List, Optional
 from urllib.parse import urlparse
@@ -16,20 +15,14 @@
 from bs4 import BeautifulSoup
 from docx2pdf import convert
 from loguru import logger
-from markdown import markdown
 from markdownify import markdownify as md
-from sklearn.feature_extraction.text import TfidfVectorizer
-from sklearn.metrics.pairwise import cosine_similarity
 from PIL import Image
 from PyQt5.QtCore import QMarginsF, QUrl
 from PyQt5.QtGui import QPageLayout, QPageSize
 from PyQt5.QtPrintSupport import QPrinter
 from PyQt5.QtWebEngineWidgets import QWebEngineView
 from PyQt5.QtWidgets import QApplication
 
-# Source: https://stackoverflow.com/a/12982689
-HTML_TAG_PATTERN = re.compile("<.*?>|&([a-z0-9]+|#[0-9]{1,6}|#x[0-9a-f]{1,6});")
-
 
 def split_pdf(input_path: str, output_dir: str, pages_per_split: int):
     paths = []
@@ -69,110 +62,6 @@ def convert_image_to_pdf(image_path: str) -> bytes:
         return pdf_buffer.getvalue()
 
 
-def remove_html_tags(text: str):
-    html = markdown(text, extensions=["tables"])
-    return re.sub(HTML_TAG_PATTERN, " ", html)
-
-
-def clean_text(txt):
-    # Remove LaTeX commands (e.g. \command, \command[args]{args})
-    txt = re.sub(r"\\[a-zA-Z]+(\[[^\]]*\])?(\{[^}]*\})?", " ", txt)
-
-    # Replace all blocks of whitespace (including tabs and newlines) with a single space
-    txt = re.sub(r"\s+", " ", txt)
-
-    # Remove all non-alphanumeric characters except spaces
-    txt = re.sub(r"[^a-zA-Z0-9 ]", " ", txt)
-
-    return txt.strip()
-
-
-def cosine_text_similarity(text1, text2):
-    texts = [clean_text(text1), clean_text(text2)]
-    vectorizer = TfidfVectorizer().fit_transform(texts)
-    return cosine_similarity(vectorizer[0], vectorizer[1])[0][0]
-
-
-def jaccard_text_similarity(text1: str, text2: str) -> float:
-    set1 = set(clean_text(text1).split())
-    set2 = set(clean_text(text2).split())
-    intersection = len(set1.intersection(set2))
-    union = len(set1.union(set2))
-    return intersection / union if union > 0 else 0.0
-
-
-def precision_recall_f1_score(text1: str, text2: str) -> Dict[str, float]:
-    """
-    Calculate precision, recall, and F1 score between two texts.
-    Precision = TP / (TP + FP)
-    Recall = TP / (TP + FN)
-    F1 Score = 2 * (Precision * Recall) / (Precision + Recall)
-    """
-    set1 = set(clean_text(text1).split())
-    set2 = set(clean_text(text2).split())
-
-    true_positive = len(set1.intersection(set2))
-    false_positive = len(set2 - set1)
-    false_negative = len(set1 - set2)
-
-    precision = (
-        true_positive / (true_positive + false_positive)
-        if (true_positive + false_positive) > 0
-        else 0.0
-    )
-    recall = (
-        true_positive / (true_positive + false_negative)
-        if (true_positive + false_negative) > 0
-        else 0.0
-    )
-    f1_score = (
-        2 * (precision * recall) / (precision + recall)
-        if (precision + recall) > 0
-        else 0.0
-    )
-
-    return {
-        "precision": precision,
-        "recall": recall,
-        "f1_score": f1_score,
-    }
-
-
-def calculate_similarities(
-    text1: str,
-    text2: str,
-    ignore_html: bool = True,
-    diff_save_path: str = "",
-) -> float:
-    """Calculate similarity ratio between two texts using SequenceMatcher."""
-    if ignore_html:
-        text1 = remove_html_tags(text1)
-        text2 = remove_html_tags(text2)
-
-    text1 = clean_text(clean_text(text1))
-    text2 = clean_text(clean_text(text2))
-
-    similarities = {}
-
-    sm = SequenceMatcher(None, text1, text2)
-    similarities["sequence_matcher"] = sm.ratio()
-    # Save the diff and the texts for debugging
-    if diff_save_path:
-        with open(diff_save_path, "w") as f:
-            f.write(f"Text 1:\n{text1}\n\n")
-            f.write(f"Text 2:\n{text2}\n\n")
-            f.write("Differences:\n")
-            for tag, i1, i2, j1, j2 in sm.get_opcodes():
-                if tag == "equal":
-                    continue
-                f.write(f"{tag} {text1[i1:i2]} -> {text2[j1:j2]}\n")
-    similarities["cosine"] = cosine_text_similarity(text1, text2)
-    similarities["jaccard"] = jaccard_text_similarity(text1, text2)
-    similarities.update(precision_recall_f1_score(text1, text2))
-
-    return similarities
-
-
 def convert_pdf_page_to_image(
     pdf_document: pypdfium2.PdfDocument, page_number: int
 ) -> bytes: