Suya OCR API — vLLM-backed, OpenAI-compatible OCR service

FastAPI service wrapping the Surya-OCR-2 model (datalab-to) served through vLLM:
legacy /v1/api/ai/* endpoints, an OpenAI-compatible /v1/chat/completions endpoint,
a coalescing request batcher, a local OCR CLI, Docker packaging, multilingual
example outputs, and quantization/concurrency benchmarks.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

surya/common/__init__.py

@@ -0,0 +1,3 @@
+
+
+

surya/common/blank.py

@@ -0,0 +1,64 @@
+"""Pixel-content heuristics for detecting blank or near-uniform image regions.
+
+Used by both the layout predictor (drop hallucinated layout blocks over empty
+space) and the recognition predictor (drop hallucinated text blocks from
+full-page OCR, decide whether an empty full-page output is a correct blank-page
+read or a failure).
+
+Two signals, combined:
+  * near-white fraction — most pixels have every RGB channel above a threshold
+  * pixel-value standard deviation — the region is essentially one color
+    (catches uniform-color fills that the white check misses)
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from PIL import Image
+
+
+# Per-channel value at/above which a pixel is considered "near-white".
+# Tolerates the small noise typical of PDF renders at 96 DPI.
+BLANK_WHITE_THRESHOLD = 245
+# Fraction of pixels that must be near-white for a region to count as blank.
+BLANK_PIXEL_FRACTION = 0.99
+# Pixel-value std below which a region is "essentially one color" regardless
+# of what that color is (catches solid-fill rectangles, dark banners, etc.).
+UNIFORM_COLOR_STD = 8.0
+
+
+def near_white_fraction(
+    image: Image.Image, white_threshold: int = BLANK_WHITE_THRESHOLD
+) -> float:
+    """Fraction of pixels where every RGB channel ≥ ``white_threshold``."""
+    arr = np.asarray(image.convert("RGB"))
+    if arr.size == 0:
+        return 0.0
+    return float(np.all(arr >= white_threshold, axis=-1).mean())
+
+
+def is_blank_region(
+    image: Image.Image,
+    *,
+    white_threshold: int = BLANK_WHITE_THRESHOLD,
+    blank_pixel_fraction: float = BLANK_PIXEL_FRACTION,
+    uniform_color_std: float = UNIFORM_COLOR_STD,
+) -> bool:
+    """True iff the image is essentially blank — either mostly near-white or
+    near-uniform color. Use this on a per-block crop or a whole page.
+
+    Returns False for empty (0-pixel) crops so callers don't accidentally
+    treat a degenerate bbox as blank.
+    """
+    arr = np.asarray(image.convert("RGB"))
+    if arr.size == 0:
+        return False
+    if np.all(arr >= white_threshold, axis=-1).mean() > blank_pixel_fraction:
+        return True
+    # Per-channel std — a uniform solid color (e.g., red banner with RGB=(200,50,50))
+    # has each channel constant across pixels, but mixing channels inflates the
+    # aggregate std. Check each channel independently.
+    per_channel_std = arr.reshape(-1, arr.shape[-1]).std(axis=0)
+    if float(per_channel_std.max()) < uniform_color_std:
+        return True
+    return False

surya/common/load.py

@@ -0,0 +1,25 @@
+from typing import Optional, Any
+
+import torch
+
+from surya.settings import settings
+
+
+class ModelLoader:
+    def __init__(self, checkpoint: Optional[str] = None):
+        self.checkpoint = checkpoint
+
+    def model(
+        self,
+        device: torch.device | str | None = settings.TORCH_DEVICE_MODEL,
+        dtype: Optional[torch.dtype | str] = settings.MODEL_DTYPE,
+        attention_implementation: Optional[str] = None,
+    ) -> Any:
+        raise NotImplementedError()
+
+    def processor(
+        self,
+        device: torch.device | str | None = settings.TORCH_DEVICE_MODEL,
+        dtype: Optional[torch.dtype | str] = settings.MODEL_DTYPE,
+    ) -> Any:
+        raise NotImplementedError()

surya/common/polygon.py

@@ -0,0 +1,122 @@
+import copy
+from typing import List, Optional
+
+import numpy as np
+from pydantic import BaseModel, field_validator, computed_field
+import numbers
+
+
+class PolygonBox(BaseModel):
+    polygon: List[List[float]]
+    confidence: Optional[float] = None
+
+    @field_validator("polygon", mode="before")
+    @classmethod
+    def convert_bbox_to_polygon(cls, value):
+        if isinstance(value, (list, tuple)) and len(value) == 4:
+            if all(isinstance(x, numbers.Number) for x in value):
+                value = [float(v) for v in value]
+                x_min, y_min, x_max, y_max = value
+                polygon = [
+                    [x_min, y_min],
+                    [x_max, y_min],
+                    [x_max, y_max],
+                    [x_min, y_max],
+                ]
+                return polygon
+            elif all(
+                isinstance(point, (list, tuple)) and len(point) == 2 for point in value
+            ):
+                value = [[float(v) for v in point] for point in value]
+                return value
+        elif isinstance(value, np.ndarray):
+            if value.shape == (4, 2):
+                return value.tolist()
+
+        raise ValueError(
+            f"Input must be either a bbox [x_min, y_min, x_max, y_max] or a polygon with 4 corners [(x,y), (x,y), (x,y), (x,y)].  All values must be numeric. You passed {value} of type {type(value)}.  The first value is of type {type(value[0])}."
+        )
+
+    @property
+    def height(self):
+        return self.bbox[3] - self.bbox[1]
+
+    @property
+    def width(self):
+        return self.bbox[2] - self.bbox[0]
+
+    @property
+    def area(self):
+        return self.width * self.height
+
+    @computed_field
+    @property
+    def bbox(self) -> List[float]:
+        x_coords = [point[0] for point in self.polygon]
+        y_coords = [point[1] for point in self.polygon]
+        return [min(x_coords), min(y_coords), max(x_coords), max(y_coords)]
+
+    def rescale(self, processor_size, image_size):
+        # Point is in x, y format
+        page_width, page_height = processor_size
+
+        img_width, img_height = image_size
+        width_scaler = img_width / page_width
+        height_scaler = img_height / page_height
+
+        for corner in self.polygon:
+            corner[0] = int(corner[0] * width_scaler)
+            corner[1] = int(corner[1] * height_scaler)
+
+    def round(self, divisor):
+        for corner in self.polygon:
+            corner[0] = int(corner[0] / divisor) * divisor
+            corner[1] = int(corner[1] / divisor) * divisor
+
+    def fit_to_bounds(self, bounds):
+        new_corners = copy.deepcopy(self.polygon)
+        for corner in new_corners:
+            corner[0] = max(min(corner[0], bounds[2]), bounds[0])
+            corner[1] = max(min(corner[1], bounds[3]), bounds[1])
+        self.polygon = new_corners
+
+    def expand(self, x_margin: float, y_margin: float):
+        new_polygon = []
+        x_margin = x_margin * self.width
+        y_margin = y_margin * self.height
+        for idx, poly in enumerate(self.polygon):
+            if idx == 0:
+                new_polygon.append([int(poly[0] - x_margin), int(poly[1] - y_margin)])
+            elif idx == 1:
+                new_polygon.append([int(poly[0] + x_margin), int(poly[1] - y_margin)])
+            elif idx == 2:
+                new_polygon.append([int(poly[0] + x_margin), int(poly[1] + y_margin)])
+            elif idx == 3:
+                new_polygon.append([int(poly[0] - x_margin), int(poly[1] + y_margin)])
+        self.polygon = new_polygon
+
+    def intersection_area(self, other, x_margin=0, y_margin=0):
+        x_overlap = self.x_overlap(other, x_margin)
+        y_overlap = self.y_overlap(other, y_margin)
+        return x_overlap * y_overlap
+
+    def x_overlap(self, other, x_margin=0):
+        return max(
+            0,
+            min(self.bbox[2] + x_margin, other.bbox[2] + x_margin)
+            - max(self.bbox[0] - x_margin, other.bbox[0] - x_margin),
+        )
+
+    def y_overlap(self, other, y_margin=0):
+        return max(
+            0,
+            min(self.bbox[3] + y_margin, other.bbox[3] + y_margin)
+            - max(self.bbox[1] - y_margin, other.bbox[1] - y_margin),
+        )
+
+    @property
+    def center(self):
+        return [(self.bbox[0] + self.bbox[2]) / 2, (self.bbox[1] + self.bbox[3]) / 2]
+
+    def __hash__(self):
+        return hash(tuple(self.bbox))

surya/common/predictor.py

@@ -0,0 +1,57 @@
+from typing import Optional
+
+import torch
+
+from surya.common.load import ModelLoader
+from surya.settings import settings
+
+
+class BasePredictor:
+    model_loader_cls = ModelLoader
+    batch_size: Optional[int] = None
+    default_batch_sizes = {"cpu": 1, "mps": 1, "cuda": 1}
+    torch_dtype = settings.MODEL_DTYPE
+
+    @property
+    def disable_tqdm(self) -> bool:
+        return self._disable_tqdm
+
+    @disable_tqdm.setter
+    def disable_tqdm(self, value: bool) -> None:
+        self._disable_tqdm = bool(value)
+
+    def __init__(
+        self,
+        checkpoint: Optional[str] = None,
+        device: torch.device | str | None = settings.TORCH_DEVICE_MODEL,
+        dtype: Optional[torch.dtype | str] = None,
+        attention_implementation: Optional[str] = None,
+    ):
+        if dtype is None:
+            dtype = self.torch_dtype
+
+        loader = self.model_loader_cls(checkpoint)
+        self.model = loader.model(device, dtype, attention_implementation)
+        self.processor = loader.processor()
+        self._disable_tqdm = settings.DISABLE_TQDM
+
+    def to(self, device_dtype: torch.device | str | None = None):
+        if hasattr(self, "model") and self.model:
+            self.model.to(device_dtype)
+            return
+        # Predictors that don't own a torch model (e.g. VLM-backed predictors that
+        # rely on an external server) treat .to() as a no-op.
+        if hasattr(self, "manager") and self.manager is not None:
+            return
+        raise ValueError("Model not loaded")
+
+    def get_batch_size(self):
+        batch_size = self.batch_size
+        if batch_size is None:
+            batch_size = self.default_batch_sizes["cpu"]
+            if settings.TORCH_DEVICE_MODEL in self.default_batch_sizes:
+                batch_size = self.default_batch_sizes[settings.TORCH_DEVICE_MODEL]
+        return batch_size
+
+    def __call__(self, *args, **kwargs):
+        raise NotImplementedError()

surya/common/pretrained.py

@@ -0,0 +1,22 @@
+from typing import Optional
+
+from transformers import PreTrainedModel
+from transformers.utils import is_flash_attn_2_available
+
+
+class SuryaPreTrainedModel(PreTrainedModel):
+    # No-op if we pass attention, so we can set attention however we want in the config
+    def _check_and_adjust_attn_implementation(
+        self, attn_implementation: Optional[str], **kwargs
+    ):
+        if attn_implementation is None:
+            try:
+                self._sdpa_can_dispatch(True)
+                attn_implementation = "sdpa"
+            except (ValueError, ImportError):
+                attn_implementation = "eager"
+
+            if self._supports_flash_attn and is_flash_attn_2_available():
+                attn_implementation = "flash_attention_2"
+
+        return attn_implementation

surya/common/s3.py

@@ -0,0 +1,182 @@
+import json
+import os
+import shutil
+import tempfile
+import time
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import requests
+from tqdm import tqdm
+
+from surya.logging import get_logger
+from surya.settings import settings
+
+logger = get_logger()
+
+# Lock file expiration time in seconds (10 minutes)
+LOCK_EXPIRATION = 600
+
+
+def join_urls(url1: str, url2: str):
+    url1 = url1.rstrip("/")
+    url2 = url2.lstrip("/")
+    return f"{url1}/{url2}"
+
+
+def get_model_name(pretrained_model_name_or_path: str):
+    return pretrained_model_name_or_path.split("/")[0]
+
+
+def download_file(remote_path: str, local_path: str, chunk_size: int = 1024 * 1024):
+    local_path = Path(local_path)
+    try:
+        response = requests.get(remote_path, stream=True, allow_redirects=True)
+        response.raise_for_status()  # Raise an exception for bad status codes
+
+        # Get file size from headers for progress bar
+        total_size = int(response.headers.get('content-length', 0))
+        
+        # Create progress bar with file name and size info
+        filename = local_path.name
+        pbar = tqdm(
+            total=total_size,
+            unit='B',
+            unit_scale=True,
+            unit_divisor=1024,
+            desc=f"Downloading {filename}",
+            miniters=1
+        )
+
+        with open(local_path, "wb") as f:
+            downloaded = 0
+            for chunk in response.iter_content(chunk_size=chunk_size):
+                if chunk:
+                    f.write(chunk)
+                    downloaded += len(chunk)
+                    pbar.update(len(chunk))
+        
+        pbar.close()
+        return local_path
+    except Exception as e:
+        if local_path.exists():
+            local_path.unlink()
+        logger.error(f"Download error for file {remote_path}: {str(e)}")
+        raise
+
+
+def check_manifest(local_dir: str):
+    local_dir = Path(local_dir)
+    manifest_path = local_dir / "manifest.json"
+    if not os.path.exists(manifest_path):
+        return False
+
+    try:
+        with open(manifest_path, "r") as f:
+            manifest = json.load(f)
+        for file in manifest["files"]:
+            if not os.path.exists(local_dir / file):
+                return False
+    except Exception:
+        return False
+
+    return True
+
+
+def download_directory(remote_path: str, local_dir: str):
+    model_name = get_model_name(remote_path)
+    s3_url = join_urls(settings.S3_BASE_URL, remote_path)
+    # Check to see if it's already downloaded
+    model_exists = check_manifest(local_dir)
+    if model_exists:
+        return
+
+    # Use tempfile.TemporaryDirectory to automatically clean up
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Download the manifest file
+        manifest_file = join_urls(s3_url, "manifest.json")
+        manifest_path = os.path.join(temp_dir, "manifest.json")
+        download_file(manifest_file, manifest_path)
+
+        # List and download all files
+        with open(manifest_path, "r") as f:
+            manifest = json.load(f)
+
+        pbar = tqdm(
+            desc=f"Downloading {model_name} model to {local_dir}",
+            total=len(manifest["files"]),
+        )
+
+        with ThreadPoolExecutor(
+            max_workers=settings.PARALLEL_DOWNLOAD_WORKERS
+        ) as executor:
+            futures = []
+            for file in manifest["files"]:
+                remote_file = join_urls(s3_url, file)
+                local_file = os.path.join(temp_dir, file)
+                futures.append(executor.submit(download_file, remote_file, local_file))
+
+            for future in futures:
+                future.result()
+                pbar.update(1)
+
+        pbar.close()
+
+        # Move all files to new directory
+        for file in os.listdir(temp_dir):
+            shutil.move(os.path.join(temp_dir, file), local_dir)
+
+
+class S3DownloaderMixin:
+    s3_prefix = "s3://"
+
+    @classmethod
+    def get_local_path(cls, pretrained_model_name_or_path) -> str:
+        if pretrained_model_name_or_path.startswith(cls.s3_prefix):
+            pretrained_model_name_or_path = pretrained_model_name_or_path.replace(
+                cls.s3_prefix, ""
+            )
+            cache_dir = settings.MODEL_CACHE_DIR
+            local_path = os.path.join(cache_dir, pretrained_model_name_or_path)
+            os.makedirs(local_path, exist_ok=True)
+        else:
+            local_path = ""
+        return local_path
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        # Allow loading models directly from the hub, or using s3
+        if not pretrained_model_name_or_path.startswith(cls.s3_prefix):
+            return super().from_pretrained(
+                pretrained_model_name_or_path, *args, **kwargs
+            )
+
+        local_path = cls.get_local_path(pretrained_model_name_or_path)
+        pretrained_model_name_or_path = pretrained_model_name_or_path.replace(
+            cls.s3_prefix, ""
+        )
+
+        # Retry logic for downloading the model folder
+        retries = 3
+        delay = 5
+        attempt = 0
+        success = False
+        while not success and attempt < retries:
+            try:
+                download_directory(pretrained_model_name_or_path, local_path)
+                success = True  # If download succeeded
+            except Exception as e:
+                logger.error(
+                    f"Error downloading model from {pretrained_model_name_or_path}. Attempt {attempt + 1} of {retries}. Error: {e}"
+                )
+                attempt += 1
+                if attempt < retries:
+                    logger.info(f"Retrying in {delay} seconds...")
+                    time.sleep(delay)  # Wait before retrying
+                else:
+                    logger.error(
+                        f"Failed to download {pretrained_model_name_or_path} after {retries} attempts."
+                    )
+                    raise e  # Reraise exception after max retries
+
+        return super().from_pretrained(local_path, *args, **kwargs)

surya/common/util.py

@@ -0,0 +1,44 @@
+from typing import List
+
+from surya.common.polygon import PolygonBox
+
+
+def clean_boxes(boxes: List[PolygonBox]) -> List[PolygonBox]:
+    new_boxes = []
+    for box_obj in boxes:
+        xs = [point[0] for point in box_obj.polygon]
+        ys = [point[1] for point in box_obj.polygon]
+        if max(xs) == min(xs) or max(ys) == min(ys):
+            continue
+
+        box = box_obj.bbox
+        contained = False
+        for other_box_obj in boxes:
+            if other_box_obj.polygon == box_obj.polygon:
+                continue
+
+            other_box = other_box_obj.bbox
+            if box == other_box:
+                continue
+            if (
+                box[0] >= other_box[0]
+                and box[1] >= other_box[1]
+                and box[2] <= other_box[2]
+                and box[3] <= other_box[3]
+            ):
+                contained = True
+                break
+        if not contained:
+            new_boxes.append(box_obj)
+    return new_boxes
+
+
+def expand_bbox(bbox, expansion_factor=0.01):
+    expansion_low = 1 - expansion_factor
+    expansion_high = 1 + expansion_factor
+    return [
+        bbox[0] * expansion_low,
+        bbox[1] * expansion_low,
+        bbox[2] * expansion_high,
+        bbox[3] * expansion_high,
+    ]

surya/debug/draw.py

@@ -0,0 +1,66 @@
+from PIL import ImageDraw, ImageFont
+
+from surya.debug.fonts import get_font_path
+from surya.debug.text import get_text_size
+
+
+def draw_bboxes_on_image(
+    bboxes, image, labels=None, label_font_size=10, color: str | list = "red"
+):
+    polys = []
+    for bb in bboxes:
+        # Clockwise polygon
+        poly = [[bb[0], bb[1]], [bb[2], bb[1]], [bb[2], bb[3]], [bb[0], bb[3]]]
+        polys.append(poly)
+
+    return draw_polys_on_image(
+        polys, image, labels, label_font_size=label_font_size, color=color
+    )
+
+
+def draw_polys_on_image(
+    corners,
+    image,
+    labels=None,
+    box_padding=-1,
+    label_offset=1,
+    label_font_size=10,
+    color: str | list = "red",
+):
+    draw = ImageDraw.Draw(image)
+    font_path = get_font_path()
+    label_font = ImageFont.truetype(font_path, label_font_size)
+
+    for i in range(len(corners)):
+        poly = corners[i]
+        poly = [(int(p[0]), int(p[1])) for p in poly]
+        draw.polygon(
+            poly, outline=color[i] if isinstance(color, list) else color, width=1
+        )
+
+        if labels is not None:
+            label = labels[i]
+            text_position = (
+                min([p[0] for p in poly]) + label_offset,
+                min([p[1] for p in poly]) + label_offset,
+            )
+            text_size = get_text_size(label, label_font)
+            box_position = (
+                text_position[0] - box_padding + label_offset,
+                text_position[1] - box_padding + label_offset,
+                text_position[0] + text_size[0] + box_padding + label_offset,
+                text_position[1] + text_size[1] + box_padding + label_offset,
+            )
+            try:
+                draw.rectangle(box_position, fill="white")
+            except Exception as e:
+                print(f"Error drawing rectangle at {box_position}: {e}")
+                continue
+            draw.text(
+                text_position,
+                label,
+                fill=color[i] if isinstance(color, list) else color,
+                font=label_font,
+            )
+
+    return image

surya/debug/fonts.py

@@ -0,0 +1,24 @@
+from typing import List, Optional
+import os
+import requests
+
+from surya.settings import settings
+
+
+def get_font_path(langs: Optional[List[str]] = None) -> str:
+    font_path = settings.RECOGNITION_RENDER_FONTS["all"]
+    if langs is not None:
+        for k in settings.RECOGNITION_RENDER_FONTS:
+            if k in langs and len(langs) == 1:
+                font_path = settings.RECOGNITION_RENDER_FONTS[k]
+                break
+
+    if not os.path.exists(font_path):
+        os.makedirs(os.path.dirname(font_path), exist_ok=True)
+        font_dl_path = f"{settings.RECOGNITION_FONT_DL_BASE}/{os.path.basename(font_path)}"
+        with requests.get(font_dl_path, stream=True) as r, open(font_path, 'wb') as f:
+            r.raise_for_status()
+            for chunk in r.iter_content(chunk_size=8192):
+                f.write(chunk)
+
+    return font_path

surya/debug/katex.js

@@ -0,0 +1,64 @@
+<style>
+    .katex-display-container {
+        display: inline-block;
+        max-width: 100%;
+        overflow-x: auto;
+        max-height: 100%;
+    }
+
+    .katex-inline-container {
+        display: inline-block;
+        max-width: 100%;
+        overflow-x: auto;
+        max-height: 100%;
+    }
+</style>
+<script src="https://cdn.jsdelivr.net/npm/katex@0.16.21/dist/katex.min.js" onload="setTimeout(function() {renderMath()})" async></script>
+<link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.21/dist/katex.min.css">
+<script>
+    function htmlUnescape(escapedText) {
+      const htmlEntities = {
+        '&amp;': '&',
+        '&lt;': '<',
+        '&gt;': '>',
+        '&quot;': '"',
+        '&#39;': "'",
+        '&nbsp;': ' '
+      };
+
+      return escapedText.replace(/&amp;|&lt;|&gt;|&quot;|&#39;|&nbsp;/g, match => htmlEntities[match]);
+    }
+
+    const renderMath = (function() {
+    try {
+       const mathElements = document.querySelectorAll('math');
+
+        mathElements.forEach(function(element) {
+          let mathContent = element.innerHTML.trim();
+          mathContent = htmlUnescape(mathContent);
+          const isDisplay = element.getAttribute('display') === 'block';
+
+          const container = document.createElement('span');
+          container.className = isDisplay ? 'katex-display-container' : 'katex-inline-container';
+          element.parentNode.insertBefore(container, element);
+
+          try {
+            katex.render(mathContent, container, {
+              displayMode: isDisplay,
+              throwOnError: false
+            });
+
+          } catch (err) {
+            console.error('KaTeX rendering error:', err);
+            container.textContent = mathContent; // Fallback to raw text
+          }
+
+          element.parentNode.removeChild(element);
+        });
+
+        console.log('Math rendering complete with', mathElements.length, 'expressions');
+      } catch (err) {
+        console.error('Error in renderMath function:', err);
+      }
+    });
+</script>

surya/debug/render_html.py

@@ -0,0 +1,90 @@
+import html as htmllib
+import os.path
+import re
+
+filepath = os.path.abspath(__file__)
+
+def render_text_as_html(
+        bboxes: list[list[int]],
+        texts: list[str],
+        image_size: tuple[int, int],
+        base_font_size: int = 16,
+        scaler: int = 2
+):
+    katex_path = os.path.join(os.path.dirname(filepath), "katex.js")
+    with open(katex_path, "r") as f:
+        katex_script = f.read()
+
+    html_content = []
+    image_size = tuple([int(s * scaler) for s in image_size])
+    width, height = image_size
+
+
+    html_content.append(f"""
+<!DOCTYPE html>
+<html>
+<head>
+    <style>
+        body {{
+            margin: 0;
+            padding: 0;
+            width: {width}px;
+            height: {height}px;
+            position: relative;
+            overflow: hidden;
+            background: white;
+            color: black;
+        }}
+        .text-box {{
+            position: absolute;
+            overflow: hidden;
+            display: flex;
+            justify-content: left;
+            font-family: Arial, sans-serif;
+            white-space: pre-wrap;
+        }}
+        .vertical-text {{
+          writing-mode: vertical-rl;  /* Top to bottom, right to left */
+        }}
+    </style>
+    {katex_script}
+</head>
+<body>
+""")
+
+    for i, (bbox, text) in enumerate(zip(bboxes, texts)):
+        bbox = bbox.copy()
+        bbox = [int(bb * scaler) for bb in bbox]
+        x1, y1, x2, y2 = bbox
+        width = x2 - x1
+        height = y2 - y1
+        min_dim = min(width, height)
+
+        # Scale font size based on box height
+        font_size = min(int(min_dim * 0.75), base_font_size)
+
+        # Create div with absolute positioning
+        div_style = (
+            f"left: {x1}px; "
+            f"top: {y1}px; "
+            f"width: {width}px; "
+            f"height: {height}px; "
+            f"font-size: {font_size}px;"
+        )
+
+        class_ = "text-box"
+        if height > width * 2:
+            class_ += " vertical-text"
+
+        # Determine if content is HTML/MathML or plain text
+        if "<" in text and ">" in text and re.search(r"<(html|math|div|sub|sup|i|u|mark|small|del|b|br|code)\b", text.lower()):
+            # Content is already HTML/MathML, include as-is
+            html_content.append(f'<span class="{class_}" id="box-{i}" style="{div_style}">{text}</span>')
+        else:
+            # Plain text, escape it
+            escaped_text = htmllib.escape(text)
+            html_content.append(f'<span class="{class_}" id="box-{i}" style="{div_style}">{escaped_text}</span>')
+
+    html_content.append("</body></html>")
+
+    return "\n".join(html_content), image_size

surya/debug/text.py

@@ -0,0 +1,8 @@
+from PIL import Image, ImageDraw
+
+
+def get_text_size(text, font):
+    im = Image.new(mode="P", size=(0, 0))
+    draw = ImageDraw.Draw(im)
+    _, _, width, height = draw.textbbox((0, 0), text=text, font=font)
+    return width, height

surya/detection/__init__.py

@@ -0,0 +1,147 @@
+from concurrent.futures import ThreadPoolExecutor
+from typing import List, Generator, Tuple
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from PIL import Image
+from tqdm import tqdm
+
+from surya.common.predictor import BasePredictor
+
+from surya.detection.loader import DetectionModelLoader
+from surya.detection.parallel import FakeExecutor
+from surya.detection.util import get_total_splits, split_image
+from surya.detection.schema import TextDetectionResult
+from surya.settings import settings
+from surya.detection.heatmap import parallel_get_boxes
+
+
+class DetectionPredictor(BasePredictor):
+    model_loader_cls = DetectionModelLoader
+    batch_size = settings.DETECTOR_BATCH_SIZE
+    default_batch_sizes = {"cpu": 8, "mps": 8, "cuda": 36}
+
+    def __call__(
+        self, images: List[Image.Image], batch_size=None, include_maps=False
+    ) -> List[TextDetectionResult]:
+        detection_generator = self.batch_detection(images, batch_size=batch_size)
+
+        postprocessing_futures = []
+        max_workers = min(settings.DETECTOR_POSTPROCESSING_CPU_WORKERS, len(images))
+        parallelize = (
+            not settings.IN_STREAMLIT
+            and len(images) >= settings.DETECTOR_MIN_PARALLEL_THRESH
+        )
+        executor = ThreadPoolExecutor if parallelize else FakeExecutor
+        with executor(max_workers=max_workers) as e:
+            for preds, orig_sizes in detection_generator:
+                for pred, orig_size in zip(preds, orig_sizes):
+                    postprocessing_futures.append(
+                        e.submit(parallel_get_boxes, pred, orig_size, include_maps)
+                    )
+
+        return [future.result() for future in postprocessing_futures]
+
+    def prepare_image(self, img):
+        new_size = (self.processor.size["width"], self.processor.size["height"])
+
+        # This double resize actually necessary for downstream accuracy
+        img.thumbnail(new_size, Image.Resampling.LANCZOS)
+        img = img.resize(
+            new_size, Image.Resampling.LANCZOS
+        )  # Stretch smaller dimension to fit new size
+
+        img = np.asarray(img, dtype=np.uint8)
+        img = self.processor(img)["pixel_values"][0]
+        img = torch.from_numpy(img)
+        return img
+
+    def batch_detection(
+        self, images: List, batch_size=None
+    ) -> Generator[Tuple[List[List[np.ndarray]], List[Tuple[int, int]]], None, None]:
+        assert all([isinstance(image, Image.Image) for image in images])
+        if batch_size is None:
+            batch_size = self.get_batch_size()
+        heatmap_count = self.model.config.num_labels
+
+        orig_sizes = [image.size for image in images]
+        splits_per_image = [
+            get_total_splits(size, self.processor.size["height"]) for size in orig_sizes
+        ]
+
+        batches = []
+        current_batch_size = 0
+        current_batch = []
+        for i in range(len(images)):
+            if current_batch_size + splits_per_image[i] > batch_size:
+                if len(current_batch) > 0:
+                    batches.append(current_batch)
+                current_batch = []
+                current_batch_size = 0
+            current_batch.append(i)
+            current_batch_size += splits_per_image[i]
+
+        if len(current_batch) > 0:
+            batches.append(current_batch)
+
+        for batch_idx in tqdm(
+            range(len(batches)), desc="Detecting bboxes", disable=self.disable_tqdm
+        ):
+            batch_image_idxs = batches[batch_idx]
+            batch_images = [images[j].convert("RGB") for j in batch_image_idxs]
+
+            split_index = []
+            split_heights = []
+            image_splits = []
+            for image_idx, image in enumerate(batch_images):
+                image_parts, split_height = split_image(
+                    image, self.processor.size["height"]
+                )
+                image_splits.extend(image_parts)
+                split_index.extend([image_idx] * len(image_parts))
+                split_heights.extend(split_height)
+
+            image_splits = [self.prepare_image(image) for image in image_splits]
+            # Batch images in dim 0
+            batch = torch.stack(image_splits, dim=0).to(self.model.dtype)
+
+            with settings.INFERENCE_MODE():
+                pred = self.model(pixel_values=batch.to(self.model.device))
+
+            logits = pred.logits
+            correct_shape = [
+                self.processor.size["height"],
+                self.processor.size["width"],
+            ]
+            current_shape = list(logits.shape[2:])
+            if current_shape != correct_shape:
+                logits = F.interpolate(
+                    logits, size=correct_shape, mode="bilinear", align_corners=False
+                )
+
+            logits = logits.to(torch.float32).cpu().numpy()
+            preds = []
+            for i, (idx, height) in enumerate(zip(split_index, split_heights)):
+                # If our current prediction length is below the image idx, that means we have a new image
+                # Otherwise, we need to add to the current image
+                if len(preds) <= idx:
+                    preds.append([logits[i][k] for k in range(heatmap_count)])
+                else:
+                    heatmaps = preds[idx]
+                    pred_heatmaps = [logits[i][k] for k in range(heatmap_count)]
+
+                    if height < self.processor.size["height"]:
+                        # Cut off padding to get original height
+                        pred_heatmaps = [
+                            pred_heatmap[:height, :] for pred_heatmap in pred_heatmaps
+                        ]
+
+                    for k in range(heatmap_count):
+                        heatmaps[k] = np.vstack([heatmaps[k], pred_heatmaps[k]])
+                    preds[idx] = heatmaps
+
+            yield preds, [orig_sizes[j] for j in batch_image_idxs]
+
+        torch.cuda.empty_cache()

surya/detection/heatmap.py

@@ -0,0 +1,165 @@
+from typing import List
+
+import cv2
+import numpy as np
+from PIL import Image
+
+from surya.common.util import clean_boxes
+from surya.detection import TextDetectionResult
+from surya.common.polygon import PolygonBox
+from surya.settings import settings
+
+
+def get_dynamic_thresholds(linemap, text_threshold, low_text, typical_top10_avg=0.7):
+    # Find average intensity of top 10% pixels
+    flat_map = linemap.ravel()
+    top_10_count = int(len(flat_map) * 0.9)
+    avg_intensity = np.mean(np.partition(flat_map, top_10_count)[top_10_count:])
+    scaling_factor = np.clip(avg_intensity / typical_top10_avg, 0, 1) ** (1 / 2)
+
+    low_text = np.clip(low_text * scaling_factor, 0.1, 0.6)
+    text_threshold = np.clip(text_threshold * scaling_factor, 0.15, 0.8)
+
+    return text_threshold, low_text
+
+
+def detect_boxes(linemap, text_threshold, low_text):
+    # From CRAFT - https://github.com/clovaai/CRAFT-pytorch
+    # Modified to return boxes and for speed, accuracy
+    img_h, img_w = linemap.shape
+
+    text_threshold, low_text = get_dynamic_thresholds(linemap, text_threshold, low_text)
+
+    text_score_comb = (linemap > low_text).astype(np.uint8)
+    label_count, labels, stats, centroids = cv2.connectedComponentsWithStats(
+        text_score_comb, connectivity=4
+    )
+
+    det = []
+    confidences = []
+    max_confidence = 0
+
+    for k in range(1, label_count):
+        # size filtering
+        size = stats[k, cv2.CC_STAT_AREA]
+        if size < 10:
+            continue
+
+        # make segmentation map
+        x, y, w, h = stats[
+            k,
+            [cv2.CC_STAT_LEFT, cv2.CC_STAT_TOP, cv2.CC_STAT_WIDTH, cv2.CC_STAT_HEIGHT],
+        ]
+
+        try:
+            niter = int(np.sqrt(min(w, h)))
+        except ValueError:
+            niter = 0
+
+        buffer = 1
+        sx, sy = max(0, x - niter - buffer), max(0, y - niter - buffer)
+        ex, ey = min(img_w, x + w + niter + buffer), min(img_h, y + h + niter + buffer)
+
+        mask = labels[sy:ey, sx:ex] == k
+        selected_linemap = linemap[sy:ey, sx:ex][mask]
+        if selected_linemap.size == 0:
+            continue
+
+        line_max = np.max(selected_linemap)
+
+        # thresholding
+        if line_max < text_threshold:
+            continue
+
+        segmap = mask.astype(np.uint8)
+
+        ksize = buffer + niter
+        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (ksize, ksize))
+        selected_segmap = cv2.dilate(segmap, kernel)
+
+        # make box
+        y_inds, x_inds = np.nonzero(selected_segmap)
+        x_inds += sx
+        y_inds += sy
+        np_contours = np.column_stack((x_inds, y_inds))
+        rectangle = cv2.minAreaRect(np_contours)
+        box = cv2.boxPoints(rectangle)
+
+        # align diamond-shape
+        w, h = np.linalg.norm(box[0] - box[1]), np.linalg.norm(box[1] - box[2])
+        box_ratio = max(w, h) / (min(w, h) + 1e-5)
+        if abs(1 - box_ratio) <= 0.1:
+            left, right = np_contours[:, 0].min(), np_contours[:, 0].max()
+            top, bottom = np_contours[:, 1].min(), np_contours[:, 1].max()
+            box = np.array(
+                [[left, top], [right, top], [right, bottom], [left, bottom]],
+                dtype=np.float32,
+            )
+
+        # make clock-wise order
+        startidx = box.sum(axis=1).argmin()
+        box = np.roll(box, 4 - startidx, 0)
+
+        max_confidence = max(max_confidence, line_max)
+
+        confidences.append(line_max)
+        det.append(box)
+
+    if max_confidence > 0:
+        confidences = [c / max_confidence for c in confidences]
+    return det, confidences
+
+
+def get_detected_boxes(textmap, text_threshold=None, low_text=None) -> List[PolygonBox]:
+    if text_threshold is None:
+        text_threshold = settings.DETECTOR_TEXT_THRESHOLD
+    if low_text is None:
+        low_text = settings.DETECTOR_BLANK_THRESHOLD
+
+    if textmap.dtype != np.float32:
+        textmap = textmap.astype(np.float32)
+
+    boxes, confidences = detect_boxes(textmap, text_threshold, low_text)
+    # From point form to box form
+    return [
+        PolygonBox(polygon=box, confidence=confidence)
+        for box, confidence in zip(boxes, confidences)
+    ]
+
+
+def get_and_clean_boxes(
+    textmap, processor_size, image_size, text_threshold=None, low_text=None
+) -> List[PolygonBox]:
+    bboxes = get_detected_boxes(textmap, text_threshold, low_text)
+    for bbox in bboxes:
+        bbox.rescale(processor_size, image_size)
+        bbox.fit_to_bounds([0, 0, image_size[0], image_size[1]])
+
+    bboxes = clean_boxes(bboxes)
+    return bboxes
+
+
+def parallel_get_boxes(preds, orig_sizes, include_maps=False):
+    heatmap, affinity_map = preds
+    heat_img, aff_img = None, None
+
+    if include_maps:
+        heat_img = Image.fromarray((heatmap * 255).astype(np.uint8))
+        aff_img = Image.fromarray((affinity_map * 255).astype(np.uint8))
+    heatmap_size = list(reversed(heatmap.shape))
+    bboxes = get_and_clean_boxes(heatmap, heatmap_size, orig_sizes)
+    for box in bboxes:
+        # Skip for vertical boxes
+        if box.height < 3 * box.width:
+            box.expand(x_margin=0, y_margin=settings.DETECTOR_BOX_Y_EXPAND_MARGIN)
+            box.fit_to_bounds(
+                [0, 0, orig_sizes[0], orig_sizes[1]]
+            )  # Fix any bad expands
+
+    result = TextDetectionResult(
+        bboxes=bboxes,
+        heatmap=heat_img,
+        affinity_map=aff_img,
+        image_bbox=[0, 0, orig_sizes[0], orig_sizes[1]],
+    )
+    return result

surya/detection/loader.py

@@ -0,0 +1,53 @@
+from typing import Optional
+
+import torch
+
+from surya.common.load import ModelLoader
+from surya.detection.processor import SegformerImageProcessor
+
+from surya.detection.model.config import EfficientViTConfig
+from surya.detection.model.encoderdecoder import EfficientViTForSemanticSegmentation
+from surya.logging import get_logger
+from surya.settings import settings
+
+logger = get_logger()
+
+
+class DetectionModelLoader(ModelLoader):
+    def __init__(self, checkpoint: Optional[str] = None):
+        super().__init__(checkpoint)
+
+        if self.checkpoint is None:
+            self.checkpoint = settings.DETECTOR_MODEL_CHECKPOINT
+
+    def model(
+        self,
+        device: Optional[torch.device | str] = None,
+        dtype: Optional[torch.dtype | str] = None,
+        attention_implementation: Optional[str] = None,
+    ) -> EfficientViTForSemanticSegmentation:
+        if device is None:
+            device = settings.TORCH_DEVICE_MODEL
+        if dtype is None:
+            dtype = settings.MODEL_DTYPE
+
+        config = EfficientViTConfig.from_pretrained(self.checkpoint)
+        model = EfficientViTForSemanticSegmentation.from_pretrained(
+            self.checkpoint,
+            dtype=dtype,
+            config=config,
+        )
+        model = model.to(device)
+        model = model.eval()
+
+        logger.debug(
+            f"Loaded detection model {self.checkpoint} from {EfficientViTForSemanticSegmentation.get_local_path(self.checkpoint)} onto device {device} with dtype {dtype}"
+        )
+        return model
+
+    def processor(
+        self,
+        device: Optional[torch.device | str] = None,
+        dtype: Optional[torch.dtype | str] = None,
+    ) -> SegformerImageProcessor:
+        return SegformerImageProcessor.from_pretrained(self.checkpoint)

surya/detection/model/config.py

@@ -0,0 +1,53 @@
+from transformers import PretrainedConfig
+
+from surya.common.s3 import S3DownloaderMixin
+
+
+class EfficientViTConfig(S3DownloaderMixin, PretrainedConfig):
+    r"""
+    ```"""
+
+    model_type = "efficientvit"
+
+    def __init__(
+        self,
+        num_classes=2,
+        num_channels=3,
+        widths=(32, 64, 128, 256, 512),
+        head_dim=32,
+        num_stages=4,
+        depths=(1, 1, 1, 6, 6),
+        strides=(2, 2, 2, 2, 2),
+        hidden_sizes=(32, 64, 160, 256),
+        patch_size=(7, 7),
+        hidden_dropout_prob=0.0,
+        attention_probs_dropout_prob=0.0,
+        classifier_dropout_prob=0.0,
+        layer_norm_eps=1e-6,
+        decoder_layer_hidden_size=128,
+        decoder_hidden_size=512,
+        semantic_loss_ignore_index=255,
+        initializer_range=0.02,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.num_classes = num_classes
+        self.widths = widths
+        self.head_dim = head_dim
+
+        self.num_channels = num_channels
+        self.num_stages = num_stages
+        self.depths = depths
+        self.strides = strides
+        self.hidden_sizes = hidden_sizes
+        self.patch_size = patch_size
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.classifier_dropout_prob = classifier_dropout_prob
+        self.layer_norm_eps = layer_norm_eps
+        self.decoder_hidden_size = decoder_hidden_size
+        self.decoder_layer_hidden_size = decoder_layer_hidden_size
+        self.semantic_loss_ignore_index = semantic_loss_ignore_index
+
+        self.initializer_range = initializer_range

surya/detection/model/encoderdecoder.py

@@ -0,0 +1,839 @@
+"""
+This is an implementation of efficientvit, with some modifications (decode head, etc).
+
+Original paper at https://arxiv.org/abs/2205.14756
+
+Code adapted from timm, https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/efficientvit_mit.py
+Original code (that timm adapted from) at https://github.com/mit-han-lab/efficientvit
+
+License: Apache 2
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Union, Tuple, List, Any
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.modeling_outputs import SemanticSegmenterOutput
+
+from surya.common.pretrained import SuryaPreTrainedModel
+from surya.common.s3 import S3DownloaderMixin
+from surya.detection.model.config import EfficientViTConfig
+
+
+def val2list(x: Union[List, Tuple, Any], repeat_time=1):
+    if isinstance(x, (list, tuple)):
+        return list(x)
+    return [x for _ in range(repeat_time)]
+
+
+def val2tuple(x: Union[List, Tuple, Any], min_len: int = 1, idx_repeat: int = -1):
+    # repeat elements if necessary
+    x = val2list(x)
+    if len(x) > 0:
+        x[idx_repeat:idx_repeat] = [x[idx_repeat] for _ in range(min_len - len(x))]
+
+    return tuple(x)
+
+
+def get_same_padding(
+    kernel_size: Union[int, Tuple[int, ...]],
+) -> Union[int, Tuple[int, ...]]:
+    if isinstance(kernel_size, tuple):
+        return tuple([get_same_padding(ks) for ks in kernel_size])
+    else:
+        assert kernel_size % 2 > 0, "kernel size should be odd number"
+        return kernel_size // 2
+
+
+def get_padding(kernel_size: int, stride: int = 1, dilation: int = 1) -> int:
+    padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2
+    return padding
+
+
+class ConvNormAct(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size=3,
+        stride=1,
+        dilation=1,
+        groups=1,
+        bias=False,
+        dropout=0.0,
+        norm_layer=nn.BatchNorm2d,
+        act_layer=nn.ReLU,
+    ):
+        super(ConvNormAct, self).__init__()
+        self.dropout = nn.Dropout(dropout, inplace=False)
+        padding = get_padding(kernel_size, stride, dilation)
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+            padding=padding,
+        )
+        self.norm = (
+            norm_layer(num_features=out_channels) if norm_layer else nn.Identity()
+        )
+        self.act = act_layer(inplace=True) if act_layer is not None else nn.Identity()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        x = self.act(x)
+        return x
+
+
+class DSConv(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size=3,
+        stride=1,
+        use_bias=False,
+        norm_layer=(nn.BatchNorm2d, nn.BatchNorm2d),
+        act_layer=(nn.ReLU6, None),
+    ):
+        super(DSConv, self).__init__()
+        use_bias = val2tuple(use_bias, 2)
+        norm_layer = val2tuple(norm_layer, 2)
+        act_layer = val2tuple(act_layer, 2)
+
+        self.depth_conv = ConvNormAct(
+            in_channels,
+            in_channels,
+            kernel_size,
+            stride,
+            groups=in_channels,
+            norm_layer=norm_layer[0],
+            act_layer=act_layer[0],
+            bias=use_bias[0],
+        )
+        self.point_conv = ConvNormAct(
+            in_channels,
+            out_channels,
+            1,
+            norm_layer=norm_layer[1],
+            act_layer=act_layer[1],
+            bias=use_bias[1],
+        )
+
+    def forward(self, x):
+        x = self.depth_conv(x)
+        x = self.point_conv(x)
+        return x
+
+
+class ConvBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size=3,
+        stride=1,
+        mid_channels=None,
+        expand_ratio=1,
+        use_bias=False,
+        norm_layer=(nn.BatchNorm2d, nn.BatchNorm2d),
+        act_layer=(nn.ReLU6, None),
+    ):
+        super(ConvBlock, self).__init__()
+        use_bias = val2tuple(use_bias, 2)
+        norm_layer = val2tuple(norm_layer, 2)
+        act_layer = val2tuple(act_layer, 2)
+        mid_channels = mid_channels or round(in_channels * expand_ratio)
+
+        self.conv1 = ConvNormAct(
+            in_channels,
+            mid_channels,
+            kernel_size,
+            stride,
+            norm_layer=norm_layer[0],
+            act_layer=act_layer[0],
+            bias=use_bias[0],
+        )
+        self.conv2 = ConvNormAct(
+            mid_channels,
+            out_channels,
+            kernel_size,
+            1,
+            norm_layer=norm_layer[1],
+            act_layer=act_layer[1],
+            bias=use_bias[1],
+        )
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        return x
+
+
+class MBConv(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size=3,
+        stride=1,
+        mid_channels=None,
+        expand_ratio=6,
+        use_bias=False,
+        norm_layer=(nn.BatchNorm2d, nn.BatchNorm2d, nn.BatchNorm2d),
+        act_layer=(nn.ReLU6, nn.ReLU6, None),
+    ):
+        super(MBConv, self).__init__()
+        use_bias = val2tuple(use_bias, 3)
+        norm_layer = val2tuple(norm_layer, 3)
+        act_layer = val2tuple(act_layer, 3)
+        mid_channels = mid_channels or round(in_channels * expand_ratio)
+
+        self.inverted_conv = ConvNormAct(
+            in_channels,
+            mid_channels,
+            1,
+            stride=1,
+            norm_layer=norm_layer[0],
+            act_layer=act_layer[0],
+            bias=use_bias[0],
+        )
+        self.depth_conv = ConvNormAct(
+            mid_channels,
+            mid_channels,
+            kernel_size,
+            stride=stride,
+            groups=mid_channels,
+            norm_layer=norm_layer[1],
+            act_layer=act_layer[1],
+            bias=use_bias[1],
+        )
+        self.point_conv = ConvNormAct(
+            mid_channels,
+            out_channels,
+            1,
+            norm_layer=norm_layer[2],
+            act_layer=act_layer[2],
+            bias=use_bias[2],
+        )
+
+    def forward(self, x):
+        x = self.inverted_conv(x)
+        x = self.depth_conv(x)
+        x = self.point_conv(x)
+        return x
+
+
+class FusedMBConv(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size=3,
+        stride=1,
+        mid_channels=None,
+        expand_ratio=6,
+        groups=1,
+        use_bias=False,
+        norm_layer=(nn.BatchNorm2d, nn.BatchNorm2d),
+        act_layer=(nn.ReLU6, None),
+    ):
+        super(FusedMBConv, self).__init__()
+        use_bias = val2tuple(use_bias, 2)
+        norm_layer = val2tuple(norm_layer, 2)
+        act_layer = val2tuple(act_layer, 2)
+        mid_channels = mid_channels or round(in_channels * expand_ratio)
+
+        self.spatial_conv = ConvNormAct(
+            in_channels,
+            mid_channels,
+            kernel_size,
+            stride=stride,
+            groups=groups,
+            norm_layer=norm_layer[0],
+            act_layer=act_layer[0],
+            bias=use_bias[0],
+        )
+        self.point_conv = ConvNormAct(
+            mid_channels,
+            out_channels,
+            1,
+            norm_layer=norm_layer[1],
+            act_layer=act_layer[1],
+            bias=use_bias[1],
+        )
+
+    def forward(self, x):
+        x = self.spatial_conv(x)
+        x = self.point_conv(x)
+        return x
+
+
+class LiteMLA(nn.Module):
+    """Lightweight multi-scale linear attention"""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        heads: Union[int, None] = None,
+        heads_ratio: float = 1.0,
+        dim=8,
+        use_bias=False,
+        norm_layer=(None, nn.BatchNorm2d),
+        act_layer=(None, None),
+        kernel_func=nn.ReLU,
+        scales=(5,),
+        eps=1e-5,
+    ):
+        super(LiteMLA, self).__init__()
+        self.eps = eps
+        heads = heads or int(in_channels // dim * heads_ratio)
+        total_dim = heads * dim
+        use_bias = val2tuple(use_bias, 2)
+        norm_layer = val2tuple(norm_layer, 2)
+        act_layer = val2tuple(act_layer, 2)
+
+        self.dim = dim
+        self.qkv = ConvNormAct(
+            in_channels,
+            3 * total_dim,
+            1,
+            bias=use_bias[0],
+            norm_layer=norm_layer[0],
+            act_layer=act_layer[0],
+        )
+        self.aggreg = nn.ModuleList(
+            [
+                nn.Sequential(
+                    nn.Conv2d(
+                        3 * total_dim,
+                        3 * total_dim,
+                        scale,
+                        padding=get_same_padding(scale),
+                        groups=3 * total_dim,
+                        bias=use_bias[0],
+                    ),
+                    nn.Conv2d(
+                        3 * total_dim,
+                        3 * total_dim,
+                        1,
+                        groups=3 * heads,
+                        bias=use_bias[0],
+                    ),
+                )
+                for scale in scales
+            ]
+        )
+        self.kernel_func = kernel_func(inplace=False)
+
+        self.proj = ConvNormAct(
+            total_dim * (1 + len(scales)),
+            out_channels,
+            1,
+            bias=use_bias[1],
+            norm_layer=norm_layer[1],
+            act_layer=act_layer[1],
+        )
+
+    def _attn(self, q, k, v):
+        dtype = v.dtype
+        q, k, v = q.float(), k.float(), v.float()
+        kv = k.transpose(-1, -2) @ v
+        out = q @ kv
+        out = out[..., :-1] / (out[..., -1:] + self.eps)
+        return out.to(dtype)
+
+    def forward(self, x):
+        # Shape is B, C, H, W
+        B, _, H, W = x.shape
+
+        # generate multi-scale q, k, v
+        qkv = self.qkv(x)
+        multi_scale_qkv = [qkv]
+        for op in self.aggreg:
+            multi_scale_qkv.append(op(qkv))
+        multi_scale_qkv = torch.cat(multi_scale_qkv, dim=1)
+        multi_scale_qkv = multi_scale_qkv.reshape(B, -1, 3 * self.dim, H * W).transpose(
+            -1, -2
+        )
+        # Shape for each is B, C, HW, head_dim
+        q, k, v = multi_scale_qkv.chunk(3, dim=-1)
+
+        # lightweight global attention
+        q = self.kernel_func(q)
+        k = self.kernel_func(k)
+        v = F.pad(v, (0, 1), mode="constant", value=1.0)
+
+        out = self._attn(q, k, v)
+
+        # final projection
+        out = out.transpose(-1, -2).reshape(B, -1, H, W)
+        out = self.proj(out)
+        return out
+
+
+class EfficientVitBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        heads_ratio=1.0,
+        head_dim=32,
+        expand_ratio=4,
+        norm_layer=nn.BatchNorm2d,
+        act_layer=nn.Hardswish,
+    ):
+        super(EfficientVitBlock, self).__init__()
+        self.context_module = ResidualBlock(
+            LiteMLA(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                heads_ratio=heads_ratio,
+                dim=head_dim,
+                norm_layer=(None, norm_layer),
+            ),
+            nn.Identity(),
+        )
+        self.local_module = ResidualBlock(
+            MBConv(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                expand_ratio=expand_ratio,
+                use_bias=(True, True, False),
+                norm_layer=(None, None, norm_layer),
+                act_layer=(act_layer, act_layer, None),
+            ),
+            nn.Identity(),
+        )
+
+    def forward(self, x):
+        x = self.context_module(x)
+        x = self.local_module(x)
+        return x
+
+
+class ResidualBlock(nn.Module):
+    def __init__(
+        self,
+        main: Optional[nn.Module],
+        shortcut: Optional[nn.Module] = None,
+        pre_norm: Optional[nn.Module] = None,
+    ):
+        super(ResidualBlock, self).__init__()
+        self.pre_norm = pre_norm if pre_norm is not None else nn.Identity()
+        self.main = main
+        self.shortcut = shortcut
+
+    def forward(self, x):
+        res = self.main(self.pre_norm(x))
+        if self.shortcut is not None:
+            res = res + self.shortcut(x)
+        return res
+
+
+def build_local_block(
+    in_channels: int,
+    out_channels: int,
+    stride: int,
+    kernel_size: int,
+    expand_ratio: float,
+    norm_layer: str,
+    act_layer: str,
+    fewer_norm: bool = False,
+    block_type: str = "default",
+):
+    assert block_type in ["default", "large", "fused"]
+    if expand_ratio == 1:
+        if block_type == "default":
+            block = DSConv(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                stride=stride,
+                kernel_size=kernel_size,
+                use_bias=(True, False) if fewer_norm else False,
+                norm_layer=(None, norm_layer) if fewer_norm else norm_layer,
+                act_layer=(act_layer, None),
+            )
+        else:
+            block = ConvBlock(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                stride=stride,
+                kernel_size=kernel_size,
+                use_bias=(True, False) if fewer_norm else False,
+                norm_layer=(None, norm_layer) if fewer_norm else norm_layer,
+                act_layer=(act_layer, None),
+            )
+    else:
+        if block_type == "default":
+            block = MBConv(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                stride=stride,
+                kernel_size=kernel_size,
+                expand_ratio=expand_ratio,
+                use_bias=(True, True, False) if fewer_norm else False,
+                norm_layer=(None, None, norm_layer) if fewer_norm else norm_layer,
+                act_layer=(act_layer, act_layer, None),
+            )
+        else:
+            block = FusedMBConv(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                stride=stride,
+                kernel_size=kernel_size,
+                expand_ratio=expand_ratio,
+                use_bias=(True, False) if fewer_norm else False,
+                norm_layer=(None, norm_layer) if fewer_norm else norm_layer,
+                act_layer=(act_layer, None),
+            )
+    return block
+
+
+class Stem(nn.Sequential):
+    def __init__(
+        self,
+        in_chs,
+        out_chs,
+        depth,
+        stride,
+        norm_layer,
+        act_layer,
+        block_type="default",
+    ):
+        super().__init__()
+        self.stride = stride
+
+        self.add_module(
+            "in_conv",
+            ConvNormAct(
+                in_chs,
+                out_chs,
+                kernel_size=stride + 1,
+                stride=stride,
+                norm_layer=norm_layer,
+                act_layer=act_layer,
+            ),
+        )
+        stem_block = 0
+        for _ in range(depth):
+            self.add_module(
+                f"res{stem_block}",
+                ResidualBlock(
+                    build_local_block(
+                        in_channels=out_chs,
+                        out_channels=out_chs,
+                        stride=1,
+                        kernel_size=3,
+                        expand_ratio=1,
+                        norm_layer=norm_layer,
+                        act_layer=act_layer,
+                        block_type=block_type,
+                    ),
+                    nn.Identity(),
+                ),
+            )
+            stem_block += 1
+
+
+class EfficientVitLargeStage(nn.Module):
+    def __init__(
+        self,
+        in_chs,
+        out_chs,
+        depth,
+        stride,
+        norm_layer,
+        act_layer,
+        head_dim,
+        vit_stage=False,
+        fewer_norm=False,
+    ):
+        super(EfficientVitLargeStage, self).__init__()
+        blocks = [
+            ResidualBlock(
+                build_local_block(
+                    in_channels=in_chs,
+                    out_channels=out_chs,
+                    stride=stride,
+                    kernel_size=stride + 1,
+                    expand_ratio=24 if vit_stage else 16,
+                    norm_layer=norm_layer,
+                    act_layer=act_layer,
+                    fewer_norm=vit_stage or fewer_norm,
+                    block_type="default" if fewer_norm else "fused",
+                ),
+                None,
+            )
+        ]
+        in_chs = out_chs
+
+        if vit_stage:
+            # for stage 4
+            for _ in range(depth):
+                blocks.append(
+                    EfficientVitBlock(
+                        in_channels=in_chs,
+                        head_dim=head_dim,
+                        expand_ratio=6,
+                        norm_layer=norm_layer,
+                        act_layer=act_layer,
+                    )
+                )
+        else:
+            # for stage 1, 2, 3
+            for i in range(depth):
+                blocks.append(
+                    ResidualBlock(
+                        build_local_block(
+                            in_channels=in_chs,
+                            out_channels=out_chs,
+                            stride=1,
+                            kernel_size=3,
+                            expand_ratio=4,
+                            norm_layer=norm_layer,
+                            act_layer=act_layer,
+                            fewer_norm=fewer_norm,
+                            block_type="default" if fewer_norm else "fused",
+                        ),
+                        nn.Identity(),
+                    )
+                )
+
+        self.blocks = nn.Sequential(*blocks)
+
+    def forward(self, x):
+        return self.blocks(x)
+
+
+class EfficientVitLarge(nn.Module):
+    def __init__(
+        self,
+        config: EfficientViTConfig,
+        norm_layer=nn.BatchNorm2d,
+        act_layer=nn.Hardswish,
+    ):
+        super(EfficientVitLarge, self).__init__()
+        self.grad_checkpointing = False
+        self.num_classes = config.num_classes
+        self.norm_eps = config.layer_norm_eps
+        norm_layer = partial(norm_layer, eps=self.norm_eps)
+
+        # input stem
+        self.stem = Stem(
+            config.num_channels,
+            config.widths[0],
+            config.depths[0],
+            config.strides[0],
+            norm_layer,
+            act_layer,
+            block_type="large",
+        )
+        stride = config.strides[0]
+
+        # stages
+        self.feature_info = []
+        self.stages = nn.Sequential()
+        in_channels = config.widths[0]
+        for i, (w, d, s) in enumerate(
+            zip(config.widths[1:], config.depths[1:], config.strides[1:])
+        ):
+            self.stages.append(
+                EfficientVitLargeStage(
+                    in_channels,
+                    w,
+                    depth=d,
+                    stride=s,
+                    norm_layer=norm_layer,
+                    act_layer=act_layer,
+                    head_dim=config.head_dim,
+                    vit_stage=i >= 3,
+                    fewer_norm=i >= 2,
+                )
+            )
+            stride *= s
+            in_channels = w
+            self.feature_info += [
+                dict(num_chs=in_channels, reduction=stride, module=f"stages.{i}")
+            ]
+
+        self.num_features = in_channels
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.grad_checkpointing = enable
+
+    def forward(self, x):
+        x = self.stem(x)
+        encoder_hidden_states = []
+        for i, module in enumerate(self.stages):
+            x = module(x)
+            encoder_hidden_states.append(x)
+
+        return encoder_hidden_states
+
+
+class EfficientViTPreTrainedModel(SuryaPreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = EfficientViTConfig
+    base_model_prefix = "efficientvit"
+    main_input_name = "pixel_values"
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        if isinstance(module, (nn.Linear, nn.Conv2d)):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+
+class DecodeMLP(nn.Module):
+    def __init__(self, input_dim, output_dim):
+        super().__init__()
+        self.proj = nn.Linear(input_dim, output_dim)
+
+    def forward(self, hidden_states: torch.Tensor):
+        # Input is B, C, H, W
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+        # Output is B, HW, C
+        hidden_states = self.proj(hidden_states)
+        return hidden_states
+
+
+class DecodeHead(EfficientViTPreTrainedModel):
+    def __init__(self, config: EfficientViTConfig):
+        super().__init__(config)
+
+        # linear layers which will unify the channel dimension of each of the encoder blocks to the same config.decoder_hidden_size
+        mlps = []
+        for width in config.widths[1:]:
+            mlp = DecodeMLP(
+                input_dim=width, output_dim=config.decoder_layer_hidden_size
+            )
+            mlps.append(mlp)
+        self.linear_c = nn.ModuleList(mlps)
+
+        # the following 3 layers implement the ConvModule of the original implementation
+        self.linear_fuse = nn.Conv2d(
+            in_channels=config.decoder_layer_hidden_size * config.num_stages,
+            out_channels=config.decoder_hidden_size,
+            kernel_size=1,
+            bias=False,
+        )
+        self.batch_norm = nn.BatchNorm2d(config.decoder_hidden_size)
+        self.activation = nn.ReLU()
+
+        self.dropout = nn.Dropout(config.classifier_dropout_prob)
+        self.classifier = nn.Conv2d(
+            config.decoder_hidden_size, config.num_labels, kernel_size=1
+        )
+
+        self.config = config
+
+    def forward(self, encoder_hidden_states: torch.FloatTensor) -> torch.Tensor:
+        batch_size = encoder_hidden_states[-1].shape[0]
+
+        all_hidden_states = ()
+        for encoder_hidden_state, mlp in zip(encoder_hidden_states, self.linear_c):
+            height, width = encoder_hidden_state.shape[2], encoder_hidden_state.shape[3]
+            encoder_hidden_state = mlp(encoder_hidden_state)  # Output is B, HW, C
+            # Permute to B, C, HW
+            encoder_hidden_state = encoder_hidden_state.permute(0, 2, 1)
+            encoder_hidden_state = encoder_hidden_state.reshape(
+                batch_size, -1, height, width
+            )
+            # upsample
+            encoder_hidden_state = nn.functional.interpolate(
+                encoder_hidden_state,
+                size=encoder_hidden_states[0].size()[2:],
+                mode="bilinear",
+                align_corners=False,
+            )
+            all_hidden_states += (encoder_hidden_state,)
+
+        hidden_states = self.linear_fuse(torch.cat(all_hidden_states[::-1], dim=1))
+        hidden_states = self.batch_norm(hidden_states)
+        hidden_states = self.activation(hidden_states)
+
+        # logits are of shape (batch_size, num_labels, height/4, width/4)
+        logits = self.classifier(hidden_states)
+
+        return logits
+
+
+class EfficientViTForSemanticSegmentation(
+    S3DownloaderMixin, EfficientViTPreTrainedModel
+):
+    def __init__(self, config, **kwargs):
+        super().__init__(config)
+        self.vit = EfficientVitLarge(config)
+        self.decode_head = DecodeHead(config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self, pixel_values: torch.FloatTensor
+    ) -> Union[Tuple, SemanticSegmenterOutput]:
+        # Pixel values should be B,C,H,W
+        encoder_hidden_states = self.vit(
+            pixel_values,
+        )
+
+        logits = self.decode_head(encoder_hidden_states)
+
+        # Apply sigmoid to get 0-1 output
+        logits = torch.special.expit(logits)
+
+        return SemanticSegmenterOutput(
+            loss=None, logits=logits, hidden_states=encoder_hidden_states
+        )
+
+
+class EfficientViTForSemanticLayoutSegmentation(EfficientViTPreTrainedModel):
+    def __init__(self, config, **kwargs):
+        super().__init__(config, **kwargs)
+        self.vit = EfficientVitLarge(config)
+        self.decode_head = DecodeHead(config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self, pixel_values: torch.FloatTensor
+    ) -> Union[Tuple, SemanticSegmenterOutput]:
+        # Pixel values should be B,C,H,W
+        encoder_hidden_states = self.vit(
+            pixel_values,
+        )
+
+        logits = self.decode_head(encoder_hidden_states)
+
+        # Apply sigmoid to get 0-1 output
+        logits = torch.special.expit(logits)
+
+        return SemanticSegmenterOutput(
+            loss=None, logits=logits, hidden_states=encoder_hidden_states
+        )

surya/detection/parallel.py

@@ -0,0 +1,19 @@
+class FakeFuture:
+    def __init__(self, func, *args, **kwargs):
+        self._result = func(*args, **kwargs)
+
+    def result(self):
+        return self._result
+
+class FakeExecutor:
+    def __init__(self, **kwargs):
+        pass
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *excinfo):
+        pass
+
+    def submit(self, fn, *args, **kwargs):
+        return FakeFuture(fn, *args, **kwargs)

surya/detection/processor.py

@@ -0,0 +1,317 @@
+# coding=utf-8
+# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Modified image processor class for Segformer based on transformers"""
+
+import warnings
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+
+from transformers.image_processing_utils import (
+    BaseImageProcessor,
+    BatchFeature,
+    get_size_dict,
+)
+from transformers.image_transforms import to_channel_dimension_format
+from transformers.image_utils import (
+    IMAGENET_DEFAULT_MEAN,
+    IMAGENET_DEFAULT_STD,
+    ChannelDimension,
+    ImageInput,
+    PILImageResampling,
+    infer_channel_dimension_format,
+    make_list_of_images,
+)
+from transformers.utils import TensorType
+
+
+import PIL.Image
+
+from surya.common.s3 import S3DownloaderMixin
+
+
+class SegformerImageProcessor(S3DownloaderMixin, BaseImageProcessor):
+    r"""
+    Constructs a Segformer image processor.
+
+    Args:
+        do_resize (`bool`, *optional*, defaults to `True`):
+            Whether to resize the image's (height, width) dimensions to the specified `(size["height"],
+            size["width"])`. Can be overridden by the `do_resize` parameter in the `preprocess` method.
+        size (`Dict[str, int]` *optional*, defaults to `{"height": 512, "width": 512}`):
+            Size of the output image after resizing. Can be overridden by the `size` parameter in the `preprocess`
+            method.
+        resample (`PILImageResampling`, *optional*, defaults to `Resampling.BILINEAR`):
+            Resampling filter to use if resizing the image. Can be overridden by the `resample` parameter in the
+            `preprocess` method.
+        do_rescale (`bool`, *optional*, defaults to `True`):
+            Whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by the `do_rescale`
+            parameter in the `preprocess` method.
+        rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
+            Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess`
+            method.
+        do_normalize (`bool`, *optional*, defaults to `True`):
+            Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess`
+            method.
+        image_mean (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_MEAN`):
+            Mean to use if normalizing the image. This is a float or list of floats the length of the number of
+            channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method.
+        image_std (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_STD`):
+            Standard deviation to use if normalizing the image. This is a float or list of floats the length of the
+            number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method.
+        do_reduce_labels (`bool`, *optional*, defaults to `False`):
+            Whether or not to reduce all label values of segmentation maps by 1. Usually used for datasets where 0 is
+            used for background, and background itself is not included in all classes of a dataset (e.g. ADE20k). The
+            background label will be replaced by 255. Can be overridden by the `do_reduce_labels` parameter in the
+            `preprocess` method.
+    """
+
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        size: Dict[str, int] = None,
+        resample: PILImageResampling = PILImageResampling.BILINEAR,
+        do_rescale: bool = True,
+        rescale_factor: Union[int, float] = 1 / 255,
+        do_normalize: bool = True,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        do_reduce_labels: bool = False,
+        **kwargs,
+    ) -> None:
+        if "reduce_labels" in kwargs:
+            warnings.warn(
+                "The `reduce_labels` parameter is deprecated and will be removed in a future version. Please use "
+                "`do_reduce_labels` instead.",
+                FutureWarning,
+            )
+            do_reduce_labels = kwargs.pop("reduce_labels")
+
+        super().__init__(**kwargs)
+        size = size if size is not None else {"height": 512, "width": 512}
+        size = get_size_dict(size)
+        self.do_resize = do_resize
+        self.size = size
+        self.resample = resample
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.do_normalize = do_normalize
+        self.image_mean = (
+            image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
+        )
+        self.image_std = image_std if image_std is not None else IMAGENET_DEFAULT_STD
+        self.do_reduce_labels = do_reduce_labels
+        self._valid_processor_keys = [
+            "images",
+            "segmentation_maps",
+            "do_resize",
+            "size",
+            "resample",
+            "do_rescale",
+            "rescale_factor",
+            "do_normalize",
+            "image_mean",
+            "image_std",
+            "do_reduce_labels",
+            "return_tensors",
+            "data_format",
+            "input_data_format",
+        ]
+
+    @classmethod
+    def from_dict(cls, image_processor_dict: Dict[str, Any], **kwargs):
+        """
+        Overrides the `from_dict` method from the base class to make sure `do_reduce_labels` is updated if image
+        processor is created using from_dict and kwargs e.g. `SegformerImageProcessor.from_pretrained(checkpoint,
+        reduce_labels=True)`
+        """
+        image_processor_dict = image_processor_dict.copy()
+        if "reduce_labels" in kwargs:
+            image_processor_dict["reduce_labels"] = kwargs.pop("reduce_labels")
+        return super().from_dict(image_processor_dict, **kwargs)
+
+    def _preprocess(
+        self,
+        image: ImageInput,
+        do_resize: bool,
+        do_rescale: bool,
+        do_normalize: bool,
+        size: Optional[Dict[str, int]] = None,
+        resample: PILImageResampling = None,
+        rescale_factor: Optional[float] = None,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+    ):
+        if do_rescale:
+            image = self.rescale(
+                image=image, scale=rescale_factor, input_data_format=input_data_format
+            )
+
+        if do_normalize:
+            image = self.normalize(
+                image=image,
+                mean=image_mean,
+                std=image_std,
+                input_data_format=input_data_format,
+            )
+
+        return image
+
+    def _preprocess_image(
+        self,
+        image: ImageInput,
+        do_resize: bool = None,
+        size: Dict[str, int] = None,
+        resample: PILImageResampling = None,
+        do_rescale: bool = None,
+        rescale_factor: float = None,
+        do_normalize: bool = None,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        data_format: Optional[Union[str, ChannelDimension]] = None,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+    ) -> np.ndarray:
+        """Preprocesses a single image."""
+        # All transformations expect numpy arrays.
+        if input_data_format is None:
+            input_data_format = infer_channel_dimension_format(image)
+
+        image = self._preprocess(
+            image=image,
+            do_resize=do_resize,
+            size=size,
+            resample=resample,
+            do_rescale=do_rescale,
+            rescale_factor=rescale_factor,
+            do_normalize=do_normalize,
+            image_mean=image_mean,
+            image_std=image_std,
+            input_data_format=input_data_format,
+        )
+        if data_format is not None:
+            image = to_channel_dimension_format(
+                image, data_format, input_channel_dim=input_data_format
+            )
+        return image
+
+    def __call__(self, images, segmentation_maps=None, **kwargs):
+        """
+        Preprocesses a batch of images and optionally segmentation maps.
+
+        Overrides the `__call__` method of the `Preprocessor` class so that both images and segmentation maps can be
+        passed in as positional arguments.
+        """
+        return super().__call__(images, segmentation_maps=segmentation_maps, **kwargs)
+
+    def preprocess(
+        self,
+        images: ImageInput,
+        segmentation_maps: Optional[ImageInput] = None,
+        do_resize: Optional[bool] = None,
+        size: Optional[Dict[str, int]] = None,
+        resample: PILImageResampling = None,
+        do_rescale: Optional[bool] = None,
+        rescale_factor: Optional[float] = None,
+        do_normalize: Optional[bool] = None,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        do_reduce_labels: Optional[bool] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        data_format: ChannelDimension = ChannelDimension.FIRST,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+        **kwargs,
+    ) -> PIL.Image.Image:
+        """
+        Preprocess an image or batch of images.
+
+        Args:
+            images (`ImageInput`):
+                Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
+                passing in images with pixel values between 0 and 1, set `do_rescale=False`.
+            segmentation_maps (`ImageInput`, *optional*):
+                Segmentation map to preprocess.
+            do_resize (`bool`, *optional*, defaults to `self.do_resize`):
+                Whether to resize the image.
+            size (`Dict[str, int]`, *optional*, defaults to `self.size`):
+                Size of the image after `resize` is applied.
+            resample (`int`, *optional*, defaults to `self.resample`):
+                Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`, Only
+                has an effect if `do_resize` is set to `True`.
+            do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
+                Whether to rescale the image values between [0 - 1].
+            rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
+                Rescale factor to rescale the image by if `do_rescale` is set to `True`.
+            do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
+                Whether to normalize the image.
+            image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
+                Image mean.
+            image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
+                Image standard deviation.
+            do_reduce_labels (`bool`, *optional*, defaults to `self.do_reduce_labels`):
+                Whether or not to reduce all label values of segmentation maps by 1. Usually used for datasets where 0
+                is used for background, and background itself is not included in all classes of a dataset (e.g.
+                ADE20k). The background label will be replaced by 255.
+            return_tensors (`str` or `TensorType`, *optional*):
+                The type of tensors to return. Can be one of:
+                    - Unset: Return a list of `np.ndarray`.
+                    - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
+                    - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
+                    - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
+                    - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
+            data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
+                The channel dimension format for the output image. Can be one of:
+                    - `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+                    - `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+            input_data_format (`ChannelDimension` or `str`, *optional*):
+                The channel dimension format for the input image. If unset, the channel dimension format is inferred
+                from the input image. Can be one of:
+                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
+        """
+        do_resize = do_resize if do_resize is not None else self.do_resize
+        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+        resample = resample if resample is not None else self.resample
+        size = size if size is not None else self.size
+        rescale_factor = (
+            rescale_factor if rescale_factor is not None else self.rescale_factor
+        )
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+
+        images = make_list_of_images(images)
+        images = [
+            self._preprocess_image(
+                image=img,
+                do_resize=do_resize,
+                resample=resample,
+                size=size,
+                do_rescale=do_rescale,
+                rescale_factor=rescale_factor,
+                do_normalize=do_normalize,
+                image_mean=image_mean,
+                image_std=image_std,
+                data_format=data_format,
+                input_data_format=input_data_format,
+            )
+            for img in images
+        ]
+
+        data = {"pixel_values": images}
+        return BatchFeature(data=data, tensor_type=return_tensors)

surya/detection/schema.py

@@ -0,0 +1,12 @@
+from typing import List, Optional, Any
+
+from pydantic import BaseModel
+
+from surya.common.polygon import PolygonBox
+
+
+class TextDetectionResult(BaseModel):
+    bboxes: List[PolygonBox]
+    heatmap: Optional[Any]
+    affinity_map: Optional[Any]
+    image_bbox: List[float]

surya/detection/util.py

@@ -0,0 +1,36 @@
+import math
+from PIL import ImageOps
+
+from surya.settings import settings
+
+
+def get_total_splits(image_size, height):
+    img_height = list(image_size)[1]
+    max_height = settings.DETECTOR_IMAGE_CHUNK_HEIGHT
+    if img_height > max_height:
+        num_splits = math.ceil(img_height / height)
+        return num_splits
+    return 1
+
+
+def split_image(img, height):
+    # This will not modify/return the original image - it will either crop, or copy the image
+    img_height = list(img.size)[1]
+    max_height = settings.DETECTOR_IMAGE_CHUNK_HEIGHT
+    if img_height > max_height:
+        num_splits = math.ceil(img_height / height)
+        splits = []
+        split_heights = []
+        for i in range(num_splits):
+            top = i * height
+            bottom = (i + 1) * height
+            if bottom > img_height:
+                bottom = img_height
+            cropped = img.crop((0, top, img.size[0], bottom))
+            chunk_height = bottom - top
+            if chunk_height < height:
+                cropped = ImageOps.pad(cropped, (img.size[0], height), color=255, centering=(0, 0))
+            splits.append(cropped)
+            split_heights.append(chunk_height)
+        return splits, split_heights
+    return [img.copy()], [img_height]

surya/endpoint/app.py

@@ -0,0 +1,97 @@
+import json
+import logging
+import logging.config
+import time
+import uuid
+
+import yaml
+from fastapi import FastAPI, Request
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import JSONResponse
+
+from surya.endpoint.legacy import router as legacy_router
+from surya.endpoint.openai import router as openai_router
+from surya.endpoint.service import _error_response
+
+with open("logger.yaml", "r", encoding="utf-8") as f:
+    config = yaml.safe_load(f.read())
+    logging.config.dictConfig(config)
+
+logger = logging.getLogger(__name__)
+
+app = FastAPI(
+    title="OCR SERVICE API SERVICE",
+    version="1.0",
+    docs_url="/v1/api/ai/swagger",
+    openapi_url="/v1/api/ai/openapi.json",
+)
+
+
+@app.exception_handler(RequestValidationError)
+async def validation_exception_handler(request: Request, exc: RequestValidationError):
+    request_id = getattr(request.state, "request_id", "-")
+    logger.warning(
+        "request_validation_failed request_id=%s method=%s path=%s errors=%s",
+        request_id,
+        request.method,
+        request.url.path,
+        exc.errors(),
+    )
+    return JSONResponse(
+        status_code=200,
+        content={
+            "data": [],
+            "message": json.dumps({"error info": exc.errors()}, default=str),
+            "code": 422,
+        },
+    )
+
+
+@app.middleware("http")
+async def allow_openapi_unauthorized(request: Request, call_next):
+    if request.url.path == "/v1/api/ai/openapi.json":
+        response = await call_next(request)
+        return response
+    return await call_next(request)
+
+
+@app.middleware("http")
+async def log_request_response(request, call_next):
+    request_id = uuid.uuid4().hex
+    request.state.request_id = request_id
+    start = time.perf_counter()
+    client = request.client.host if request.client else "-"
+    logger.info(
+        "request_start request_id=%s method=%s path=%s client=%s",
+        request_id,
+        request.method,
+        request.url.path,
+        client,
+    )
+    try:
+        response = await call_next(request)
+    except Exception as e:
+        duration_ms = (time.perf_counter() - start) * 1000
+        logger.exception(
+            "request_unhandled_exception request_id=%s method=%s path=%s duration_ms=%.2f",
+            request_id,
+            request.method,
+            request.url.path,
+            duration_ms,
+            exc_info=True,
+        )
+        return JSONResponse(status_code=500, content=_error_response(e))
+    duration_ms = (time.perf_counter() - start) * 1000
+    logger.info(
+        "request_end request_id=%s method=%s path=%s status_code=%s duration_ms=%.2f",
+        request_id,
+        request.method,
+        request.url.path,
+        response.status_code,
+        duration_ms,
+    )
+    return response
+
+
+app.include_router(legacy_router)
+app.include_router(openai_router)

surya/endpoint/legacy.py

@@ -0,0 +1,157 @@
+import io
+import time
+
+from fastapi import APIRouter, File, Request, UploadFile
+from PIL import Image
+
+from tools import (
+    ocr,
+    text_detection,
+    layout_detection,
+    table_recognition,
+    extract_text_from_image,
+)
+from vllm_batcher import vllm_ocr_batcher
+from vllm_tools import vllm_backend_info
+from surya.endpoint.schemas import ApiResponse, Info
+from surya.endpoint.service import (
+    _dump_model_or_list,
+    _error_response,
+    _load_base64_image,
+    _log_exception,
+    _ocr_response_data,
+    _success_response,
+    logger,
+)
+
+router = APIRouter()
+
+
+@router.post("/home")
+def home():
+    return "<h1>Welcome to SURYA OCR API!</h1>"
+
+
+@router.post("/v1/api/ai/suya_ocr", include_in_schema=False, response_model=ApiResponse)
+@router.post("/v1/api/ai/suya_ocr/", response_model=ApiResponse)
+def run_ocr(p: Info, request: Request):
+    try:
+        pil_image, pil_image_highres = _load_base64_image(p)
+        rec_img, pred, box_img = ocr(
+            pil_image,
+            pil_image_highres,
+            p.skip_text_detection,
+            p.recognize_math,
+            with_bboxes=p.ocr_with_boxes,
+        )
+        return _success_response(_ocr_response_data(pred))
+    except Exception as e:
+        results = _error_response(e)
+        _log_exception("suya_ocr", request, e, results)
+        return results
+
+
+@router.get("/v1/api/ai/suya_ocr_vllm/health", response_model=ApiResponse)
+async def suya_ocr_vllm_health():
+    return _success_response(vllm_backend_info())
+
+
+@router.post("/v1/api/ai/suya_ocr_vllm", include_in_schema=False, response_model=ApiResponse)
+@router.post("/v1/api/ai/suya_ocr_vllm/", response_model=ApiResponse)
+def run_ocr_vllm(p: Info, request: Request):
+    try:
+        pil_image, pil_image_highres = _load_base64_image(p)
+        start = time.perf_counter()
+        _, pred, _ = vllm_ocr_batcher.submit(
+            pil_image,
+            pil_image_highres,
+            skip_text_detection=p.skip_text_detection,
+            recognize_math=p.recognize_math,
+            with_bboxes=False,
+            request_id=getattr(request.state, "request_id", "-"),
+        )
+        logger.info(
+            "api_vllm_submit_wait_complete request_id=%s duration_ms=%.2f",
+            getattr(request.state, "request_id", "-"),
+            (time.perf_counter() - start) * 1000,
+        )
+        return _success_response(_ocr_response_data(pred))
+    except Exception as e:
+        results = _error_response(e)
+        _log_exception("suya_ocr_vllm", request, e, results)
+        return results
+
+
+@router.post("/v1/api/ai/suya_text_det", include_in_schema=False, response_model=ApiResponse)
+@router.post("/v1/api/ai/suya_text_det/", response_model=ApiResponse)
+def run_text_det(p: Info, request: Request):
+    try:
+        pil_image, _ = _load_base64_image(p)
+        det_img, text_pred = text_detection(pil_image)
+        text_lines = text_pred.model_dump(exclude=["heatmap", "affinity_map"])
+        return _success_response({"text_lines": text_lines})
+    except Exception as e:
+        results = _error_response(e)
+        _log_exception("suya_text_det", request, e, results)
+        return results
+
+
+@router.post("/v1/api/ai/suya_layout_det", include_in_schema=False, response_model=ApiResponse)
+@router.post("/v1/api/ai/suya_layout_det/", response_model=ApiResponse)
+def run_layout_det(p: Info, request: Request):
+    try:
+        pil_image, _ = _load_base64_image(p)
+        layout_img, pred = layout_detection(pil_image)
+        text_lines = pred.model_dump(exclude=["segmentation_map"])
+        return _success_response({"text_lines": text_lines})
+    except Exception as e:
+        results = _error_response(e)
+        _log_exception("suya_layout_det", request, e, results)
+        return results
+
+
+@router.post("/v1/api/ai/suya_table_rec", include_in_schema=False, response_model=ApiResponse)
+@router.post("/v1/api/ai/suya_table_rec/", response_model=ApiResponse)
+def run_table_rec(p: Info, request: Request):
+    try:
+        pil_image, pil_image_highres = _load_base64_image(p)
+
+        table_img, pred = table_recognition(
+            pil_image, pil_image_highres, p.skip_table_detection
+        )
+
+        text_json = _dump_model_or_list(pred)
+        text_lines = ""
+        if isinstance(pred, list):
+            text_lines = "\n".join(
+                [item.html for item in pred if getattr(item, "html", None)]
+            )
+        elif hasattr(pred, "text_lines"):
+            text_lines = "\n".join([p.text for p in pred.text_lines])
+
+        return _success_response({"ocr_text_json": text_json, "text_lines": text_lines})
+    except Exception as e:
+        results = _error_response(e)
+        _log_exception("suya_table_rec", request, e, results)
+        return results
+
+
+@router.post("/image2text", response_model=ApiResponse)
+async def image_to_text(request: Request, file: UploadFile = File(...)):
+    try:
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents))
+        logger.info(
+            "image2text_upload request_id=%s filename=%s content_type=%s size_bytes=%s image_size=%s",
+            getattr(request.state, "request_id", "-"),
+            file.filename,
+            file.content_type,
+            len(contents),
+            image.size,
+        )
+        full_text = extract_text_from_image(image)
+        return _success_response(full_text)
+    except Exception as e:
+        results = _error_response(e)
+        _log_exception("image2text", request, e, results)
+        return results

surya/endpoint/openai.py

@@ -0,0 +1,108 @@
+import time
+import uuid
+from typing import Any, Dict, List, Literal, Optional, Union
+
+from fastapi import APIRouter, Request
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel, ConfigDict
+
+from surya.endpoint import service
+
+router = APIRouter()
+
+
+class _ImageUrl(BaseModel):
+    model_config = ConfigDict(extra="allow")
+    url: str
+
+
+class _ContentPart(BaseModel):
+    model_config = ConfigDict(extra="allow")
+    type: str
+    image_url: Optional[_ImageUrl] = None
+    text: Optional[str] = None
+
+
+class _ChatMessage(BaseModel):
+    model_config = ConfigDict(extra="allow")
+    role: str
+    content: Union[str, List[_ContentPart]]
+
+
+class ChatCompletionRequest(BaseModel):
+    # The OpenAI SDK merges `extra_body={...}` into the TOP LEVEL of the request
+    # JSON, so the OCR control flags live here rather than nested under a key.
+    model_config = ConfigDict(extra="allow")
+    model: str = "surya-ocr"
+    messages: List[_ChatMessage]
+    stream: Optional[bool] = False
+    skip_text_detection: bool = False
+    recognize_math: bool = False
+    skip_table_detection: bool = False
+    ocr_with_boxes: bool = True
+    mode: Literal["block", "full_page", "table"] = "block"
+
+
+def _openai_error(message: str, err_type: str, status_code: int) -> JSONResponse:
+    return JSONResponse(
+        status_code=status_code,
+        content={"error": {"message": message, "type": err_type, "code": None}},
+    )
+
+
+def _build_chat_completion(
+    content: str, model: str, ocr_json: Any, elapsed: Optional[float]
+) -> Dict[str, Any]:
+    # Token accounting happens inside vLLM and is not surfaced to this layer, so
+    # usage counts are reported as 0 (prompt_tokens intentionally 0).
+    return {
+        "id": f"chatcmpl-{uuid.uuid4().hex}",
+        "object": "chat.completion",
+        "created": int(time.time()),
+        "model": model,
+        "choices": [
+            {
+                "index": 0,
+                "finish_reason": "stop",
+                "message": {"role": "assistant", "content": content},
+            }
+        ],
+        "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
+        "surya": {"ocr_text_json": ocr_json, "elapsed_seconds": elapsed},
+    }
+
+
+@router.post("/v1/chat/completions")
+def chat_completions(req: ChatCompletionRequest, request: Request):
+    request_id = getattr(request.state, "request_id", "-")
+    if req.stream:
+        return _openai_error(
+            "streaming is not supported by this OCR endpoint",
+            "invalid_request_error",
+            400,
+        )
+    try:
+        image = service.extract_last_image(req.messages)
+    except service.OcrRequestError as exc:
+        return _openai_error(str(exc), "invalid_request_error", 400)
+
+    try:
+        if req.mode == "table":
+            data = service.table_image(image, req.skip_table_detection)
+        else:
+            skip = req.skip_text_detection or req.mode == "full_page"
+            data = service.ocr_via_batcher(
+                image,
+                skip_text_detection=skip,
+                recognize_math=req.recognize_math,
+                request_id=request_id,
+            )
+    except service.OcrRequestError as exc:
+        return _openai_error(str(exc), "invalid_request_error", 400)
+    except Exception as exc:
+        service._log_exception("chat_completions", request, exc, {})
+        return _openai_error(str(exc), "internal_error", 500)
+
+    content = data.get("text_lines", "") or ""
+    ocr_json = data.get("ocr_text_json") if req.ocr_with_boxes else None
+    return _build_chat_completion(content, req.model, ocr_json, data.get("elapsed_seconds"))

surya/endpoint/schemas.py

@@ -0,0 +1,46 @@
+import base64
+import binascii
+from typing import Any, Literal, Optional
+
+from pydantic import (
+    BaseModel,
+    Field,
+    PrivateAttr,
+    StrictBool,
+    field_validator,
+    model_validator,
+)
+
+
+class Info(BaseModel):
+    _decoded_file: bytes = PrivateAttr(default=b"")
+
+    file: str  # base64 string
+    type: Literal["png", "jpg", "jpeg", "gif"]  # image type like 'png', 'jpg'
+    skip_text_detection: Optional[StrictBool] = Field(False)
+    skip_table_detection: Optional[StrictBool] = Field(False)
+    recognize_math: Optional[StrictBool] = Field(False)
+    ocr_with_boxes: Optional[StrictBool] = Field(True)
+
+    @field_validator("type", mode="before")
+    @classmethod
+    def normalize_type(cls, value: str) -> str:
+        if not isinstance(value, str):
+            raise ValueError("type must be a string")
+        return value.lower()
+
+    @model_validator(mode="after")
+    def decode_base64_file(self):
+        if not self.file:
+            raise ValueError("file must be a non-empty base64 string")
+        try:
+            self._decoded_file = base64.b64decode(self.file, validate=True)
+        except binascii.Error as exc:
+            raise ValueError("file must be valid base64") from exc
+        return self
+
+
+class ApiResponse(BaseModel):
+    data: Any
+    message: str
+    code: int

surya/endpoint/service.py

@@ -0,0 +1,166 @@
+import base64
+import binascii
+import io
+import json
+import logging
+import time
+import traceback
+from typing import Any, Dict
+
+import requests
+from fastapi import Request
+from PIL import Image
+
+import tools
+from vllm_batcher import vllm_ocr_batcher
+from vllm_tools import page_ocr_to_response
+
+logger = logging.getLogger("surya.endpoint")
+
+
+class OcrRequestError(Exception):
+    """Client-side error in an OCR request. Maps to HTTP 400."""
+
+
+# ---- helpers moved verbatim from api.py ------------------------------------
+
+def _load_base64_image(p, copy_highres: bool = False):
+    start = time.perf_counter()
+    filetype = p.type.lower()
+    allowed_types = ["png", "jpg", "jpeg", "gif"]
+    if filetype not in allowed_types:
+        raise ValueError("Unsupported file type")
+    pil_image = Image.open(io.BytesIO(p._decoded_file)).convert("RGB")
+    logger.info(
+        "api_decode_image duration_ms=%.2f size_bytes=%s image_size=%s copy_highres=%s",
+        (time.perf_counter() - start) * 1000,
+        len(p._decoded_file),
+        pil_image.size,
+        copy_highres,
+    )
+    return pil_image, pil_image.copy() if copy_highres else pil_image
+
+
+def _error_response(e: Exception):
+    tb = traceback.extract_tb(e.__traceback__)
+    frame = tb[-1] if tb else None
+    return {
+        "data": [],
+        "message": json.dumps(
+            {
+                "error info": str(e),
+                "error at": frame.filename if frame else "unknown",
+                "errort line at": frame.lineno if frame else 0,
+            }
+        ),
+        "code": 500,
+    }
+
+
+def _success_response(data: Any) -> Dict[str, Any]:
+    return {"data": data, "message": "success", "code": 200}
+
+
+def _log_exception(endpoint: str, request: Request, e: Exception, results: Dict[str, Any]):
+    request_id = getattr(request.state, "request_id", "-")
+    logger.exception(
+        "endpoint_failed request_id=%s endpoint=%s code=%s message=%s",
+        request_id,
+        endpoint,
+        results.get("code"),
+        results.get("message"),
+        exc_info=True,
+    )
+
+
+def _ocr_response_data(pred):
+    if hasattr(pred, "text_lines"):
+        text_json = pred.model_dump()
+        text_lines = "\n".join([p.text for p in pred.text_lines])
+        return {"ocr_text_json": text_json, "text_lines": text_lines}
+    return page_ocr_to_response(pred)
+
+
+def _dump_model_or_list(value):
+    if hasattr(value, "model_dump"):
+        return value.model_dump()
+    if isinstance(value, list):
+        return [
+            item.model_dump() if hasattr(item, "model_dump") else item
+            for item in value
+        ]
+    return value
+
+
+# ---- new helpers for the OpenAI endpoint -----------------------------------
+
+def load_image_from_url(url: str) -> Image.Image:
+    """Decode an OpenAI `image_url` (data: base64 URL or http(s): URL) to RGB."""
+    if url.startswith("data:"):
+        header, _, b64 = url.partition(",")
+        if "base64" not in header or not b64:
+            raise OcrRequestError("image_url data URL must be base64-encoded")
+        try:
+            raw = base64.b64decode(b64, validate=True)
+        except binascii.Error as exc:
+            raise OcrRequestError("image_url contains invalid base64") from exc
+    elif url.startswith("http://") or url.startswith("https://"):
+        try:
+            resp = requests.get(url, timeout=10)
+            resp.raise_for_status()
+        except Exception as exc:
+            raise OcrRequestError(f"could not fetch image_url: {exc}") from exc
+        raw = resp.content
+    else:
+        raise OcrRequestError("image_url must be a data: or http(s): URL")
+    try:
+        return Image.open(io.BytesIO(raw)).convert("RGB")
+    except Exception as exc:
+        raise OcrRequestError(f"could not decode image: {exc}") from exc
+
+
+def extract_last_image(messages) -> Image.Image:
+    """Return the last image found across the chat messages' content parts."""
+    url = None
+    for msg in messages:
+        content = msg.content
+        if isinstance(content, list):
+            for part in content:
+                if getattr(part, "type", None) == "image_url" and part.image_url:
+                    url = part.image_url.url
+    if url is None:
+        raise OcrRequestError("no image_url content part found in messages")
+    return load_image_from_url(url)
+
+
+def ocr_via_batcher(
+    image: Image.Image,
+    *,
+    skip_text_detection: bool,
+    recognize_math: bool,
+    request_id: str,
+) -> Dict[str, Any]:
+    """Run page OCR through the shared request batcher. The annotated image is
+    discarded, so we request with_bboxes=False to maximise batch coalescing with
+    the primary /suya_ocr_vllm endpoint."""
+    _, pred, _ = vllm_ocr_batcher.submit(
+        image,
+        image,
+        skip_text_detection=skip_text_detection,
+        recognize_math=recognize_math,
+        with_bboxes=False,
+        request_id=request_id,
+    )
+    return _ocr_response_data(pred)
+
+
+def table_image(image: Image.Image, skip_table_detection: bool) -> Dict[str, Any]:
+    """Run table-structure recognition; mirrors the /suya_table_rec response shape."""
+    _, pred = tools.table_recognition(image, image.copy(), skip_table_detection)
+    text_json = _dump_model_or_list(pred)
+    text_lines = ""
+    if isinstance(pred, list):
+        text_lines = "\n".join([item.html for item in pred if getattr(item, "html", None)])
+    elif hasattr(pred, "text_lines"):
+        text_lines = "\n".join([line.text for line in pred.text_lines])
+    return {"ocr_text_json": text_json, "text_lines": text_lines, "elapsed_seconds": None}

surya/inference/__init__.py

@@ -0,0 +1,116 @@
+"""Surya inference manager.
+
+One process owns one SuryaInferenceManager. The manager wraps a single backend
+(vllm | llamacpp) which speaks OpenAI-compatible chat completions.
+
+Predictors take the manager via explicit injection at construction time.
+"""
+
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+from typing import List, Optional
+
+from surya.inference.backends.base import Backend
+from surya.inference.schema import BatchInputItem, BatchOutputItem
+from surya.logging import get_logger
+from surya.settings import settings
+from surya.timing import timing_span
+
+logger = get_logger()
+
+
+def _has_nvidia_gpu() -> bool:
+    """True if an NVIDIA GPU is present on this host.
+
+    We deliberately do *not* rely solely on ``torch.cuda.is_available()``:
+    the installed torch wheel's CUDA build can be newer than the host driver
+    (PyPI's default wheel tracks the latest CUDA), in which case torch reports
+    no CUDA even on a perfectly good GPU box. That would silently route us to
+    the CPU llama.cpp backend on a machine that should be running vllm. So we
+    take torch's word when it *does* see CUDA, and otherwise fall back to
+    probing for the GPU directly via ``nvidia-smi``.
+    """
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            return True
+    except Exception:
+        pass
+
+    # Instant, load-independent check: the NVIDIA device node only exists when
+    # a GPU + driver are present. Preferred over nvidia-smi because nvidia-smi
+    # can block for several seconds on a GPU under heavy load, which would race
+    # a timeout and falsely report "no GPU".
+    if os.path.exists("/dev/nvidia0"):
+        return True
+
+    nvidia_smi = shutil.which("nvidia-smi")
+    if not nvidia_smi:
+        return False
+    try:
+        result = subprocess.run(
+            [nvidia_smi, "-L"], capture_output=True, text=True, timeout=15
+        )
+        return result.returncode == 0 and "GPU" in result.stdout
+    except Exception:
+        return False
+
+
+def _autodetect_backend() -> str:
+    if settings.SURYA_INFERENCE_BACKEND:
+        return settings.SURYA_INFERENCE_BACKEND
+    # NVIDIA GPU → vllm, mps/cpu → llamacpp
+    if _has_nvidia_gpu():
+        return "vllm"
+    return "llamacpp"
+
+
+def _build_backend(method: str) -> Backend:
+    method = method.lower()
+    if method == "vllm":
+        from surya.inference.backends.vllm import VllmBackend
+
+        return VllmBackend()
+    if method == "llamacpp":
+        from surya.inference.backends.llamacpp import LlamaCppBackend
+        return LlamaCppBackend()
+    raise ValueError(
+        f"Unknown inference backend {method!r}. Supported: 'vllm', 'llamacpp'."
+    )
+
+
+class SuryaInferenceManager:
+    """Single entry point for VLM inference. Construct once per process."""
+
+    def __init__(self, method: Optional[str] = None, lazy: bool = True):
+        self.method = method or _autodetect_backend()
+        self.backend: Backend = _build_backend(self.method)
+        if not lazy:
+            self.backend.start()
+
+    def start(self) -> None:
+        self.backend.start()
+
+    def stop(self) -> None:
+        self.backend.stop()
+
+    def generate(self, batch: List[BatchInputItem]) -> List[BatchOutputItem]:
+        with timing_span("surya_manager_generate", backend=self.method, item_count=len(batch)):
+            return self.backend.generate(batch)
+
+
+# Module-level lazy singleton for callers that don't want explicit construction
+# (notebooks, ad-hoc scripts). Surya's own models.py and marker should use
+# explicit construction.
+_default_manager: Optional[SuryaInferenceManager] = None
+
+
+def get_default_manager() -> SuryaInferenceManager:
+    global _default_manager
+    if _default_manager is None:
+        _default_manager = SuryaInferenceManager()
+    return _default_manager

surya/inference/backends/__init__.py

@@ -0,0 +1,2 @@
+from surya.inference.backends.base import Backend as Backend
+from surya.inference.backends.base import ServerHandle as ServerHandle

surya/inference/backends/base.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List
+
+
+from surya.inference.schema import BatchInputItem, BatchOutputItem
+
+
+@dataclass
+class ServerHandle:
+    base_url: str  # e.g. "http://127.0.0.1:8765/v1"
+    model_name: str  # what gets passed in OpenAI `model` field
+    spawned_by_us: bool  # if True, we manage atexit cleanup
+
+
+class Backend:
+    """Abstract backend. Concrete backends own server lifecycle + generation."""
+
+    name: str  # "vllm" | "llamacpp"
+
+    def start(self) -> ServerHandle:
+        """Idempotent: probe → attach if alive, else spawn. Returns handle."""
+        raise NotImplementedError
+
+    def stop(self) -> None:
+        """Stop the server if we spawned it."""
+        raise NotImplementedError
+
+    def generate(self, batch: List[BatchInputItem]) -> List[BatchOutputItem]:
+        raise NotImplementedError

surya/inference/backends/llamacpp.py

@@ -0,0 +1,207 @@
+"""llama.cpp backend: spawns the upstream `llama-server` binary natively.
+
+Install:
+- macOS:    brew install llama.cpp     (Metal build, MPS)
+- Linux:    brew install llama.cpp  OR  github.com/ggml-org/llama.cpp/releases
+"""
+
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+from pathlib import Path
+from typing import List, Optional
+
+from huggingface_hub import hf_hub_download
+from openai import OpenAI
+
+from surya.inference.backends.base import Backend, ServerHandle
+from surya.inference.backends.openai_client import chat_completions_batch
+from surya.inference.backends.spawn import (
+    SpawnHandle,
+    SpawnError,
+    attach_or_spawn,
+)
+from surya.inference.schema import BatchInputItem, BatchOutputItem
+from surya.logging import get_logger
+from surya.settings import settings
+
+logger = get_logger()
+
+
+def _resolve_llama_server_binary() -> str:
+    binary = settings.LLAMA_CPP_BINARY
+    if binary and os.path.isfile(binary):
+        return binary
+    found = shutil.which(binary or "llama-server")
+    if found:
+        return found
+    raise SpawnError(
+        "llama-server binary not found. Install with:\n"
+        "  macOS:  brew install llama.cpp\n"
+        "  Linux:  brew install llama.cpp  OR download from\n"
+        "          https://github.com/ggml-org/llama.cpp/releases\n"
+        "Or set LLAMA_CPP_BINARY in your env to the binary path."
+    )
+
+
+def _download_gguf_files() -> tuple[str, str]:
+    """Download model + mmproj GGUFs from HF Hub. Returns local paths."""
+    repo = settings.SURYA_GGUF_REPO
+    model_file = settings.SURYA_GGUF_MODEL_FILE
+    mmproj_file = settings.SURYA_GGUF_MMPROJ_FILE
+    logger.info(f"Downloading {model_file} and {mmproj_file} from {repo}")
+    model_path = hf_hub_download(repo_id=repo, filename=model_file)
+    mmproj_path = hf_hub_download(repo_id=repo, filename=mmproj_file)
+    return model_path, mmproj_path
+
+
+def _health_url(port: int) -> str:
+    return f"http://{settings.SURYA_INFERENCE_HOST}:{port}"
+
+
+def _openai_url(port: int) -> str:
+    return f"http://{settings.SURYA_INFERENCE_HOST}:{port}/v1"
+
+
+class LlamaCppBackend(Backend):
+    name = "llamacpp"
+
+    def __init__(self):
+        self.handle: Optional[ServerHandle] = None
+        self._client: Optional[OpenAI] = None
+
+    def start(self) -> ServerHandle:
+        if self.handle is not None:
+            return self.handle
+
+        # If user pinned an external server, attach without spawning.
+        # No binary or GGUF download needed in that case.
+        if settings.SURYA_INFERENCE_URL:
+            spawned = attach_or_spawn(
+                backend=self.name,
+                expected_model_name=settings.SURYA_MODEL_CHECKPOINT,
+                spawn_fn=lambda port: SpawnHandle(
+                    pid=None, cleanup_id="", cleanup_kind="process"
+                ),  # never called
+                health_url_for=_health_url,
+                openai_url_for=_openai_url,
+                startup_timeout=settings.SURYA_INFERENCE_STARTUP_TIMEOUT,
+            )
+            self.handle = ServerHandle(
+                base_url=spawned.base_url,
+                model_name=spawned.model_name,
+                spawned_by_us=spawned.spawned_by_us,
+            )
+            self._client = OpenAI(api_key="EMPTY", base_url=self.handle.base_url)
+            return self.handle
+
+        binary = _resolve_llama_server_binary()
+
+        # Pre-download GGUFs so the spawn doesn't race the download
+        if (
+            settings.SURYA_GGUF_LOCAL_MODEL_PATH
+            and settings.SURYA_GGUF_LOCAL_MMPROJ_PATH
+        ):
+            model_path = settings.SURYA_GGUF_LOCAL_MODEL_PATH
+            mmproj_path = settings.SURYA_GGUF_LOCAL_MMPROJ_PATH
+        else:
+            model_path, mmproj_path = _download_gguf_files()
+
+        # Total KV-cache budget. llama-server divides --ctx-size across
+        # --parallel slots, so a too-small total silently truncates outputs
+        # once each slot's share fills. Scale with parallel by default;
+        # SURYA_INFERENCE_CTX_SIZE overrides to a fixed value if set.
+        parallel = settings.SURYA_INFERENCE_PARALLEL
+        per_slot = settings.SURYA_INFERENCE_CTX_PER_SLOT
+        ctx_size = settings.SURYA_INFERENCE_CTX_SIZE
+        if ctx_size is None:
+            ctx_size = max(16384, parallel * per_slot)
+        effective_per_slot = ctx_size // max(parallel, 1)
+        logger.info(
+            f"llama-server ctx-size={ctx_size} "
+            f"(~{effective_per_slot}/slot × {parallel} parallel slots)"
+        )
+        if effective_per_slot < per_slot:
+            logger.warning(
+                f"per-slot ctx ({effective_per_slot}) is below recommended "
+                f"{per_slot}; outputs may truncate. Raise "
+                f"SURYA_INFERENCE_CTX_SIZE or SURYA_INFERENCE_CTX_PER_SLOT, "
+                f"or lower SURYA_INFERENCE_PARALLEL."
+            )
+
+        def spawn_fn(port: int) -> SpawnHandle:
+            cmd = [
+                binary,
+                "-m",
+                model_path,
+                "--mmproj",
+                mmproj_path,
+                "-ngl",
+                str(settings.LLAMA_CPP_NGL),
+                "--host",
+                settings.SURYA_INFERENCE_HOST,
+                "--port",
+                str(port),
+                "--parallel",
+                str(parallel),
+                "--ctx-size",
+                str(ctx_size),
+                "--no-mmproj-offload" if settings.LLAMA_CPP_NO_MMPROJ_OFFLOAD else "",
+                "--alias",
+                settings.SURYA_MODEL_CHECKPOINT,
+                "--jinja",
+            ]
+            cmd = [c for c in cmd if c]
+            for extra in (settings.LLAMA_CPP_EXTRA_ARGS or "").split():
+                cmd.append(extra)
+            logger.info(f"Spawning: {' '.join(cmd)}")
+            log_path = Path("~/.cache/datalab/surya/llamacpp_server.log").expanduser()
+            log_path.parent.mkdir(parents=True, exist_ok=True)
+            log_fp = open(log_path, "ab")
+            proc = subprocess.Popen(
+                cmd,
+                stdout=log_fp,
+                stderr=subprocess.STDOUT,
+                start_new_session=True,
+            )
+            return SpawnHandle(
+                pid=proc.pid, cleanup_id=str(proc.pid), cleanup_kind="process"
+            )
+
+        spawned = attach_or_spawn(
+            backend=self.name,
+            expected_model_name=settings.SURYA_MODEL_CHECKPOINT,
+            spawn_fn=spawn_fn,
+            health_url_for=_health_url,
+            openai_url_for=_openai_url,
+            startup_timeout=settings.SURYA_INFERENCE_STARTUP_TIMEOUT,
+        )
+        self.handle = ServerHandle(
+            base_url=spawned.base_url,
+            model_name=spawned.model_name,
+            spawned_by_us=spawned.spawned_by_us,
+        )
+        self._client = OpenAI(
+            api_key="EMPTY",
+            base_url=self.handle.base_url,
+        )
+        return self.handle
+
+    def stop(self) -> None:
+        # atexit handler in spawn.py owns cleanup; nothing to do here.
+        self.handle = None
+        self._client = None
+
+    def generate(self, batch: List[BatchInputItem]) -> List[BatchOutputItem]:
+        if self.handle is None or self._client is None:
+            self.start()
+        return chat_completions_batch(
+            batch,
+            client=self._client,
+            model_name=self.handle.model_name,
+            timeout=settings.SURYA_INFERENCE_TIMEOUT_SECONDS,
+            max_workers=settings.SURYA_INFERENCE_PARALLEL,
+            request_logprobs_default=settings.SURYA_INFERENCE_LOGPROBS,
+        )

surya/inference/backends/openai_client.py

@@ -0,0 +1,257 @@
+"""Shared OpenAI-compatible chat completions client. Used by vllm + llama.cpp.
+
+Both servers expose `/v1/chat/completions` with the same request/response shape,
+so this module is the single point of HTTP contact for both backends.
+"""
+
+from __future__ import annotations
+
+import base64
+import io
+import math
+import os
+import time
+from concurrent.futures import ThreadPoolExecutor
+from typing import List, Optional
+
+from PIL import Image
+
+from surya.inference.prompts import PROMPT_MAPPING
+from surya.inference.schema import (
+    BatchInputItem,
+    BatchOutputItem,
+    GenerationResult,
+)
+from surya.inference.util import detect_repeat_token, scale_to_fit
+from surya.logging import get_logger
+from surya.timing import get_current_timing, set_current_timing, timing_span
+
+logger = get_logger()
+
+
+def resolve_max_workers(batch_len: int, max_inflight: int) -> int:
+    """Concurrent HTTP workers for a batch: as many as the batch needs, capped
+    by max_inflight so we keep vLLM's sequence slots full without over-queueing."""
+    return max(1, min(batch_len, max_inflight))
+
+
+def encode_image_b64(image: Image.Image) -> tuple[str, str]:
+    image_format = os.getenv("SUYA_VLLM_IMAGE_FORMAT", "JPEG").upper()
+    if image_format not in {"JPEG", "PNG"}:
+        raise ValueError("SUYA_VLLM_IMAGE_FORMAT must be JPEG or PNG")
+
+    buf = io.BytesIO()
+    if image_format == "JPEG":
+        quality = int(os.getenv("SUYA_VLLM_JPEG_QUALITY", "92"))
+        image.save(buf, format="JPEG", quality=quality, subsampling=0)
+        mime_type = "image/jpeg"
+    else:
+        image.save(buf, format="PNG")
+        mime_type = "image/png"
+    view = buf.getbuffer()
+    try:
+        return base64.b64encode(view).decode("ascii"), mime_type
+    finally:
+        view.release()
+
+
+def _build_messages(image: Image.Image, prompt: str):
+    with timing_span("openai_encode_image_b64", image_size=image.size):
+        image_b64, mime_type = encode_image_b64(image)
+    return [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image_url",
+                    "image_url": {"url": f"data:{mime_type};base64,{image_b64}"},
+                },
+                {"type": "text", "text": prompt},
+            ],
+        }
+    ]
+
+
+def _mean_token_prob(logprobs_content) -> Optional[float]:
+    if not logprobs_content:
+        return None
+    probs = []
+    for tok in logprobs_content:
+        lp = (
+            tok.get("logprob")
+            if isinstance(tok, dict)
+            else getattr(tok, "logprob", None)
+        )
+        if lp is None:
+            continue
+        probs.append(math.exp(lp))
+    if not probs:
+        return None
+    return sum(probs) / len(probs)
+
+
+def _generate_one(
+    item: BatchInputItem,
+    client,
+    model_name: str,
+    max_tokens_default: int,
+    temperature: float,
+    top_p: float,
+    timeout: float,
+    request_logprobs_default: bool,
+) -> GenerationResult:
+    prompt = item.prompt or PROMPT_MAPPING[item.prompt_type]
+    with timing_span("openai_scale_image", prompt_type=item.prompt_type, image_size=item.image.size):
+        image = scale_to_fit(item.image)
+    with timing_span("openai_build_messages", prompt_type=item.prompt_type):
+        messages = _build_messages(image, prompt)
+
+    max_tokens = item.max_tokens or max_tokens_default
+    request_logprobs = item.request_logprobs or request_logprobs_default
+
+    kwargs = dict(
+        model=model_name,
+        messages=messages,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        timeout=timeout,
+    )
+    if request_logprobs:
+        kwargs["logprobs"] = True
+
+    # Structured output: prefer OpenAI-standard response_format (works on both
+    # vllm and llama.cpp). Fall back to vllm's extra_body for guided_regex.
+    if item.guided_json is not None:
+        kwargs["response_format"] = {
+            "type": "json_schema",
+            "json_schema": {
+                "name": "structured_output",
+                "schema": item.guided_json,
+                "strict": True,
+            },
+        }
+    if item.guided_regex is not None:
+        kwargs.setdefault("extra_body", {})["guided_regex"] = item.guided_regex
+
+    try:
+        with timing_span(
+            "openai_chat_completion",
+            prompt_type=item.prompt_type,
+            max_tokens=max_tokens,
+        ):
+            completion = client.chat.completions.create(**kwargs)
+        raw = completion.choices[0].message.content or ""
+        token_count = completion.usage.completion_tokens if completion.usage else 0
+        with timing_span(
+            "openai_parse_completion",
+            prompt_type=item.prompt_type,
+            token_count=token_count,
+        ):
+            mean_p = None
+            logprobs_content = None
+            if request_logprobs:
+                choice = completion.choices[0]
+                lp = getattr(choice, "logprobs", None)
+                if lp is not None:
+                    content = getattr(lp, "content", None)
+                    if content is not None:
+                        logprobs_content = [
+                            c.model_dump() if hasattr(c, "model_dump") else c
+                            for c in content
+                        ]
+                        mean_p = _mean_token_prob(content)
+        return GenerationResult(
+            raw=raw,
+            token_count=token_count,
+            error=False,
+            mean_token_prob=mean_p,
+            logprobs=logprobs_content,
+        )
+    except Exception as e:
+        logger.warning(f"Inference error: {e}")
+        return GenerationResult(raw="", token_count=0, error=True)
+
+
+def _should_retry(
+    result: GenerationResult,
+    retries: int,
+    max_retries: int,
+) -> bool:
+    if retries >= max_retries:
+        return False
+    if result.error:
+        return True
+    has_repeat = detect_repeat_token(result.raw) or (
+        len(result.raw) > 50 and detect_repeat_token(result.raw, cut_from_end=50)
+    )
+    return has_repeat
+
+
+def chat_completions_batch(
+    batch: List[BatchInputItem],
+    client,
+    model_name: str,
+    max_tokens_default: int = 2048,
+    temperature: float = 0.0,
+    top_p: float = 0.1,
+    timeout: float = 600.0,
+    max_workers: Optional[int] = None,
+    max_retries: int = 3,
+    request_logprobs_default: bool = True,
+) -> List[BatchOutputItem]:
+    """Run a batch of items through the chat completions endpoint with concurrent workers."""
+    if not batch:
+        return []
+    if max_workers is None:
+        max_workers = min(64, len(batch))
+    collector = get_current_timing()
+
+    def _process(item: BatchInputItem) -> BatchOutputItem:
+        if collector is not None:
+            set_current_timing(collector)
+        result = _generate_one(
+            item,
+            client=client,
+            model_name=model_name,
+            max_tokens_default=max_tokens_default,
+            temperature=temperature,
+            top_p=top_p,
+            timeout=timeout,
+            request_logprobs_default=request_logprobs_default,
+        )
+        retries = 0
+        while _should_retry(result, retries, max_retries):
+            backoff = 1.5 * (retries + 1) if result.error else 0
+            if backoff:
+                time.sleep(backoff)
+            retry_temp = min(temperature + 0.2 * (retries + 1), 0.8)
+            retry_top_p = 0.95 if not result.error else top_p
+            result = _generate_one(
+                item,
+                client=client,
+                model_name=model_name,
+                max_tokens_default=max_tokens_default,
+                temperature=retry_temp,
+                top_p=retry_top_p,
+                timeout=timeout,
+                request_logprobs_default=request_logprobs_default,
+            )
+            retries += 1
+        return BatchOutputItem(
+            raw=result.raw,
+            token_count=result.token_count,
+            error=result.error,
+            mean_token_prob=result.mean_token_prob,
+            logprobs=result.logprobs,
+            metadata=item.metadata,
+        )
+
+    with timing_span(
+        "openai_batch_threadpool",
+        item_count=len(batch),
+        max_workers=max_workers,
+        max_retries=max_retries,
+    ):
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            return list(executor.map(_process, batch))

surya/inference/backends/spawn.py

@@ -0,0 +1,351 @@
+"""Server lifecycle: probe, filelock, sentinel, atexit cleanup.
+
+Pattern: probe `/health` → if alive return handle → else acquire lock, re-probe,
+spawn detached, write sentinel, register atexit kill (only the spawner cleans up).
+"""
+
+from __future__ import annotations
+
+import atexit
+import json
+import os
+import socket
+import subprocess
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, Optional
+
+import httpx
+
+from surya.logging import get_logger
+from surya.settings import settings
+
+logger = get_logger()
+
+
+def _cache_dir() -> Path:
+    base = Path(os.path.expanduser("~/.cache/datalab/surya"))
+    base.mkdir(parents=True, exist_ok=True)
+    return base
+
+
+def _sentinel_path(backend: str) -> Path:
+    return _cache_dir() / f"{backend}_server.json"
+
+
+def _lock_path(backend: str) -> Path:
+    return _cache_dir() / f"{backend}_server.lock"
+
+
+def find_free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def probe_health(base_url: str, timeout: float = 1.0) -> bool:
+    """Returns True if the server reports healthy at /health."""
+    try:
+        # llama.cpp returns 200 on /health when ready; vllm returns 200 on /health too.
+        with httpx.Client(timeout=timeout) as client:
+            r = client.get(f"{base_url}/health")
+            return r.status_code == 200
+    except Exception:
+        return False
+
+
+def wait_for_health(
+    base_url: str, total_timeout: float = 300.0, interval: float = 1.0
+) -> bool:
+    deadline = time.time() + total_timeout
+    while time.time() < deadline:
+        if probe_health(base_url):
+            return True
+        time.sleep(interval)
+    return False
+
+
+def probe_model_id(openai_base: str, timeout: float = 5.0) -> Optional[str]:
+    """Returns the model id reported by the running server, or None on failure."""
+    try:
+        with httpx.Client(timeout=timeout) as client:
+            r = client.get(f"{openai_base}/models")
+            r.raise_for_status()
+            data = r.json()
+            models = data.get("data") or []
+            if models:
+                return models[0].get("id")
+    except Exception:
+        return None
+    return None
+
+
+@dataclass
+class SpawnedServer:
+    base_url: str  # full openai base, e.g. "http://127.0.0.1:8765/v1"
+    health_url: str  # base for /health, e.g. "http://127.0.0.1:8765"
+    model_name: str  # what to pass as `model`
+    pid: Optional[int]
+    backend: str
+    spawned_by_us: bool
+
+
+class SpawnError(RuntimeError):
+    pass
+
+
+def _read_sentinel(backend: str) -> Optional[dict]:
+    p = _sentinel_path(backend)
+    if not p.exists():
+        return None
+    try:
+        return json.loads(p.read_text())
+    except Exception:
+        return None
+
+
+def _write_sentinel(backend: str, data: dict) -> None:
+    _sentinel_path(backend).write_text(json.dumps(data))
+
+
+def _delete_sentinel(backend: str) -> None:
+    p = _sentinel_path(backend)
+    if p.exists():
+        try:
+            p.unlink()
+        except Exception:
+            pass
+
+
+def _stop_process(pid: int, name: str) -> None:
+    try:
+        # Graceful first
+        os.kill(pid, 15)  # SIGTERM
+        for _ in range(20):
+            try:
+                os.kill(pid, 0)  # still alive?
+            except ProcessLookupError:
+                logger.info(f"Stopped {name} (pid {pid})")
+                return
+            time.sleep(0.5)
+        # Hard
+        os.kill(pid, 9)
+        logger.warning(f"Force-killed {name} (pid {pid})")
+    except ProcessLookupError:
+        pass
+    except Exception as e:
+        logger.warning(f"Failed to stop {name} (pid {pid}): {e}")
+
+
+def _capture_server_logs(handle: "SpawnHandle", tail: int = 100) -> str:
+    """Best-effort tail of a server's logs, for surfacing startup failures."""
+    try:
+        if handle.cleanup_kind == "docker":
+            r = subprocess.run(
+                ["docker", "logs", "--tail", str(tail), handle.cleanup_id],
+                capture_output=True,
+                text=True,
+                timeout=15,
+            )
+            return (r.stdout or "") + (r.stderr or "") or "(no docker logs)"
+        # llama.cpp process backend logs to this file (see llamacpp.py)
+        log_path = Path("~/.cache/datalab/surya/llamacpp_server.log").expanduser()
+        if log_path.exists():
+            lines = log_path.read_text(errors="replace").splitlines()
+            return "\n".join(lines[-tail:]) or "(empty log)"
+    except Exception as e:
+        return f"(could not capture logs: {e})"
+    return "(no logs available)"
+
+
+def _stop_docker_container(name: str) -> None:
+    try:
+        subprocess.run(
+            ["docker", "stop", name], check=False, capture_output=True, timeout=30
+        )
+        logger.info(f"Stopped docker container {name}")
+    except Exception as e:
+        logger.warning(f"Failed to stop docker container {name}: {e}")
+
+
+def attach_or_spawn(
+    backend: str,
+    expected_model_name: str,
+    spawn_fn: Callable[[int], "SpawnHandle"],
+    health_url_for: Callable[[int], str],
+    openai_url_for: Callable[[int], str],
+    startup_timeout: float = 600.0,
+) -> SpawnedServer:
+    """Generic attach-or-spawn with file lock and sentinel.
+
+    `spawn_fn(port)` must launch the server detached and return a SpawnHandle
+    with `pid` (int or None for docker) and a `cleanup_id` (e.g. container name).
+    """
+    # 0. If user pinned an external URL, attach without lock
+    if settings.SURYA_INFERENCE_URL:
+        base_url = settings.SURYA_INFERENCE_URL.rstrip("/")
+        health_url = base_url[: -len("/v1")] if base_url.endswith("/v1") else base_url
+        if not probe_health(health_url):
+            raise SpawnError(
+                f"SURYA_INFERENCE_URL={base_url} is not reachable at /health. "
+                "Start the server or unset the variable."
+            )
+        model_name = probe_model_id(base_url) or expected_model_name
+        if model_name != expected_model_name:
+            raise SpawnError(
+                f"Model mismatch at {base_url}: expected {expected_model_name!r}, got {model_name!r}. "
+                "Stop the running server or unset SURYA_INFERENCE_URL."
+            )
+        return SpawnedServer(
+            base_url=base_url,
+            health_url=health_url,
+            model_name=model_name,
+            pid=None,
+            backend=backend,
+            spawned_by_us=False,
+        )
+
+    # 1. Probe sentinel without lock
+    existing = _read_sentinel(backend)
+    if existing:
+        port = existing.get("port")
+        pid = existing.get("pid")
+        if port and probe_health(health_url_for(port)):
+            running_model = probe_model_id(openai_url_for(port)) or expected_model_name
+            if running_model != expected_model_name:
+                raise SpawnError(
+                    f"Existing {backend} server on port {port} serves {running_model!r}, "
+                    f"expected {expected_model_name!r}. Stop it before continuing."
+                )
+            logger.info(f"Attaching to existing {backend} server on port {port}")
+            return SpawnedServer(
+                base_url=openai_url_for(port),
+                health_url=health_url_for(port),
+                model_name=running_model,
+                pid=pid,
+                backend=backend,
+                spawned_by_us=False,
+            )
+        else:
+            _delete_sentinel(backend)
+
+    if not settings.SURYA_INFERENCE_AUTOSTART:
+        raise SpawnError(
+            f"No running {backend} server and SURYA_INFERENCE_AUTOSTART is False. "
+            "Set the variable to True or start the server manually."
+        )
+
+    # 2. Acquire filelock to prevent races
+    try:
+        from filelock import FileLock
+    except ImportError as e:
+        raise SpawnError(
+            "filelock is required for server spawn. pip install filelock"
+        ) from e
+
+    lock = FileLock(str(_lock_path(backend)), timeout=120)
+    with lock:
+        # Re-check sentinel inside the lock
+        existing = _read_sentinel(backend)
+        if existing:
+            port = existing.get("port")
+            if port and probe_health(health_url_for(port)):
+                running_model = (
+                    probe_model_id(openai_url_for(port)) or expected_model_name
+                )
+                if running_model != expected_model_name:
+                    raise SpawnError(
+                        f"Existing {backend} server on port {port} serves {running_model!r}, "
+                        f"expected {expected_model_name!r}."
+                    )
+                return SpawnedServer(
+                    base_url=openai_url_for(port),
+                    health_url=health_url_for(port),
+                    model_name=running_model,
+                    pid=existing.get("pid"),
+                    backend=backend,
+                    spawned_by_us=False,
+                )
+
+        # 3. Spawn fresh
+        port = settings.SURYA_INFERENCE_PORT or find_free_port()
+        logger.info(f"Spawning {backend} server on port {port}")
+        spawn_handle = spawn_fn(port)
+
+        # 4. Write sentinel
+        _write_sentinel(
+            backend,
+            {
+                "port": port,
+                "pid": spawn_handle.pid,
+                "model": expected_model_name,
+                "backend": backend,
+                "cleanup_id": spawn_handle.cleanup_id,
+                "cleanup_kind": spawn_handle.cleanup_kind,
+            },
+        )
+
+        # 5. Register atexit cleanup (only spawner). Skipped when keep-alive is
+        # set so the server outlives this process and later commands attach to
+        # it via the sentinel. (_cleanup is still callable below on startup
+        # failure, where we always tear a half-started server down.)
+        def _cleanup():
+            try:
+                if spawn_handle.cleanup_kind == "docker":
+                    _stop_docker_container(spawn_handle.cleanup_id)
+                elif spawn_handle.cleanup_kind == "process":
+                    if spawn_handle.pid:
+                        _stop_process(spawn_handle.pid, backend)
+            finally:
+                _delete_sentinel(backend)
+
+        if settings.SURYA_INFERENCE_KEEP_ALIVE:
+            logger.info(
+                f"keep-alive: {backend} server on port {port} will stay up "
+                f"after exit (cleanup_id={spawn_handle.cleanup_id!r})"
+            )
+        else:
+            atexit.register(_cleanup)
+
+        # 6. Wait for health
+        health_url = health_url_for(port)
+        if not wait_for_health(health_url, total_timeout=startup_timeout):
+            # Grab the server's own logs *before* cleanup tears the (--rm)
+            # container down, otherwise the actual failure reason is lost and
+            # all the caller sees is this timeout.
+            logs = _capture_server_logs(spawn_handle)
+            _cleanup()
+            raise SpawnError(
+                f"{backend} server failed to become healthy at {health_url} "
+                f"within {startup_timeout}s.\n"
+                f"--- last {backend} server logs ---\n{logs}"
+            )
+
+        # 7. Verify model name
+        running_model = probe_model_id(openai_url_for(port))
+        if running_model and running_model != expected_model_name:
+            logger.warning(
+                f"{backend} server reports model={running_model!r} "
+                f"but expected {expected_model_name!r}; using reported name."
+            )
+            expected_model_name = running_model
+
+        logger.info(
+            f"{backend} server ready on port {port} (model={expected_model_name})"
+        )
+        return SpawnedServer(
+            base_url=openai_url_for(port),
+            health_url=health_url,
+            model_name=expected_model_name,
+            pid=spawn_handle.pid,
+            backend=backend,
+            spawned_by_us=True,
+        )
+
+
+@dataclass
+class SpawnHandle:
+    pid: Optional[int]
+    cleanup_id: str  # container name for docker, str(pid) for process
+    cleanup_kind: str  # "docker" | "process"

surya/inference/backends/vllm.py

@@ -0,0 +1,224 @@
+"""vllm backend: spawns the vllm/vllm-openai docker image with MTP=2."""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+import shutil
+import subprocess
+from typing import List, Optional
+
+from openai import OpenAI
+
+from surya.inference.backends.base import Backend, ServerHandle
+from surya.inference.backends.openai_client import chat_completions_batch, resolve_max_workers
+from surya.inference.backends.spawn import (
+    SpawnHandle,
+    SpawnError,
+    attach_or_spawn,
+)
+from surya.inference.schema import BatchInputItem, BatchOutputItem
+from surya.logging import get_logger
+from surya.settings import settings
+from surya.timing import timing_span
+
+logger = get_logger()
+
+
+# 24GB baseline (re-tune for surya-2 once benchmarks land)
+BASELINE_VRAM_GB = 24
+BASELINE_MAX_BATCHED_TOKENS = 8192
+BASELINE_MAX_NUM_SEQS = 32
+
+GPU_VRAM_GB = {
+    "b300": 270,
+    "b200": 180,
+    "h200": 141,
+    "h100": 80,
+    "a100-80": 80,
+    "a100": 40,
+    "a100-40": 40,
+    "l40s": 48,
+    "a10": 24,
+    "l4": 24,
+    "5090": 32,
+    "4090": 24,
+    "3090": 24,
+    "t4": 16,
+}
+
+
+def _gpu_settings(gpu: str) -> tuple[int, int]:
+    vram = GPU_VRAM_GB.get(gpu)
+    if vram is None:
+        available = ", ".join(sorted(GPU_VRAM_GB.keys()))
+        raise SpawnError(f"Unknown VLLM_GPU_TYPE {gpu!r}. Available: {available}")
+    ratio = vram / BASELINE_VRAM_GB
+    raw_tokens = BASELINE_MAX_BATCHED_TOKENS * ratio
+    max_batched_tokens = max(1024, 2 ** math.floor(math.log2(raw_tokens)))
+    max_num_seqs = max(8, (int(BASELINE_MAX_NUM_SEQS * ratio) // 8) * 8)
+    return max_batched_tokens, max_num_seqs
+
+
+def _resolve_docker_binary() -> str:
+    found = shutil.which("docker")
+    if found:
+        return found
+    raise SpawnError(
+        "docker binary not found. Install Docker (https://docs.docker.com/get-docker/) "
+        "and ensure the daemon is running."
+    )
+
+
+def _health_url(port: int) -> str:
+    return f"http://{settings.SURYA_INFERENCE_HOST}:{port}"
+
+
+def _openai_url(port: int) -> str:
+    return f"http://{settings.SURYA_INFERENCE_HOST}:{port}/v1"
+
+
+class VllmBackend(Backend):
+    name = "vllm"
+
+    def __init__(self):
+        self.handle: Optional[ServerHandle] = None
+        self._client: Optional[OpenAI] = None
+
+    def start(self) -> ServerHandle:
+        if self.handle is not None:
+            return self.handle
+
+        # If user pinned an external server, attach without spawning docker.
+        if settings.SURYA_INFERENCE_URL:
+            spawned = attach_or_spawn(
+                backend=self.name,
+                expected_model_name=settings.SURYA_MODEL_CHECKPOINT,
+                spawn_fn=lambda port: SpawnHandle(
+                    pid=None, cleanup_id="", cleanup_kind="docker"
+                ),
+                health_url_for=_health_url,
+                openai_url_for=_openai_url,
+                startup_timeout=settings.SURYA_INFERENCE_STARTUP_TIMEOUT,
+            )
+            self.handle = ServerHandle(
+                base_url=spawned.base_url,
+                model_name=spawned.model_name,
+                spawned_by_us=spawned.spawned_by_us,
+            )
+            self._client = OpenAI(
+                api_key=settings.VLLM_API_KEY, base_url=self.handle.base_url
+            )
+            return self.handle
+
+        if os.getenv("SUYA_ALLOW_NESTED_DOCKER", "false").lower() not in {"1", "true", "yes"}:
+            raise SpawnError(
+                "Nested Docker vLLM startup is disabled. Start vLLM in this "
+                "container and set SURYA_INFERENCE_URL, for example "
+                "http://127.0.0.1:8000/v1."
+            )
+
+        docker = _resolve_docker_binary()
+        max_batched_tokens, max_num_seqs = _gpu_settings(settings.VLLM_GPU_TYPE)
+
+        def spawn_fn(port: int) -> SpawnHandle:
+            container_name = f"surya-vllm-{port}"
+            hf_cache = os.path.expanduser(settings.DOCKER_HF_CACHE_PATH)
+            cmd = [
+                docker,
+                "run",
+                "--rm",
+                "-d",
+                "--name",
+                container_name,
+                "--runtime",
+                "nvidia",
+                "--gpus",
+                f"device={settings.VLLM_GPUS}",
+                "-v",
+                f"{hf_cache}:/root/.cache/huggingface",
+                "-p",
+                f"{port}:8000",
+                "--ipc=host",
+                settings.VLLM_DOCKER_IMAGE,
+                "--model",
+                settings.SURYA_MODEL_CHECKPOINT,
+                "--no-enforce-eager",
+                "--max-num-seqs",
+                str(max_num_seqs),
+                "--dtype",
+                settings.VLLM_DTYPE,
+                "--max-model-len",
+                str(settings.VLLM_MAX_MODEL_LEN),
+                "--max-num-batched-tokens",
+                str(max_batched_tokens),
+                "--gpu-memory-utilization",
+                str(settings.VLLM_GPU_MEMORY_UTILIZATION),
+                "--enable-prefix-caching",
+                "--mm-processor-kwargs",
+                json.dumps({"min_pixels": 3136, "max_pixels": 6291456}),
+                "--served-model-name",
+                settings.SURYA_MODEL_CHECKPOINT,
+            ]
+            if settings.VLLM_ENABLE_MTP:
+                spec_config = json.dumps(
+                    {
+                        "method": "mtp",
+                        "num_speculative_tokens": settings.VLLM_MTP_TOKENS,
+                    }
+                )
+                cmd.extend(["--speculative-config", spec_config])
+            for extra in (settings.VLLM_EXTRA_ARGS or "").split():
+                cmd.append(extra)
+            logger.info(f"Spawning: {' '.join(cmd)}")
+            result = subprocess.run(cmd, capture_output=True, text=True, check=False)
+            if result.returncode != 0:
+                raise SpawnError(f"docker run failed: {result.stderr or result.stdout}")
+            return SpawnHandle(
+                pid=None, cleanup_id=container_name, cleanup_kind="docker"
+            )
+
+        spawned = attach_or_spawn(
+            backend=self.name,
+            expected_model_name=settings.SURYA_MODEL_CHECKPOINT,
+            spawn_fn=spawn_fn,
+            health_url_for=_health_url,
+            openai_url_for=_openai_url,
+            startup_timeout=settings.SURYA_INFERENCE_STARTUP_TIMEOUT,
+        )
+        self.handle = ServerHandle(
+            base_url=spawned.base_url,
+            model_name=spawned.model_name,
+            spawned_by_us=spawned.spawned_by_us,
+        )
+        self._client = OpenAI(
+            api_key=settings.VLLM_API_KEY,
+            base_url=self.handle.base_url,
+        )
+        return self.handle
+
+    def stop(self) -> None:
+        self.handle = None
+        self._client = None
+
+    def generate(self, batch: List[BatchInputItem]) -> List[BatchOutputItem]:
+        if self.handle is None or self._client is None:
+            with timing_span("vllm_backend_start"):
+                self.start()
+        with timing_span(
+            "vllm_backend_generate",
+            item_count=len(batch),
+            parallel=settings.SURYA_INFERENCE_PARALLEL,
+        ):
+            return chat_completions_batch(
+                batch,
+                client=self._client,
+                model_name=self.handle.model_name,
+                timeout=settings.SURYA_INFERENCE_TIMEOUT_SECONDS,
+                max_workers=resolve_max_workers(
+                    len(batch), settings.SURYA_INFERENCE_MAX_INFLIGHT
+                ),
+                max_retries=settings.SURYA_INFERENCE_MAX_RETRIES,
+                request_logprobs_default=settings.SURYA_INFERENCE_LOGPROBS,
+            )

surya/inference/parsers.py

@@ -0,0 +1,182 @@
+"""Parsers for the three task outputs."""
+
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import dataclass
+from typing import List, Tuple
+
+
+from surya.logging import get_logger
+
+logger = get_logger()
+
+
+# ---- Layout (LAYOUT_PROMPT) -------------------------------------------------
+
+
+@dataclass
+class ParsedLayoutBlock:
+    label: str
+    bbox: Tuple[float, float, float, float]  # 0-1000 normalized
+    count: int  # multiple of 50, model's token estimate
+
+
+_JSON_ARRAY_RE = re.compile(r"\[.*\]", re.DOTALL)
+
+
+def _strip_fences(text: str) -> str:
+    cleaned = text.strip()
+    if cleaned.startswith("```"):
+        cleaned = re.sub(r"^```[a-zA-Z]*\n", "", cleaned)
+        cleaned = re.sub(r"\n```\s*$", "", cleaned)
+    return cleaned
+
+
+def _coerce_bbox(bbox) -> Tuple[float, float, float, float]:
+    if isinstance(bbox, str):
+        parts = [float(x) for x in bbox.replace(",", " ").split()]
+    else:
+        parts = [float(x) for x in bbox]
+    if len(parts) != 4:
+        raise ValueError(f"Bad bbox: {bbox!r}")
+    return (parts[0], parts[1], parts[2], parts[3])
+
+
+def _coerce_count(value) -> int:
+    if value is None:
+        return 0
+    try:
+        return max(0, int(value))
+    except (TypeError, ValueError):
+        return 0
+
+
+def parse_layout(text: str) -> List[ParsedLayoutBlock]:
+    """Pull the JSON array out of LAYOUT_PROMPT output and convert to typed blocks.
+
+    Tolerates code fences, missing fields, and stringified bboxes.
+    """
+    cleaned = _strip_fences(text)
+    m = _JSON_ARRAY_RE.search(cleaned)
+    if not m:
+        raise ValueError(f"No JSON array found in layout output: {text[:500]!r}")
+    raw = json.loads(m.group(0))
+    out: List[ParsedLayoutBlock] = []
+    for item in raw:
+        try:
+            bbox = _coerce_bbox(item["bbox"])
+        except (KeyError, ValueError) as e:
+            logger.warning(f"Skipping layout block with bad bbox: {e}")
+            continue
+        label = str(item.get("label", "block"))
+        count = _coerce_count(item.get("count"))
+        out.append(ParsedLayoutBlock(label=label, bbox=bbox, count=count))
+    return out
+
+
+# ---- Table rec (TABLE_REC_PROMPT) ------------------------------------------
+
+
+@dataclass
+class ParsedTableElement:
+    label: str  # "Row" or "Col"
+    bbox: Tuple[float, float, float, float]
+
+
+def parse_table_rec(text: str) -> List[ParsedTableElement]:
+    """Parse JSON array of {label: "Row"|"Col", bbox: "x0 y0 x1 y1"} from
+    TABLE_REC_PROMPT output. Returns a flat list of Row + Col elements;
+    cell derivation is the caller's job."""
+    cleaned = _strip_fences(text)
+    m = _JSON_ARRAY_RE.search(cleaned)
+    if not m:
+        raise ValueError(f"No JSON array found in table_rec output: {text[:500]!r}")
+    raw = json.loads(m.group(0))
+    out: List[ParsedTableElement] = []
+    for item in raw:
+        label = str(item.get("label", "")).strip()
+        if label not in ("Row", "Col"):
+            continue
+        try:
+            bbox = _coerce_bbox(item["bbox"])
+        except (KeyError, ValueError):
+            continue
+        out.append(ParsedTableElement(label=label, bbox=bbox))
+    return out
+
+
+# ---- Block HTML (BLOCK_PROMPT for full table path / general block path) ---
+
+
+def clean_block_html(html: str) -> str:
+    """Light cleanup of model-emitted HTML for a single block.
+
+    Strips code fences, leading/trailing whitespace. Does NOT validate against
+    ALLOWED_TAGS — the model is expected to comply, and downstream consumers
+    can sanitize further if needed.
+    """
+    cleaned = _strip_fences(html).strip()
+    return cleaned
+
+
+# ---- Full-page fallback (HIGH_ACCURACY_BBOX_PROMPT) -----------------------
+
+
+@dataclass
+class ParsedFullPageBlock:
+    label: str
+    bbox: Tuple[float, float, float, float]  # 0-1000 normalized
+    html: str  # inner HTML of the wrapping div
+
+
+def parse_full_page_html(text: str) -> List[ParsedFullPageBlock]:
+    """Parse output of HIGH_ACCURACY_BBOX_PROMPT — top-level <div data-bbox=...
+    data-label=...>inner HTML</div> blocks. Returns one entry per top-level div."""
+    from bs4 import BeautifulSoup
+
+    cleaned = _strip_fences(text).strip()
+    if not cleaned:
+        return []
+    # The model outputs a sequence of top-level divs (no surrounding root).
+    # BeautifulSoup parses fine without one.
+    soup = BeautifulSoup(cleaned, "html.parser")
+    divs = soup.find_all("div", recursive=False)
+    out: List[ParsedFullPageBlock] = []
+    for div in divs:
+        label = div.get("data-label")
+        bbox_str = div.get("data-bbox")
+        if not label or not bbox_str:
+            continue
+        try:
+            parts = [float(x) for x in bbox_str.split()]
+        except ValueError:
+            continue
+        if len(parts) != 4:
+            continue
+        # Strip nested data-bbox attrs from the inner HTML so downstream
+        # consumers don't see model debug info on every child element.
+        for tag in div.find_all(attrs={"data-bbox": True}):
+            del tag["data-bbox"]
+        for tag in div.find_all(attrs={"data-label": True}):
+            del tag["data-label"]
+        inner = "".join(str(c) for c in div.contents).strip()
+        out.append(
+            ParsedFullPageBlock(
+                label=str(label),
+                bbox=(parts[0], parts[1], parts[2], parts[3]),
+                html=inner,
+            )
+        )
+    return out
+
+
+def denorm_bbox(bbox, img_w: int, img_h: int, scale: int = 1000):
+    x0, y0, x1, y1 = bbox
+    return (
+        x0 / scale * img_w,
+        y0 / scale * img_h,
+        x1 / scale * img_w,
+        y1 / scale * img_h,
+    )

surya/inference/prompts.py

@@ -0,0 +1,158 @@
+"""Prompt strings for surya2. The exact wording is the model's training-time
+contract — do not paraphrase without retraining."""
+
+from surya.inference.schema import PROMPT_TYPE_BLOCK as PROMPT_TYPE_BLOCK
+from surya.inference.schema import (
+    PROMPT_TYPE_HIGH_ACCURACY_BBOX as PROMPT_TYPE_HIGH_ACCURACY_BBOX,
+)
+from surya.inference.schema import PROMPT_TYPE_LAYOUT as PROMPT_TYPE_LAYOUT
+from surya.inference.schema import PROMPT_TYPE_TABLE_REC as PROMPT_TYPE_TABLE_REC
+
+ALLOWED_TAGS = [
+    "math",
+    "br",
+    "i",
+    "b",
+    "u",
+    "del",
+    "sup",
+    "sub",
+    "table",
+    "tr",
+    "td",
+    "p",
+    "th",
+    "div",
+    "pre",
+    "h1",
+    "h2",
+    "h3",
+    "h4",
+    "h5",
+    "ul",
+    "ol",
+    "li",
+    "input",
+    "a",
+    "span",
+    "img",
+    "hr",
+    "tbody",
+    "small",
+    "caption",
+    "strong",
+    "thead",
+    "big",
+    "code",
+    "chem",
+]
+
+ALLOWED_ATTRIBUTES = [
+    "class",
+    "colspan",
+    "rowspan",
+    "display",
+    "checked",
+    "type",
+    "border",
+    "value",
+    "style",
+    "href",
+    "alt",
+    "align",
+    "data-bbox",
+    "data-label",
+]
+
+# Block labels we don't run OCR on.
+SKIP_OCR_LABELS = {"Figure", "Image", "Diagram", "Blank-Page"}
+
+LAYOUT_PROMPT = (
+    "Output the layout of this image as JSON. Each entry is a dict with "
+    '"label", "bbox", and "count" fields. Bbox is x0 y0 x1 y1, normalized 0-1000.'
+)
+
+BLOCK_PROMPT = "OCR this block image to HTML."
+
+TABLE_REC_PROMPT = (
+    "Output the table rows then columns as JSON. Each entry is a dict with "
+    '"label" ("Row" or "Col") and "bbox" (x0 y0 x1 y1, normalized 0-1000).'
+)
+
+HIGH_ACCURACY_BBOX_PROMPT = (
+    "OCR this image to HTML. Each block is a div with data-label and data-bbox "
+    "(x0 y0 x1 y1, normalized 0-1000)."
+)
+
+
+PROMPT_MAPPING = {
+    "layout": LAYOUT_PROMPT,
+    "block": BLOCK_PROMPT,
+    "table_rec": TABLE_REC_PROMPT,
+    "high_accuracy_bbox": HIGH_ACCURACY_BBOX_PROMPT,
+}
+
+
+# JSON schema for LAYOUT_PROMPT — enforced via vllm guided decoding so the
+# model can't emit malformed JSON. bbox is a "x0 y0 x1 y1" string (model's
+# training-time format); count is a non-negative integer.
+LAYOUT_LABEL_SET = [
+    "Caption",
+    "Footnote",
+    "Equation-Block",
+    "List-Group",
+    "Page-Header",
+    "Page-Footer",
+    "Image",
+    "Section-Header",
+    "Table",
+    "Text",
+    "Complex-Block",
+    "Code-Block",
+    "Form",
+    "Table-Of-Contents",
+    "Figure",
+    "Chemical-Block",
+    "Diagram",
+    "Bibliography",
+    "Blank-Page",
+]
+
+LAYOUT_JSON_SCHEMA = {
+    "type": "array",
+    "maxItems": 200,
+    "items": {
+        "type": "object",
+        "properties": {
+            "label": {"type": "string", "enum": LAYOUT_LABEL_SET},
+            "bbox": {
+                "type": "string",
+                "pattern": r"^\d{1,4} \d{1,4} \d{1,4} \d{1,4}$",
+            },
+            "count": {"type": "integer", "minimum": 0, "maximum": 10000},
+        },
+        "required": ["label", "bbox", "count"],
+        "additionalProperties": False,
+    },
+}
+
+
+# JSON schema for TABLE_REC_PROMPT — array of {label: Row|Col, bbox: "x0 y0 x1 y1"}.
+TABLE_REC_LABEL_SET = ["Row", "Col"]
+
+TABLE_REC_JSON_SCHEMA = {
+    "type": "array",
+    "maxItems": 200,
+    "items": {
+        "type": "object",
+        "properties": {
+            "label": {"type": "string", "enum": TABLE_REC_LABEL_SET},
+            "bbox": {
+                "type": "string",
+                "pattern": r"^\d{1,4} \d{1,4} \d{1,4} \d{1,4}$",
+            },
+        },
+        "required": ["label", "bbox"],
+        "additionalProperties": False,
+    },
+}

surya/inference/schema.py

@@ -0,0 +1,45 @@
+from dataclasses import dataclass, field
+from typing import Any, List, Optional
+
+from PIL import Image
+
+
+PROMPT_TYPE_LAYOUT = "layout"
+PROMPT_TYPE_BLOCK = "block"
+PROMPT_TYPE_TABLE_REC = "table_rec"
+PROMPT_TYPE_HIGH_ACCURACY_BBOX = "high_accuracy_bbox"
+
+
+@dataclass
+class BatchInputItem:
+    image: Image.Image
+    prompt_type: str
+    prompt: Optional[str] = None  # If set, overrides the default prompt for prompt_type
+    max_tokens: Optional[int] = None
+    request_logprobs: bool = False
+    # vllm-native guided decoding — JSON schema, regex, or grammar string.
+    # When set, the server constrains the decode tokens to match the schema.
+    guided_json: Optional[dict] = None
+    guided_regex: Optional[str] = None
+    metadata: dict = field(default_factory=dict)  # Free-form, passes through to output
+
+
+@dataclass
+class GenerationResult:
+    raw: str
+    token_count: int
+    error: bool = False
+    # Mean of exp(logprob) across response tokens, if logprobs requested
+    mean_token_prob: Optional[float] = None
+    # Per-token logprobs (raw OpenAI-style content list), if requested - phase 2 use
+    logprobs: Optional[List[Any]] = None
+
+
+@dataclass
+class BatchOutputItem:
+    raw: str
+    token_count: int
+    error: bool
+    mean_token_prob: Optional[float] = None
+    logprobs: Optional[List[Any]] = None
+    metadata: dict = field(default_factory=dict)

surya/inference/util.py

@@ -0,0 +1,96 @@
+from typing import Tuple
+
+from PIL import Image
+
+
+def scale_to_fit(
+    img: Image.Image,
+    max_size: Tuple[int, int] = (3072, 2048),
+    min_size: Tuple[int, int] = (1792, 28),
+    grid_size: int = 28,
+) -> Image.Image:
+    resample_method = Image.Resampling.LANCZOS
+
+    width, height = img.size
+
+    if width <= 0 or height <= 0:
+        return img
+
+    original_ar = width / height
+    current_pixels = width * height
+    max_pixels = max_size[0] * max_size[1]
+    min_pixels = min_size[0] * min_size[1]
+
+    scale = 1.0
+    if current_pixels > max_pixels:
+        scale = (max_pixels / current_pixels) ** 0.5
+    elif current_pixels < min_pixels:
+        scale = (min_pixels / current_pixels) ** 0.5
+
+    w_blocks = max(1, round((width * scale) / grid_size))
+    h_blocks = max(1, round((height * scale) / grid_size))
+
+    while (w_blocks * h_blocks * grid_size * grid_size) > max_pixels:
+        if w_blocks == 1 and h_blocks == 1:
+            break
+
+        if w_blocks == 1:
+            h_blocks -= 1
+            continue
+        if h_blocks == 1:
+            w_blocks -= 1
+            continue
+
+        ar_w_loss = abs(((w_blocks - 1) / h_blocks) - original_ar)
+        ar_h_loss = abs((w_blocks / (h_blocks - 1)) - original_ar)
+
+        if ar_w_loss < ar_h_loss:
+            w_blocks -= 1
+        else:
+            h_blocks -= 1
+
+    new_width = w_blocks * grid_size
+    new_height = h_blocks * grid_size
+
+    if (new_width, new_height) == (width, height):
+        return img
+
+    return img.resize((new_width, new_height), resample=resample_method)
+
+
+def detect_repeat_token(
+    predicted_tokens: str,
+    base_max_repeats: int = 4,
+    window_size: int = 500,
+    cut_from_end: int = 0,
+    scaling_factor: float = 3.0,
+) -> bool:
+    if cut_from_end > 0:
+        predicted_tokens = predicted_tokens[:-cut_from_end]
+
+    for seq_len in range(1, window_size // 2 + 1):
+        candidate_seq = predicted_tokens[-seq_len:]
+
+        max_repeats = int(base_max_repeats * (1 + scaling_factor / seq_len))
+
+        repeat_count = 0
+        pos = len(predicted_tokens) - seq_len
+        if pos < 0:
+            continue
+
+        while pos >= 0:
+            if predicted_tokens[pos : pos + seq_len] == candidate_seq:
+                repeat_count += 1
+                pos -= seq_len
+            else:
+                break
+
+        if repeat_count > max_repeats:
+            return True
+
+    return False
+
+
+def image_token_budget(block_count: int, ceiling: int = 4096, floor: int = 64) -> int:
+    """Per-block max_tokens: count + 100, clamped to [floor, ceiling]."""
+    return min(max(block_count + 100, floor), ceiling)

surya/input/load.py

@@ -0,0 +1,77 @@
+from typing import List
+import PIL
+
+from surya.input.processing import open_pdf, get_page_images
+from surya.logging import get_logger
+from surya.settings import settings
+import os
+import filetype
+from PIL import Image
+
+logger = get_logger()
+
+
+def get_name_from_path(path):
+    return os.path.basename(path).split(".")[0]
+
+
+def load_pdf(pdf_path, page_range: List[int] | None = None, dpi=settings.IMAGE_DPI):
+    doc = open_pdf(pdf_path)
+    last_page = len(doc)
+
+    if page_range:
+        assert all([0 <= page < last_page for page in page_range]), (
+            f"Invalid page range: {page_range}"
+        )
+    else:
+        page_range = list(range(last_page))
+
+    images = get_page_images(doc, page_range, dpi=dpi)
+    doc.close()
+    names = [get_name_from_path(pdf_path) for _ in page_range]
+    return images, names
+
+
+def load_image(image_path):
+    image = Image.open(image_path).convert("RGB")
+    name = get_name_from_path(image_path)
+    return [image], [name]
+
+
+def load_from_file(
+    input_path, page_range: List[int] | None = None, dpi=settings.IMAGE_DPI
+):
+    input_type = filetype.guess(input_path)
+    if input_type and input_type.extension == "pdf":
+        return load_pdf(input_path, page_range, dpi=dpi)
+    else:
+        return load_image(input_path)
+
+
+def load_from_folder(
+    folder_path, page_range: List[int] | None = None, dpi=settings.IMAGE_DPI
+):
+    image_paths = [
+        os.path.join(folder_path, image_name)
+        for image_name in os.listdir(folder_path)
+        if not image_name.startswith(".")
+    ]
+    image_paths = [ip for ip in image_paths if not os.path.isdir(ip)]
+
+    images = []
+    names = []
+    for path in image_paths:
+        extension = filetype.guess(path)
+        if extension and extension.extension == "pdf":
+            image, name = load_pdf(path, page_range, dpi=dpi)
+            images.extend(image)
+            names.extend(name)
+        else:
+            try:
+                image, name = load_image(path)
+                images.extend(image)
+                names.extend(name)
+            except PIL.UnidentifiedImageError:
+                logger.warning(f"Could not load image {path}")
+                continue
+    return images, names

surya/input/processing.py

@@ -0,0 +1,17 @@
+from typing import List
+
+import pypdfium2
+
+from surya.settings import settings
+
+
+def open_pdf(pdf_filepath):
+    return pypdfium2.PdfDocument(pdf_filepath)
+
+
+def get_page_images(doc, indices: List, dpi=settings.IMAGE_DPI):
+    images = [
+        doc[i].render(scale=dpi / 72, draw_annots=False).to_pil() for i in indices
+    ]
+    images = [image.convert("RGB") for image in images]
+    return images

surya/layout/__init__.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+from typing import List, Optional
+
+from PIL import Image
+
+from surya.common.blank import is_blank_region
+from surya.inference import SuryaInferenceManager, get_default_manager
+from surya.inference.parsers import denorm_bbox, parse_layout
+from surya.inference.prompts import LAYOUT_JSON_SCHEMA, PROMPT_TYPE_LAYOUT
+from surya.inference.schema import BatchInputItem
+from surya.layout.label import LAYOUT_PRED_RELABEL, TEXT_LABELS
+from surya.layout.schema import LayoutBox, LayoutResult
+from surya.logging import get_logger
+from surya.settings import settings
+from surya.timing import timing_span
+
+logger = get_logger()
+
+
+class LayoutPredictor:
+    """Run LAYOUT_PROMPT on full pages, parse JSON, return LayoutResult per image."""
+
+    def __init__(self, manager: Optional[SuryaInferenceManager] = None):
+        self.manager = manager  # If None, get_default_manager() is used at call time
+        self._disable_tqdm = settings.DISABLE_TQDM
+
+    @property
+    def disable_tqdm(self) -> bool:
+        return self._disable_tqdm
+
+    @disable_tqdm.setter
+    def disable_tqdm(self, value: bool) -> None:
+        self._disable_tqdm = bool(value)
+
+    def to(self, *args, **kwargs):
+        # Manager-backed; .to() is a no-op for compatibility with BasePredictor callers.
+        return
+
+    def __call__(
+        self,
+        images: List[Image.Image],
+        target_image_sizes: Optional[List[tuple]] = None,
+        max_tokens: Optional[int] = None,
+    ) -> List[LayoutResult]:
+        """Run layout on a batch of images.
+
+        target_image_sizes: optional list of (width, height) tuples — if
+        provided, bboxes are denormalized to these sizes instead of each
+        input image's size. Useful when layout runs on a low-DPI render but
+        you want bboxes in the OCR image's coordinate space.
+        """
+        if not images:
+            return []
+        manager = self.manager or get_default_manager()
+
+        max_tokens = max_tokens or settings.SURYA_MAX_TOKENS_LAYOUT
+        guided = LAYOUT_JSON_SCHEMA if settings.SURYA_GUIDED_LAYOUT else None
+        with timing_span("layout_build_batch", image_count=len(images), max_tokens=max_tokens):
+            batch = [
+                BatchInputItem(
+                    image=img,
+                    prompt_type=PROMPT_TYPE_LAYOUT,
+                    max_tokens=max_tokens,
+                    guided_json=guided,
+                )
+                for img in images
+            ]
+        with timing_span("layout_manager_generate", item_count=len(batch)):
+            outputs = manager.generate(batch)
+
+        if target_image_sizes is not None and len(target_image_sizes) != len(images):
+            raise ValueError("target_image_sizes must match images length")
+
+        with timing_span("layout_parse_outputs", item_count=len(outputs)):
+            results: List[LayoutResult] = []
+            for idx, (img, out) in enumerate(zip(images, outputs)):
+                if target_image_sizes is not None:
+                    w, h = target_image_sizes[idx]
+                else:
+                    w, h = img.size
+                page_bbox = [0, 0, float(w), float(h)]
+                if out.error or not out.raw:
+                    results.append(
+                        LayoutResult(
+                            bboxes=[], image_bbox=page_bbox, raw=out.raw, error=True
+                        )
+                    )
+                    continue
+                try:
+                    parsed = parse_layout(out.raw)
+                except Exception as e:
+                    logger.warning(f"Layout parse failed: {e}; raw[:300]={out.raw[:300]!r}")
+                    results.append(
+                        LayoutResult(
+                            bboxes=[], image_bbox=page_bbox, raw=out.raw, error=True
+                        )
+                    )
+                    continue
+
+                confidence = out.mean_token_prob if out.mean_token_prob is not None else 1.0
+                img_w, img_h = img.size
+                boxes: List[LayoutBox] = []
+                dropped_blank = 0
+                for blk in parsed:
+                    canon = LAYOUT_PRED_RELABEL.get(blk.label, blk.label)
+                    # Drop text-labeled blocks the model hallucinated over an
+                    # essentially-blank region (mostly white OR near-uniform
+                    # color). Visual blocks (Picture / Figure / Table / etc.)
+                    # are allowed to be uniform — that's normal content.
+                    if canon in TEXT_LABELS:
+                        img_bbox = denorm_bbox(
+                            blk.bbox, img_w, img_h, scale=settings.BBOX_SCALE
+                        )
+                        x0, y0, x1, y1 = (max(0, int(v)) for v in img_bbox)
+                        if x1 > x0 and y1 > y0:
+                            if is_blank_region(img.crop((x0, y0, x1, y1))):
+                                dropped_blank += 1
+                                continue
+                    pixel_bbox = denorm_bbox(blk.bbox, w, h, scale=settings.BBOX_SCALE)
+                    boxes.append(
+                        LayoutBox(
+                            polygon=list(pixel_bbox),
+                            label=canon,
+                            raw_label=blk.label,
+                            position=len(boxes),
+                            count=blk.count,
+                            confidence=confidence,
+                        )
+                    )
+                if dropped_blank:
+                    logger.info(
+                        f"dropped {dropped_blank} text-labeled layout block(s) over "
+                        f"blank/uniform regions"
+                    )
+                results.append(
+                    LayoutResult(
+                        bboxes=boxes, image_bbox=page_bbox, raw=out.raw, error=False
+                    )
+                )
+        return results

surya/layout/label.py

@@ -0,0 +1,44 @@
+"""Surya2 layout labels emitted by the model + canonicalization to surya's
+public label vocabulary."""
+
+# Canonical text-bearing labels — used by blank-region filters to decide
+# which blocks may be dropped when their underlying image region is empty.
+# Excludes Picture/Figure/Diagram/Table/Form/Equation/etc., which can legitimately
+# contain whitespace or solid fills.
+TEXT_LABELS = frozenset(
+    {
+        "Text",
+        "SectionHeader",
+        "PageHeader",
+        "PageFooter",
+        "Caption",
+        "Footnote",
+        "Code",
+        "Bibliography",
+    }
+)
+
+
+# Canonicalize raw model labels to public surya label names. Marker and other
+# downstream consumers depend on these names.
+LAYOUT_PRED_RELABEL = {
+    "Caption": "Caption",
+    "Footnote": "Footnote",
+    "Equation-Block": "Equation",
+    "List-Group": "ListGroup",
+    "Page-Header": "PageHeader",
+    "Page-Footer": "PageFooter",
+    "Image": "Picture",
+    "Section-Header": "SectionHeader",
+    "Table": "Table",
+    "Text": "Text",
+    "Complex-Block": "Figure",
+    "Code-Block": "Code",
+    "Form": "Form",
+    "Table-Of-Contents": "TableOfContents",
+    "Figure": "Figure",
+    "Chemical-Block": "ChemicalBlock",
+    "Diagram": "Diagram",
+    "Bibliography": "Bibliography",
+    "Blank-Page": "BlankPage",
+}

surya/layout/schema.py

@@ -0,0 +1,19 @@
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+from surya.common.polygon import PolygonBox
+
+
+class LayoutBox(PolygonBox):
+    label: str  # canonicalized via LAYOUT_PRED_RELABEL
+    raw_label: str  # original model label, before canonicalization
+    position: int  # reading order index
+    count: int = 0  # model's token estimate for OCR output (multiple of 50)
+
+
+class LayoutResult(BaseModel):
+    bboxes: List[LayoutBox]
+    image_bbox: List[float]
+    raw: Optional[str] = None  # raw model output, useful for debugging
+    error: bool = False

surya/logging.py

@@ -0,0 +1,27 @@
+import logging
+import warnings
+from surya.settings import settings
+
+
+def configure_logging():
+    logger = get_logger()
+
+    # Remove any existing handlers to prevent duplicates
+    for handler in logger.handlers[:]:
+        logger.removeHandler(handler)
+
+    # Add our handler
+    handler = logging.StreamHandler()
+    formatter = logging.Formatter("%(asctime)s [%(levelname)s] %(name)s: %(message)s")
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+
+    # Prevent propagation to parent loggers to avoid double logging
+    logger.propagate = False
+
+    logger.setLevel(settings.LOGLEVEL)
+    warnings.simplefilter(action="ignore", category=FutureWarning)
+
+
+def get_logger():
+    return logging.getLogger("surya")

surya/ocr_error/__init__.py

@@ -0,0 +1,52 @@
+import math
+from typing import List, Optional
+
+from tqdm import tqdm
+
+from surya.common.predictor import BasePredictor
+from surya.ocr_error.loader import OCRErrorModelLoader
+from surya.ocr_error.model.config import ID2LABEL
+from surya.ocr_error.schema import OCRErrorDetectionResult
+from surya.settings import settings
+
+
+class OCRErrorPredictor(BasePredictor):
+    model_loader_cls = OCRErrorModelLoader
+    batch_size = settings.OCR_ERROR_BATCH_SIZE
+    default_batch_sizes = {"cpu": 8, "mps": 8, "cuda": 64}
+
+    def __call__(self, texts: List[str], batch_size: Optional[int] = None):
+        return self.batch_ocr_error_detection(texts, batch_size)
+
+    def batch_ocr_error_detection(
+        self, texts: List[str], batch_size: Optional[int] = None
+    ):
+        if batch_size is None:
+            batch_size = self.get_batch_size()
+
+        num_batches = math.ceil(len(texts) / batch_size)
+        texts_processed = self.processor(
+            texts, padding="longest", truncation=True, return_tensors="pt"
+        )
+        predictions = []
+        for batch_idx in tqdm(
+            range(num_batches),
+            desc="Running OCR Error Detection",
+            disable=self.disable_tqdm,
+        ):
+            start_idx, end_idx = batch_idx * batch_size, (batch_idx + 1) * batch_size
+            batch_input_ids = texts_processed.input_ids[start_idx:end_idx].to(
+                self.model.device
+            )
+            batch_attention_mask = texts_processed.attention_mask[start_idx:end_idx].to(
+                self.model.device
+            )
+
+            with settings.INFERENCE_MODE():
+                pred = self.model(batch_input_ids, attention_mask=batch_attention_mask)
+                logits = pred.logits.argmax(dim=1).cpu().tolist()
+                predictions.extend(logits)
+
+        return OCRErrorDetectionResult(
+            texts=texts, labels=[ID2LABEL[p] for p in predictions]
+        )

surya/ocr_error/loader.py

@@ -0,0 +1,48 @@
+from typing import Optional
+
+
+from surya.common.load import ModelLoader
+from surya.logging import get_logger
+from surya.ocr_error.model.config import DistilBertConfig
+from surya.ocr_error.model.encoder import DistilBertForSequenceClassification
+from surya.ocr_error.tokenizer import DistilBertTokenizer
+from surya.settings import settings
+
+logger = get_logger()
+
+
+class OCRErrorModelLoader(ModelLoader):
+    def __init__(self, checkpoint: Optional[str] = None):
+        super().__init__(checkpoint)
+
+        if self.checkpoint is None:
+            self.checkpoint = settings.OCR_ERROR_MODEL_CHECKPOINT
+
+    def model(
+        self,
+        device=settings.TORCH_DEVICE_MODEL,
+        dtype=settings.MODEL_DTYPE,
+        attention_implementation: Optional[str] = None,
+    ) -> DistilBertForSequenceClassification:
+        if device is None:
+            device = settings.TORCH_DEVICE_MODEL
+        if dtype is None:
+            dtype = settings.MODEL_DTYPE
+
+        config = DistilBertConfig.from_pretrained(self.checkpoint)
+        model = (
+            DistilBertForSequenceClassification.from_pretrained(
+                self.checkpoint,
+                dtype=dtype,
+                config=config,
+            )
+            .to(device)
+            .eval()
+        )
+
+        return model
+
+    def processor(
+        self, device=settings.TORCH_DEVICE_MODEL, dtype=settings.MODEL_DTYPE
+    ) -> DistilBertTokenizer:
+        return DistilBertTokenizer.from_pretrained(self.checkpoint)

surya/ocr_error/model/config.py

@@ -0,0 +1,68 @@
+from collections import OrderedDict
+from typing import Mapping
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.onnx import OnnxConfig
+
+from surya.common.s3 import S3DownloaderMixin
+
+ID2LABEL = {
+    0: 'good',
+    1: 'bad'
+}
+
+class DistilBertConfig(S3DownloaderMixin, PretrainedConfig):
+    model_type = "distilbert"
+    attribute_map = {
+        "hidden_size": "dim",
+        "num_attention_heads": "n_heads",
+        "num_hidden_layers": "n_layers",
+    }
+
+    def __init__(
+        self,
+        vocab_size=30522,
+        max_position_embeddings=512,
+        sinusoidal_pos_embds=False,
+        n_layers=6,
+        n_heads=12,
+        dim=768,
+        hidden_dim=4 * 768,
+        dropout=0.1,
+        attention_dropout=0.1,
+        activation="gelu",
+        initializer_range=0.02,
+        qa_dropout=0.1,
+        seq_classif_dropout=0.2,
+        pad_token_id=0,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.sinusoidal_pos_embds = sinusoidal_pos_embds
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.dim = dim
+        self.hidden_dim = hidden_dim
+        self.dropout = dropout
+        self.attention_dropout = attention_dropout
+        self.activation = activation
+        self.initializer_range = initializer_range
+        self.qa_dropout = qa_dropout
+        self.seq_classif_dropout = seq_classif_dropout
+        super().__init__(**kwargs, pad_token_id=pad_token_id)
+
+
+class DistilBertOnnxConfig(OnnxConfig):
+    @property
+    def inputs(self) -> Mapping[str, Mapping[int, str]]:
+        if self.task == "multiple-choice":
+            dynamic_axis = {0: "batch", 1: "choice", 2: "sequence"}
+        else:
+            dynamic_axis = {0: "batch", 1: "sequence"}
+        return OrderedDict(
+            [
+                ("input_ids", dynamic_axis),
+                ("attention_mask", dynamic_axis),
+            ]
+        )

surya/ocr_error/model/encoder.py

@@ -0,0 +1,910 @@
+from __future__ import annotations
+
+import math
+from typing import Optional, Set, List, Tuple, Union, Dict
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F, MSELoss, CrossEntropyLoss, BCEWithLogitsLoss
+from transformers import apply_chunking_to_forward
+from transformers.activations import get_activation
+from transformers.modeling_outputs import BaseModelOutput, SequenceClassifierOutput
+from transformers.pytorch_utils import (
+    find_pruneable_heads_and_indices,
+    prune_linear_layer,
+)
+
+from transformers.utils import (
+    is_flash_attn_greater_or_equal_2_10,
+)
+
+from surya.common.pretrained import SuryaPreTrainedModel
+
+from surya.common.s3 import S3DownloaderMixin
+from surya.ocr_error.model.config import DistilBertConfig
+
+
+def _get_unpad_data(attention_mask):
+    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
+    return (
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+    )
+
+
+def create_sinusoidal_embeddings(n_pos: int, dim: int, out: torch.Tensor):
+    position_enc = np.array(
+        [
+            [pos / np.power(10000, 2 * (j // 2) / dim) for j in range(dim)]
+            for pos in range(n_pos)
+        ]
+    )
+    out.requires_grad = False
+    out[:, 0::2] = torch.FloatTensor(np.sin(position_enc[:, 0::2]))
+    out[:, 1::2] = torch.FloatTensor(np.cos(position_enc[:, 1::2]))
+    out.detach_()
+
+
+class Embeddings(nn.Module):
+    def __init__(self, config: DistilBertConfig):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(
+            config.vocab_size, config.dim, padding_idx=config.pad_token_id
+        )
+        self.position_embeddings = nn.Embedding(
+            config.max_position_embeddings, config.dim
+        )
+
+        self.LayerNorm = nn.LayerNorm(config.dim, eps=1e-12)
+        self.dropout = nn.Dropout(config.dropout)
+        self.register_buffer(
+            "position_ids",
+            torch.arange(config.max_position_embeddings).expand((1, -1)),
+            persistent=False,
+        )
+
+    def forward(
+        self, input_ids: torch.Tensor, input_embeds: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Parameters:
+            input_ids (torch.Tensor):
+                torch.tensor(bs, max_seq_length) The token ids to embed.
+            input_embeds (*optional*, torch.Tensor):
+                The pre-computed word embeddings. Can only be passed if the input ids are `None`.
+
+
+        Returns: torch.tensor(bs, max_seq_length, dim) The embedded tokens (plus position embeddings, no token_type
+        embeddings)
+        """
+        if input_ids is not None:
+            input_embeds = self.word_embeddings(input_ids)  # (bs, max_seq_length, dim)
+
+        seq_length = input_embeds.size(1)
+
+        # Setting the position-ids to the registered buffer in constructor, it helps
+        # when tracing the model without passing position-ids, solves
+        # isues similar to issue #5664
+        if hasattr(self, "position_ids"):
+            position_ids = self.position_ids[:, :seq_length]
+        else:
+            position_ids = torch.arange(
+                seq_length, dtype=torch.long, device=input_ids.device
+            )  # (max_seq_length)
+            position_ids = position_ids.unsqueeze(0).expand_as(
+                input_ids
+            )  # (bs, max_seq_length)
+
+        position_embeddings = self.position_embeddings(
+            position_ids
+        )  # (bs, max_seq_length, dim)
+
+        embeddings = input_embeds + position_embeddings  # (bs, max_seq_length, dim)
+        embeddings = self.LayerNorm(embeddings)  # (bs, max_seq_length, dim)
+        embeddings = self.dropout(embeddings)  # (bs, max_seq_length, dim)
+        return embeddings
+
+
+class MultiHeadSelfAttention(nn.Module):
+    def __init__(self, config: DistilBertConfig):
+        super().__init__()
+        self.config = config
+
+        self.n_heads = config.n_heads
+        self.dim = config.dim
+        self.dropout = nn.Dropout(p=config.attention_dropout)
+        self.is_causal = False
+
+        # Have an even number of multi heads that divide the dimensions
+        if self.dim % self.n_heads != 0:
+            # Raise value errors for even multi-head attention nodes
+            raise ValueError(
+                f"self.n_heads: {self.n_heads} must divide self.dim: {self.dim} evenly"
+            )
+
+        self.q_lin = nn.Linear(in_features=config.dim, out_features=config.dim)
+        self.k_lin = nn.Linear(in_features=config.dim, out_features=config.dim)
+        self.v_lin = nn.Linear(in_features=config.dim, out_features=config.dim)
+        self.out_lin = nn.Linear(in_features=config.dim, out_features=config.dim)
+
+        self.pruned_heads: Set[int] = set()
+        self.attention_head_size = self.dim // self.n_heads
+
+    def prune_heads(self, heads: List[int]):
+        if len(heads) == 0:
+            return
+        heads, index = find_pruneable_heads_and_indices(
+            heads, self.n_heads, self.attention_head_size, self.pruned_heads
+        )
+        # Prune linear layers
+        self.q_lin = prune_linear_layer(self.q_lin, index)
+        self.k_lin = prune_linear_layer(self.k_lin, index)
+        self.v_lin = prune_linear_layer(self.v_lin, index)
+        self.out_lin = prune_linear_layer(self.out_lin, index, dim=1)
+        # Update hyper params
+        self.n_heads = self.n_heads - len(heads)
+        self.dim = self.attention_head_size * self.n_heads
+        self.pruned_heads = self.pruned_heads.union(heads)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        mask: torch.Tensor,
+        head_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, ...]:
+        """
+        Parameters:
+            query: torch.tensor(bs, seq_length, dim)
+            key: torch.tensor(bs, seq_length, dim)
+            value: torch.tensor(bs, seq_length, dim)
+            mask: torch.tensor(bs, seq_length)
+
+        Returns:
+            weights: torch.tensor(bs, n_heads, seq_length, seq_length) Attention weights context: torch.tensor(bs,
+            seq_length, dim) Contextualized layer. Optional: only if `output_attentions=True`
+        """
+        bs, q_length, dim = query.size()
+        k_length = key.size(1)
+        # assert dim == self.dim, f'Dimensions do not match: {dim} input vs {self.dim} configured'
+        # assert key.size() == value.size()
+
+        dim_per_head = self.dim // self.n_heads
+
+        mask_reshp = (bs, 1, 1, k_length)
+
+        def shape(x: torch.Tensor) -> torch.Tensor:
+            """separate heads"""
+            return x.view(bs, -1, self.n_heads, dim_per_head).transpose(1, 2)
+
+        def unshape(x: torch.Tensor) -> torch.Tensor:
+            """group heads"""
+            return (
+                x.transpose(1, 2).contiguous().view(bs, -1, self.n_heads * dim_per_head)
+            )
+
+        q = shape(self.q_lin(query))  # (bs, n_heads, q_length, dim_per_head)
+        k = shape(self.k_lin(key))  # (bs, n_heads, k_length, dim_per_head)
+        v = shape(self.v_lin(value))  # (bs, n_heads, k_length, dim_per_head)
+
+        q = q / math.sqrt(dim_per_head)  # (bs, n_heads, q_length, dim_per_head)
+        scores = torch.matmul(q, k.transpose(2, 3))  # (bs, n_heads, q_length, k_length)
+        mask = (
+            (mask == 0).view(mask_reshp).expand_as(scores)
+        )  # (bs, n_heads, q_length, k_length)
+        scores = scores.masked_fill(
+            mask, torch.tensor(torch.finfo(scores.dtype).min)
+        )  # (bs, n_heads, q_length, k_length)
+
+        weights = nn.functional.softmax(
+            scores, dim=-1
+        )  # (bs, n_heads, q_length, k_length)
+        weights = self.dropout(weights)  # (bs, n_heads, q_length, k_length)
+
+        # Mask heads if we want to
+        if head_mask is not None:
+            weights = weights * head_mask
+
+        context = torch.matmul(weights, v)  # (bs, n_heads, q_length, dim_per_head)
+        context = unshape(context)  # (bs, q_length, dim)
+        context = self.out_lin(context)  # (bs, q_length, dim)
+
+        if output_attentions:
+            return (context, weights)
+        else:
+            return (context,)
+
+
+class DistilBertFlashAttention2(MultiHeadSelfAttention):
+    """
+    DistilBert flash attention module. This module inherits from `MultiHeadSelfAttention` as the weights of the module
+    stays untouched. The only required change would be on the forward pass where it needs to correctly call the public
+    API of flash attention and deal with padding tokens in case the input contains any of them.
+    """
+
+    # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
+        # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
+        # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
+        self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        mask: torch.Tensor,
+        head_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, ...]:
+        """
+        Parameters:
+            query: torch.tensor(bs, seq_length, dim)
+            key: torch.tensor(bs, seq_length, dim)
+            value: torch.tensor(bs, seq_length, dim)
+            mask: torch.tensor(bs, seq_length)
+
+        Returns:
+            weights: torch.tensor(bs, n_heads, seq_length, seq_length) Attention weights context: torch.tensor(bs,
+            seq_length, dim) Contextualized layer. Optional: only if `output_attentions=True`
+        """
+        batch_size, q_length, dim = query.size()
+
+        dim_per_head = self.dim // self.n_heads
+
+        def reshape(x: torch.Tensor) -> torch.Tensor:
+            """separate heads"""
+            return x.view(batch_size, -1, self.n_heads, dim_per_head)
+
+        # Flash attention requires the input to have the shape
+        # batch_size x seq_length x head_dim x hidden_dim
+        query_states = reshape(self.q_lin(query))
+        key_states = reshape(self.k_lin(key))
+        value_states = reshape(self.v_lin(value))
+
+        attn_dropout = self.config.attention_dropout if self.training else 0.0
+
+        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+        # therefore the input hidden states gets silently casted in float32. Hence, we need
+        # cast them back in the correct dtype just to be sure everything works as expected.
+        # This might slowdown training & inference so it is recommended to not cast the LayerNorms
+        # in fp32. (LlamaRMSNorm handles it correctly)
+
+        if query_states.dtype == torch.float32:
+            if torch.is_autocast_enabled():
+                target_dtype = torch.get_autocast_gpu_dtype()
+            # Handle the case where the model is quantized
+            elif hasattr(self.config, "_pre_quantization_dtype"):
+                target_dtype = self.config._pre_quantization_dtype
+            else:
+                target_dtype = self.q_lin.weight.dtype
+
+            query_states = query_states.to(target_dtype)
+            key_states = key_states.to(target_dtype)
+            value_states = value_states.to(target_dtype)
+
+        attn_weights = self._flash_attention_forward(
+            query_states, key_states, value_states, mask, q_length, dropout=attn_dropout
+        )
+
+        attn_weights_reshaped = attn_weights.reshape(
+            batch_size, q_length, self.n_heads * dim_per_head
+        )
+        attn_output = self.out_lin(attn_weights_reshaped)
+
+        if output_attentions:
+            return (attn_output, attn_weights)
+        else:
+            return (attn_output,)
+
+    # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward with causal=True->causal=False
+    def _flash_attention_forward(
+        self,
+        query_states,
+        key_states,
+        value_states,
+        attention_mask,
+        query_length,
+        dropout=0.0,
+        softmax_scale=None,
+    ):
+        """
+        Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
+        first unpad the input, then computes the attention scores and pad the final attention scores.
+
+        Args:
+            query_states (`torch.Tensor`):
+                Input query states to be passed to Flash Attention API
+            key_states (`torch.Tensor`):
+                Input key states to be passed to Flash Attention API
+            value_states (`torch.Tensor`):
+                Input value states to be passed to Flash Attention API
+            attention_mask (`torch.Tensor`):
+                The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
+                position of padding tokens and 1 for the position of non-padding tokens.
+            dropout (`float`):
+                Attention dropout
+            softmax_scale (`float`, *optional*):
+                The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
+        """
+        from flash_attn import flash_attn_func, flash_attn_varlen_func
+        from flash_attn.bert_padding import pad_input
+
+        if not self._flash_attn_uses_top_left_mask:
+            causal = self.is_causal
+        else:
+            # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
+            causal = self.is_causal and query_length != 1
+
+        # Contains at least one padding token in the sequence
+        if attention_mask is not None:
+            batch_size = query_states.shape[0]
+            (
+                query_states,
+                key_states,
+                value_states,
+                indices_q,
+                cu_seq_lens,
+                max_seq_lens,
+            ) = self._upad_input(
+                query_states, key_states, value_states, attention_mask, query_length
+            )
+
+            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+            max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+
+            attn_output_unpad = flash_attn_varlen_func(
+                query_states,
+                key_states,
+                value_states,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_in_batch_q,
+                max_seqlen_k=max_seqlen_in_batch_k,
+                dropout_p=dropout,
+                softmax_scale=softmax_scale,
+                causal=causal,
+            )
+
+            attn_output = pad_input(
+                attn_output_unpad, indices_q, batch_size, query_length
+            )
+        else:
+            attn_output = flash_attn_func(
+                query_states,
+                key_states,
+                value_states,
+                dropout,
+                softmax_scale=softmax_scale,
+                causal=causal,
+            )
+
+        return attn_output
+
+    # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input with num_heads->n_heads
+    def _upad_input(
+        self, query_layer, key_layer, value_layer, attention_mask, query_length
+    ):
+        from flash_attn.bert_padding import index_first_axis, unpad_input
+
+        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
+        batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
+
+        key_layer = index_first_axis(
+            key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim),
+            indices_k,
+        )
+        value_layer = index_first_axis(
+            value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim),
+            indices_k,
+        )
+        if query_length == kv_seq_len:
+            query_layer = index_first_axis(
+                query_layer.reshape(batch_size * kv_seq_len, self.n_heads, head_dim),
+                indices_k,
+            )
+            cu_seqlens_q = cu_seqlens_k
+            max_seqlen_in_batch_q = max_seqlen_in_batch_k
+            indices_q = indices_k
+        elif query_length == 1:
+            max_seqlen_in_batch_q = 1
+            cu_seqlens_q = torch.arange(
+                batch_size + 1, dtype=torch.int32, device=query_layer.device
+            )  # There is a memcpy here, that is very bad.
+            indices_q = cu_seqlens_q[:-1]
+            query_layer = query_layer.squeeze(1)
+        else:
+            # The -q_len: slice assumes left padding.
+            attention_mask = attention_mask[:, -query_length:]
+            query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(
+                query_layer, attention_mask
+            )
+
+        return (
+            query_layer,
+            key_layer,
+            value_layer,
+            indices_q,
+            (cu_seqlens_q, cu_seqlens_k),
+            (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+        )
+
+
+class FFN(nn.Module):
+    def __init__(self, config: DistilBertConfig):
+        super().__init__()
+        self.dropout = nn.Dropout(p=config.dropout)
+        self.chunk_size_feed_forward = config.chunk_size_feed_forward
+        self.seq_len_dim = 1
+        self.lin1 = nn.Linear(in_features=config.dim, out_features=config.hidden_dim)
+        self.lin2 = nn.Linear(in_features=config.hidden_dim, out_features=config.dim)
+        self.activation = get_activation(config.activation)
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return apply_chunking_to_forward(
+            self.ff_chunk, self.chunk_size_feed_forward, self.seq_len_dim, input
+        )
+
+    def ff_chunk(self, input: torch.Tensor) -> torch.Tensor:
+        x = self.lin1(input)
+        x = self.activation(x)
+        x = self.lin2(x)
+        x = self.dropout(x)
+        return x
+
+
+DISTILBERT_ATTENTION_CLASSES = {
+    "eager": MultiHeadSelfAttention,
+    "flash_attention_2": DistilBertFlashAttention2,
+}
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: DistilBertConfig):
+        super().__init__()
+
+        # Have an even number of Configure multi-heads
+        if config.dim % config.n_heads != 0:
+            raise ValueError(
+                f"config.n_heads {config.n_heads} must divide config.dim {config.dim} evenly"
+            )
+
+        self.attention = DISTILBERT_ATTENTION_CLASSES[config._attn_implementation](
+            config
+        )
+        self.sa_layer_norm = nn.LayerNorm(normalized_shape=config.dim, eps=1e-12)
+
+        self.ffn = FFN(config)
+        self.output_layer_norm = nn.LayerNorm(normalized_shape=config.dim, eps=1e-12)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, ...]:
+        """
+        Parameters:
+            x: torch.tensor(bs, seq_length, dim)
+            attn_mask: torch.tensor(bs, seq_length)
+
+        Returns:
+            sa_weights: torch.tensor(bs, n_heads, seq_length, seq_length) The attention weights ffn_output:
+            torch.tensor(bs, seq_length, dim) The output of the transformer block contextualization.
+        """
+        # Self-Attention
+        sa_output = self.attention(
+            query=x,
+            key=x,
+            value=x,
+            mask=attn_mask,
+            head_mask=head_mask,
+            output_attentions=output_attentions,
+        )
+        if output_attentions:
+            sa_output, sa_weights = (
+                sa_output  # (bs, seq_length, dim), (bs, n_heads, seq_length, seq_length)
+            )
+        else:  # To handle these `output_attentions` or `output_hidden_states` cases returning tuples
+            sa_output = sa_output[0]
+
+        sa_output = self.sa_layer_norm(sa_output + x)  # (bs, seq_length, dim)
+
+        # Feed Forward Network
+        ffn_output = self.ffn(sa_output)  # (bs, seq_length, dim)
+        ffn_output: torch.Tensor = self.output_layer_norm(
+            ffn_output + sa_output
+        )  # (bs, seq_length, dim)
+
+        output = (ffn_output,)
+        if output_attentions:
+            output = (sa_weights,) + output
+        return output
+
+
+class Transformer(nn.Module):
+    def __init__(self, config: DistilBertConfig):
+        super().__init__()
+        self.n_layers = config.n_layers
+        self.layer = nn.ModuleList(
+            [TransformerBlock(config) for _ in range(config.n_layers)]
+        )
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+        return_dict: Optional[bool] = None,
+    ) -> Union[BaseModelOutput, Tuple[torch.Tensor, ...]]:  # docstyle-ignore
+        """
+        Parameters:
+            x: torch.tensor(bs, seq_length, dim) Input sequence embedded.
+            attn_mask: torch.tensor(bs, seq_length) Attention mask on the sequence.
+
+        Returns:
+            hidden_state: torch.tensor(bs, seq_length, dim) Sequence of hidden states in the last (top)
+            layer all_hidden_states: Tuple[torch.tensor(bs, seq_length, dim)]
+                Tuple of length n_layers with the hidden states from each layer.
+                Optional: only if output_hidden_states=True
+            all_attentions: Tuple[torch.tensor(bs, n_heads, seq_length, seq_length)]
+                Tuple of length n_layers with the attention weights from each layer
+                Optional: only if output_attentions=True
+        """
+        all_hidden_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        hidden_state = x
+        for i, layer_module in enumerate(self.layer):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_state,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    layer_module.__call__,
+                    hidden_state,
+                    attn_mask,
+                    head_mask[i],
+                    output_attentions,
+                )
+            else:
+                layer_outputs = layer_module(
+                    hidden_state,
+                    attn_mask,
+                    head_mask[i],
+                    output_attentions,
+                )
+
+            hidden_state = layer_outputs[-1]
+
+            if output_attentions:
+                if len(layer_outputs) != 2:
+                    raise ValueError(
+                        f"The length of the layer_outputs should be 2, but it is {len(layer_outputs)}"
+                    )
+
+                attentions = layer_outputs[0]
+                all_attentions = all_attentions + (attentions,)
+            else:
+                if len(layer_outputs) != 1:
+                    raise ValueError(
+                        f"The length of the layer_outputs should be 1, but it is {len(layer_outputs)}"
+                    )
+
+        # Add last layer
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_state,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [hidden_state, all_hidden_states, all_attentions]
+                if v is not None
+            )
+        return BaseModelOutput(
+            last_hidden_state=hidden_state,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions,
+        )
+
+
+class DistilBertPreTrainedModel(SuryaPreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = DistilBertConfig
+    load_tf_weights = None
+    base_model_prefix = "distilbert"
+    supports_gradient_checkpointing = True
+    _supports_flash_attn_2 = True
+
+    def _init_weights(self, module: nn.Module):
+        """Initialize the weights."""
+        if isinstance(module, nn.Linear):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        elif isinstance(module, Embeddings) and self.config.sinusoidal_pos_embds:
+            create_sinusoidal_embeddings(
+                self.config.max_position_embeddings,
+                self.config.dim,
+                module.position_embeddings.weight,
+            )
+
+
+class DistilBertModel(DistilBertPreTrainedModel):
+    def __init__(self, config: DistilBertConfig):
+        super().__init__(config)
+
+        self.embeddings = Embeddings(config)  # Embeddings
+        self.transformer = Transformer(config)  # Encoder
+        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_position_embeddings(self) -> nn.Embedding:
+        """
+        Returns the position embeddings
+        """
+        return self.embeddings.position_embeddings
+
+    def resize_position_embeddings(self, new_num_position_embeddings: int):
+        """
+        Resizes position embeddings of the model if `new_num_position_embeddings != config.max_position_embeddings`.
+
+        Arguments:
+            new_num_position_embeddings (`int`):
+                The number of new position embedding matrix. If position embeddings are learned, increasing the size
+                will add newly initialized vectors at the end, whereas reducing the size will remove vectors from the
+                end. If position embeddings are not learned (*e.g.* sinusoidal position embeddings), increasing the
+                size will add correct vectors at the end following the position encoding algorithm, whereas reducing
+                the size will remove vectors from the end.
+        """
+        num_position_embeds_diff = (
+            new_num_position_embeddings - self.config.max_position_embeddings
+        )
+
+        # no resizing needs to be done if the length stays the same
+        if num_position_embeds_diff == 0:
+            return
+
+        self.config.max_position_embeddings = new_num_position_embeddings
+
+        old_position_embeddings_weight = (
+            self.embeddings.position_embeddings.weight.clone()
+        )
+
+        self.embeddings.position_embeddings = nn.Embedding(
+            self.config.max_position_embeddings, self.config.dim
+        )
+
+        if self.config.sinusoidal_pos_embds:
+            create_sinusoidal_embeddings(
+                n_pos=self.config.max_position_embeddings,
+                dim=self.config.dim,
+                out=self.position_embeddings.weight,
+            )
+        else:
+            with torch.no_grad():
+                if num_position_embeds_diff > 0:
+                    self.embeddings.position_embeddings.weight[
+                        :-num_position_embeds_diff
+                    ] = nn.Parameter(old_position_embeddings_weight)
+                else:
+                    self.embeddings.position_embeddings.weight = nn.Parameter(
+                        old_position_embeddings_weight[:num_position_embeds_diff]
+                    )
+        # move position_embeddings to correct device
+        self.embeddings.position_embeddings.to(self.device)
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, new_embeddings: nn.Embedding):
+        self.embeddings.word_embeddings = new_embeddings
+
+    def _prune_heads(self, heads_to_prune: Dict[int, List[List[int]]]):
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+        class PreTrainedModel
+        """
+        for layer, heads in heads_to_prune.items():
+            self.transformer.layer[layer].attention.prune_heads(heads)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[BaseModelOutput, Tuple[torch.Tensor, ...]]:
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time"
+            )
+        elif input_ids is not None:
+            self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
+            input_shape = input_ids.size()
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+        # Prepare head mask if needed
+        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
+
+        embeddings = self.embeddings(input_ids, inputs_embeds)  # (bs, seq_length, dim)
+
+        if self._use_flash_attention_2:
+            attention_mask = (
+                attention_mask
+                if (attention_mask is not None and 0 in attention_mask)
+                else None
+            )
+        else:
+            if attention_mask is None:
+                attention_mask = torch.ones(
+                    input_shape, device=device
+                )  # (bs, seq_length)
+
+        return self.transformer(
+            x=embeddings,
+            attn_mask=attention_mask,
+            head_mask=head_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+
+class DistilBertForSequenceClassification(S3DownloaderMixin, DistilBertPreTrainedModel):
+    def __init__(self, config: DistilBertConfig, **kwargs):
+        super().__init__(config, **kwargs)
+        self.num_labels = config.num_labels
+        self.config = config
+
+        self.distilbert = DistilBertModel(config)
+        self.pre_classifier = nn.Linear(config.dim, config.dim)
+        self.classifier = nn.Linear(config.dim, config.num_labels)
+        self.dropout = nn.Dropout(config.seq_classif_dropout)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_position_embeddings(self) -> nn.Embedding:
+        """
+        Returns the position embeddings
+        """
+        return self.distilbert.get_position_embeddings()
+
+    def resize_position_embeddings(self, new_num_position_embeddings: int):
+        """
+        Resizes position embeddings of the model if `new_num_position_embeddings != config.max_position_embeddings`.
+
+        Arguments:
+            new_num_position_embeddings (`int`):
+                The number of new position embedding matrix. If position embeddings are learned, increasing the size
+                will add newly initialized vectors at the end, whereas reducing the size will remove vectors from the
+                end. If position embeddings are not learned (*e.g.* sinusoidal position embeddings), increasing the
+                size will add correct vectors at the end following the position encoding algorithm, whereas reducing
+                the size will remove vectors from the end.
+        """
+        self.distilbert.resize_position_embeddings(new_num_position_embeddings)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[SequenceClassifierOutput, Tuple[torch.Tensor, ...]]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        distilbert_output = self.distilbert(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_state = distilbert_output[0]  # (bs, seq_len, dim)
+        pooled_output = hidden_state[:, 0]  # (bs, dim)
+        pooled_output = self.pre_classifier(pooled_output)  # (bs, dim)
+        pooled_output = nn.ReLU()(pooled_output)  # (bs, dim)
+        pooled_output = self.dropout(pooled_output)  # (bs, dim)
+        logits = self.classifier(pooled_output)  # (bs, num_labels)
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (
+                    labels.dtype == torch.long or labels.dtype == torch.int
+                ):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+
+        if not return_dict:
+            output = (logits,) + distilbert_output[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=distilbert_output.hidden_states,
+            attentions=distilbert_output.attentions,
+        )

surya/ocr_error/schema.py

@@ -0,0 +1,8 @@
+from typing import List
+
+from pydantic import BaseModel
+
+
+class OCRErrorDetectionResult(BaseModel):
+    texts: List[str]
+    labels: List[str]

surya/ocr_error/tokenizer.py

@@ -0,0 +1,525 @@
+import collections
+import os
+import unicodedata
+from typing import List, Optional, Tuple
+
+from transformers.tokenization_utils import (
+    PreTrainedTokenizer,
+    _is_control,
+    _is_punctuation,
+    _is_whitespace,
+)
+
+from surya.common.s3 import S3DownloaderMixin
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
+
+
+# Copied from transformers.models.bert.tokenization_bert.load_vocab
+def load_vocab(vocab_file):
+    """Loads a vocabulary file into a dictionary."""
+    vocab = collections.OrderedDict()
+    with open(vocab_file, "r", encoding="utf-8") as reader:
+        tokens = reader.readlines()
+    for index, token in enumerate(tokens):
+        token = token.rstrip("\n")
+        vocab[token] = index
+    return vocab
+
+
+# Copied from transformers.models.bert.tokenization_bert.whitespace_tokenize
+def whitespace_tokenize(text):
+    """Runs basic whitespace cleaning and splitting on a piece of text."""
+    text = text.strip()
+    if not text:
+        return []
+    tokens = text.split()
+    return tokens
+
+
+class DistilBertTokenizer(S3DownloaderMixin, PreTrainedTokenizer):
+    r"""
+    Construct a DistilBERT tokenizer. Based on WordPiece.
+
+    This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
+    this superclass for more information regarding those methods.
+
+    Args:
+        vocab_file (`str`):
+            File containing the vocabulary.
+        do_lower_case (`bool`, *optional*, defaults to `True`):
+            Whether or not to lowercase the input when tokenizing.
+        do_basic_tokenize (`bool`, *optional*, defaults to `True`):
+            Whether or not to do basic tokenization before WordPiece.
+        never_split (`Iterable`, *optional*):
+            Collection of tokens which will never be split during tokenization. Only has an effect when
+            `do_basic_tokenize=True`
+        unk_token (`str`, *optional*, defaults to `"[UNK]"`):
+            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+            token instead.
+        sep_token (`str`, *optional*, defaults to `"[SEP]"`):
+            The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
+            sequence classification or for a text and a question for question answering. It is also used as the last
+            token of a sequence built with special tokens.
+        pad_token (`str`, *optional*, defaults to `"[PAD]"`):
+            The token used for padding, for example when batching sequences of different lengths.
+        cls_token (`str`, *optional*, defaults to `"[CLS]"`):
+            The classifier token which is used when doing sequence classification (classification of the whole sequence
+            instead of per-token classification). It is the first token of the sequence when built with special tokens.
+        mask_token (`str`, *optional*, defaults to `"[MASK]"`):
+            The token used for masking values. This is the token used when training this model with masked language
+            modeling. This is the token which the model will try to predict.
+        tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
+            Whether or not to tokenize Chinese characters.
+
+            This should likely be deactivated for Japanese (see this
+            [issue](https://github.com/huggingface/transformers/issues/328)).
+        strip_accents (`bool`, *optional*):
+            Whether or not to strip all accents. If this option is not specified, then it will be determined by the
+            value for `lowercase` (as in the original BERT).
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file,
+        do_lower_case=True,
+        do_basic_tokenize=True,
+        never_split=None,
+        unk_token="[UNK]",
+        sep_token="[SEP]",
+        pad_token="[PAD]",
+        cls_token="[CLS]",
+        mask_token="[MASK]",
+        tokenize_chinese_chars=True,
+        strip_accents=None,
+        **kwargs,
+    ):
+        if not os.path.isfile(vocab_file):
+            raise ValueError(
+                f"Can't find a vocabulary file at path '{vocab_file}'. To load the vocabulary from a Google pretrained"
+                " model use `tokenizer = DistilBertTokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`"
+            )
+        self.vocab = load_vocab(vocab_file)
+        self.ids_to_tokens = collections.OrderedDict(
+            [(ids, tok) for tok, ids in self.vocab.items()]
+        )
+        self.do_basic_tokenize = do_basic_tokenize
+        if do_basic_tokenize:
+            self.basic_tokenizer = BasicTokenizer(
+                do_lower_case=do_lower_case,
+                never_split=never_split,
+                tokenize_chinese_chars=tokenize_chinese_chars,
+                strip_accents=strip_accents,
+            )
+        self.wordpiece_tokenizer = WordpieceTokenizer(
+            vocab=self.vocab, unk_token=str(unk_token)
+        )
+
+        super().__init__(
+            do_lower_case=do_lower_case,
+            do_basic_tokenize=do_basic_tokenize,
+            never_split=never_split,
+            unk_token=unk_token,
+            sep_token=sep_token,
+            pad_token=pad_token,
+            cls_token=cls_token,
+            mask_token=mask_token,
+            tokenize_chinese_chars=tokenize_chinese_chars,
+            strip_accents=strip_accents,
+            **kwargs,
+        )
+
+    @property
+    # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.do_lower_case
+    def do_lower_case(self):
+        return self.basic_tokenizer.do_lower_case
+
+    @property
+    # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.vocab_size
+    def vocab_size(self):
+        return len(self.vocab)
+
+    # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.get_vocab
+    def get_vocab(self):
+        return dict(self.vocab, **self.added_tokens_encoder)
+
+    # Copied from transformers.models.bert.tokenization_bert.BertTokenizer._tokenize
+    def _tokenize(self, text, split_special_tokens=False):
+        split_tokens = []
+        if self.do_basic_tokenize:
+            for token in self.basic_tokenizer.tokenize(
+                text,
+                never_split=self.all_special_tokens
+                if not split_special_tokens
+                else None,
+            ):
+                # If the token is part of the never_split set
+                if token in self.basic_tokenizer.never_split:
+                    split_tokens.append(token)
+                else:
+                    split_tokens += self.wordpiece_tokenizer.tokenize(token)
+        else:
+            split_tokens = self.wordpiece_tokenizer.tokenize(text)
+        return split_tokens
+
+    # Copied from transformers.models.bert.tokenization_bert.BertTokenizer._convert_token_to_id
+    def _convert_token_to_id(self, token):
+        """Converts a token (str) in an id using the vocab."""
+        return self.vocab.get(token, self.vocab.get(self.unk_token))
+
+    # Copied from transformers.models.bert.tokenization_bert.BertTokenizer._convert_id_to_token
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        return self.ids_to_tokens.get(index, self.unk_token)
+
+    # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.convert_tokens_to_string
+    def convert_tokens_to_string(self, tokens):
+        """Converts a sequence of tokens (string) in a single string."""
+        out_string = " ".join(tokens).replace(" ##", "").strip()
+        return out_string
+
+    # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.build_inputs_with_special_tokens
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
+        adding special tokens. A BERT sequence has the following format:
+
+        - single sequence: `[CLS] X [SEP]`
+        - pair of sequences: `[CLS] A [SEP] B [SEP]`
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs to which the special tokens will be added.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
+        """
+        if token_ids_1 is None:
+            return [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
+        cls = [self.cls_token_id]
+        sep = [self.sep_token_id]
+        return cls + token_ids_0 + sep + token_ids_1 + sep
+
+    # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.get_special_tokens_mask
+    def get_special_tokens_mask(
+        self,
+        token_ids_0: List[int],
+        token_ids_1: Optional[List[int]] = None,
+        already_has_special_tokens: bool = False,
+    ) -> List[int]:
+        """
+        Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
+        special tokens using the tokenizer `prepare_for_model` method.
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+            already_has_special_tokens (`bool`, *optional*, defaults to `False`):
+                Whether or not the token list is already formatted with special tokens for the model.
+
+        Returns:
+            `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
+        """
+
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(
+                token_ids_0=token_ids_0,
+                token_ids_1=token_ids_1,
+                already_has_special_tokens=True,
+            )
+
+        if token_ids_1 is not None:
+            return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
+        return [1] + ([0] * len(token_ids_0)) + [1]
+
+    # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.create_token_type_ids_from_sequences
+    def create_token_type_ids_from_sequences(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Create a mask from the two sequences passed to be used in a sequence-pair classification task. A BERT sequence
+        pair mask has the following format:
+
+        ```
+        0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
+        | first sequence    | second sequence |
+        ```
+
+        If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
+        """
+        sep = [self.sep_token_id]
+        cls = [self.cls_token_id]
+        if token_ids_1 is None:
+            return len(cls + token_ids_0 + sep) * [0]
+        return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
+
+    # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.save_vocabulary
+    def save_vocabulary(
+        self, save_directory: str, filename_prefix: Optional[str] = None
+    ) -> Tuple[str]:
+        index = 0
+        if os.path.isdir(save_directory):
+            vocab_file = os.path.join(
+                save_directory,
+                (filename_prefix + "-" if filename_prefix else "")
+                + VOCAB_FILES_NAMES["vocab_file"],
+            )
+        else:
+            vocab_file = (
+                filename_prefix + "-" if filename_prefix else ""
+            ) + save_directory
+        with open(vocab_file, "w", encoding="utf-8") as writer:
+            for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]):
+                if index != token_index:
+                    # logger.warning(
+                    #     f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
+                    #     " Please check that the vocabulary is not corrupted!"
+                    # )
+                    index = token_index
+                writer.write(token + "\n")
+                index += 1
+        return (vocab_file,)
+
+
+# Copied from transformers.models.bert.tokenization_bert.BasicTokenizer
+class BasicTokenizer(object):
+    """
+    Constructs a BasicTokenizer that will run basic tokenization (punctuation splitting, lower casing, etc.).
+
+    Args:
+        do_lower_case (`bool`, *optional*, defaults to `True`):
+            Whether or not to lowercase the input when tokenizing.
+        never_split (`Iterable`, *optional*):
+            Collection of tokens which will never be split during tokenization. Only has an effect when
+            `do_basic_tokenize=True`
+        tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
+            Whether or not to tokenize Chinese characters.
+
+            This should likely be deactivated for Japanese (see this
+            [issue](https://github.com/huggingface/transformers/issues/328)).
+        strip_accents (`bool`, *optional*):
+            Whether or not to strip all accents. If this option is not specified, then it will be determined by the
+            value for `lowercase` (as in the original BERT).
+        do_split_on_punc (`bool`, *optional*, defaults to `True`):
+            In some instances we want to skip the basic punctuation splitting so that later tokenization can capture
+            the full context of the words, such as contractions.
+    """
+
+    def __init__(
+        self,
+        do_lower_case=True,
+        never_split=None,
+        tokenize_chinese_chars=True,
+        strip_accents=None,
+        do_split_on_punc=True,
+    ):
+        if never_split is None:
+            never_split = []
+        self.do_lower_case = do_lower_case
+        self.never_split = set(never_split)
+        self.tokenize_chinese_chars = tokenize_chinese_chars
+        self.strip_accents = strip_accents
+        self.do_split_on_punc = do_split_on_punc
+
+    def tokenize(self, text, never_split=None):
+        """
+        Basic Tokenization of a piece of text. For sub-word tokenization, see WordPieceTokenizer.
+
+        Args:
+            never_split (`List[str]`, *optional*)
+                Kept for backward compatibility purposes. Now implemented directly at the base class level (see
+                [`PreTrainedTokenizer.tokenize`]) List of token not to split.
+        """
+        # union() returns a new set by concatenating the two sets.
+        never_split = (
+            self.never_split.union(set(never_split))
+            if never_split
+            else self.never_split
+        )
+        text = self._clean_text(text)
+
+        # This was added on November 1st, 2018 for the multilingual and Chinese
+        # models. This is also applied to the English models now, but it doesn't
+        # matter since the English models were not trained on any Chinese data
+        # and generally don't have any Chinese data in them (there are Chinese
+        # characters in the vocabulary because Wikipedia does have some Chinese
+        # words in the English Wikipedia.).
+        if self.tokenize_chinese_chars:
+            text = self._tokenize_chinese_chars(text)
+        # prevents treating the same character with different unicode codepoints as different characters
+        unicode_normalized_text = unicodedata.normalize("NFC", text)
+        orig_tokens = whitespace_tokenize(unicode_normalized_text)
+        split_tokens = []
+        for token in orig_tokens:
+            if token not in never_split:
+                if self.do_lower_case:
+                    token = token.lower()
+                    if self.strip_accents is not False:
+                        token = self._run_strip_accents(token)
+                elif self.strip_accents:
+                    token = self._run_strip_accents(token)
+            split_tokens.extend(self._run_split_on_punc(token, never_split))
+
+        output_tokens = whitespace_tokenize(" ".join(split_tokens))
+        return output_tokens
+
+    def _run_strip_accents(self, text):
+        """Strips accents from a piece of text."""
+        text = unicodedata.normalize("NFD", text)
+        output = []
+        for char in text:
+            cat = unicodedata.category(char)
+            if cat == "Mn":
+                continue
+            output.append(char)
+        return "".join(output)
+
+    def _run_split_on_punc(self, text, never_split=None):
+        """Splits punctuation on a piece of text."""
+        if not self.do_split_on_punc or (
+            never_split is not None and text in never_split
+        ):
+            return [text]
+        chars = list(text)
+        i = 0
+        start_new_word = True
+        output = []
+        while i < len(chars):
+            char = chars[i]
+            if _is_punctuation(char):
+                output.append([char])
+                start_new_word = True
+            else:
+                if start_new_word:
+                    output.append([])
+                start_new_word = False
+                output[-1].append(char)
+            i += 1
+
+        return ["".join(x) for x in output]
+
+    def _tokenize_chinese_chars(self, text):
+        """Adds whitespace around any CJK character."""
+        output = []
+        for char in text:
+            cp = ord(char)
+            if self._is_chinese_char(cp):
+                output.append(" ")
+                output.append(char)
+                output.append(" ")
+            else:
+                output.append(char)
+        return "".join(output)
+
+    def _is_chinese_char(self, cp):
+        """Checks whether CP is the codepoint of a CJK character."""
+        # This defines a "chinese character" as anything in the CJK Unicode block:
+        #   https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
+        #
+        # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
+        # despite its name. The modern Korean Hangul alphabet is a different block,
+        # as is Japanese Hiragana and Katakana. Those alphabets are used to write
+        # space-separated words, so they are not treated specially and handled
+        # like the all of the other languages.
+        if (
+            (cp >= 0x4E00 and cp <= 0x9FFF)
+            or (cp >= 0x3400 and cp <= 0x4DBF)  #
+            or (cp >= 0x20000 and cp <= 0x2A6DF)  #
+            or (cp >= 0x2A700 and cp <= 0x2B73F)  #
+            or (cp >= 0x2B740 and cp <= 0x2B81F)  #
+            or (cp >= 0x2B820 and cp <= 0x2CEAF)  #
+            or (cp >= 0xF900 and cp <= 0xFAFF)
+            or (cp >= 0x2F800 and cp <= 0x2FA1F)  #
+        ):  #
+            return True
+
+        return False
+
+    def _clean_text(self, text):
+        """Performs invalid character removal and whitespace cleanup on text."""
+        output = []
+        for char in text:
+            cp = ord(char)
+            if cp == 0 or cp == 0xFFFD or _is_control(char):
+                continue
+            if _is_whitespace(char):
+                output.append(" ")
+            else:
+                output.append(char)
+        return "".join(output)
+
+
+# Copied from transformers.models.bert.tokenization_bert.WordpieceTokenizer
+class WordpieceTokenizer(object):
+    """Runs WordPiece tokenization."""
+
+    def __init__(self, vocab, unk_token, max_input_chars_per_word=100):
+        self.vocab = vocab
+        self.unk_token = unk_token
+        self.max_input_chars_per_word = max_input_chars_per_word
+
+    def tokenize(self, text):
+        """
+        Tokenizes a piece of text into its word pieces. This uses a greedy longest-match-first algorithm to perform
+        tokenization using the given vocabulary.
+
+        For example, `input = "unaffable"` wil return as output `["un", "##aff", "##able"]`.
+
+        Args:
+            text: A single token or whitespace separated tokens. This should have
+                already been passed through *BasicTokenizer*.
+
+        Returns:
+            A list of wordpiece tokens.
+        """
+
+        output_tokens = []
+        for token in whitespace_tokenize(text):
+            chars = list(token)
+            if len(chars) > self.max_input_chars_per_word:
+                output_tokens.append(self.unk_token)
+                continue
+
+            is_bad = False
+            start = 0
+            sub_tokens = []
+            while start < len(chars):
+                end = len(chars)
+                cur_substr = None
+                while start < end:
+                    substr = "".join(chars[start:end])
+                    if start > 0:
+                        substr = "##" + substr
+                    if substr in self.vocab:
+                        cur_substr = substr
+                        break
+                    end -= 1
+                if cur_substr is None:
+                    is_bad = True
+                    break
+                sub_tokens.append(cur_substr)
+                start = end
+
+            if is_bad:
+                output_tokens.append(self.unk_token)
+            else:
+                output_tokens.extend(sub_tokens)
+        return output_tokens

surya/recognition/__init__.py

@@ -0,0 +1,404 @@
+"""RecognitionPredictor: per-block OCR via BLOCK_PROMPT.
+
+Given page images and corresponding LayoutResult (or any list of LayoutBox),
+crops each block, runs BLOCK_PROMPT, returns PageOCRResult per page.
+"""
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+from PIL import Image
+
+from surya.common.blank import is_blank_region
+from surya.inference import SuryaInferenceManager, get_default_manager
+from surya.inference.parsers import clean_block_html, parse_full_page_html
+from surya.inference.prompts import (
+    PROMPT_TYPE_BLOCK,
+    PROMPT_TYPE_HIGH_ACCURACY_BBOX,
+    SKIP_OCR_LABELS,
+)
+from surya.inference.schema import BatchInputItem
+from surya.inference.util import image_token_budget
+from surya.layout.label import LAYOUT_PRED_RELABEL, TEXT_LABELS
+from surya.layout.schema import LayoutResult
+from surya.logging import get_logger
+from surya.recognition.schema import (
+    BlockOCRResult,
+    PageOCRResult,
+)
+from surya.settings import settings
+from surya.timing import timing_span
+
+logger = get_logger()
+
+
+# Surya's canonical labels we shouldn't OCR (mirrors model-emitted SKIP_OCR_LABELS
+# after canonicalization).
+SKIP_CANON_LABELS = {LAYOUT_PRED_RELABEL.get(lbl, lbl) for lbl in SKIP_OCR_LABELS}
+
+
+def _crop_block(image: Image.Image, polygon, pad: int = 4) -> Image.Image:
+    xs = [p[0] for p in polygon]
+    ys = [p[1] for p in polygon]
+    x0 = max(0, int(min(xs)) - pad)
+    y0 = max(0, int(min(ys)) - pad)
+    x1 = min(image.size[0], int(max(xs)) + pad)
+    y1 = min(image.size[1], int(max(ys)) + pad)
+    if x1 <= x0 or y1 <= y0:
+        return image.crop((0, 0, 1, 1))
+    return image.crop((x0, y0, x1, y1))
+
+
+def _drop_blank_text_blocks(
+    image: Image.Image,
+    blocks: List[BlockOCRResult],
+) -> List[BlockOCRResult]:
+    """Drop text-labeled blocks whose source page region is essentially blank.
+
+    Full-page OCR can emit text divs for regions that are visually empty
+    (margins, gutter space) — the model hallucinates a paragraph where there
+    is none. We crop the region, count near-white pixels, and drop the block
+    when the fraction exceeds ``blank_pixel_fraction``. Only text-like labels
+    (see ``TEXT_LABELS``) are eligible: tables, forms, equations, and visual
+    blocks may legitimately contain large whitespace and are left untouched.
+    """
+    kept: List[BlockOCRResult] = []
+    dropped = 0
+    for blk in blocks:
+        if blk.label not in TEXT_LABELS or blk.skipped or blk.error:
+            kept.append(blk)
+            continue
+        crop = _crop_block(image, blk.polygon)
+        if not is_blank_region(crop):
+            kept.append(blk)
+            continue
+        dropped += 1
+    if dropped:
+        logger.info(f"dropped {dropped} blank text block(s) from full-page OCR")
+    return kept
+
+
+def _detect_repeat_loop(
+    text: str,
+    base_max_repeats: int = 4,
+    window_size: int = 500,
+    scaling_factor: float = 3.0,
+) -> bool:
+    """True iff the tail of ``text`` ends in a repeating sequence.
+
+    Ported from chandra's detect_repeat_token. For each candidate length
+    1..window_size/2, takes that many trailing chars and counts consecutive
+    identical preceding blocks. Shorter loops need many repeats to count;
+    longer ones only need a few. Catches the typical decoder failure mode
+    where a page output gets stuck emitting the same div / phrase until it
+    hits max_tokens.
+    """
+    if not text:
+        return False
+    for seq_len in range(1, window_size // 2 + 1):
+        candidate = text[-seq_len:]
+        max_repeats = int(base_max_repeats * (1 + scaling_factor / seq_len))
+        repeats = 0
+        pos = len(text) - seq_len
+        while pos >= 0 and text[pos : pos + seq_len] == candidate:
+            repeats += 1
+            pos -= seq_len
+        if repeats > max_repeats:
+            return True
+    return False
+
+
+class RecognitionPredictor:
+    """Per-block OCR. Construct with a SuryaInferenceManager (or rely on default)."""
+
+    def __init__(self, manager: Optional[SuryaInferenceManager] = None):
+        self.manager = manager
+        self._disable_tqdm = settings.DISABLE_TQDM
+
+    @property
+    def disable_tqdm(self) -> bool:
+        return self._disable_tqdm
+
+    @disable_tqdm.setter
+    def disable_tqdm(self, value: bool) -> None:
+        self._disable_tqdm = bool(value)
+
+    def to(self, *args, **kwargs):
+        return
+
+    def __call__(
+        self,
+        images: List[Image.Image],
+        layout_results: Optional[List[LayoutResult]] = None,
+        *,
+        full_page: Optional[bool] = None,
+    ) -> List[PageOCRResult]:
+        """Run OCR on each page.
+
+        Mode resolution:
+          - ``full_page=None`` (default): block mode if ``layout_results`` is
+            given, else full-page mode. This is the most-do-what-I-mean form.
+          - ``full_page=True``: full-page OCR (single HIGH_ACCURACY_BBOX_PROMPT
+            request per page). ``layout_results`` is ignored — a warning is
+            logged if it was supplied.
+          - ``full_page=False``: block mode (per-layout-block OCR request).
+            ``layout_results`` is required.
+
+        Full-page is the more accurate path; block mode is for callers that
+        specifically need per-block crops (e.g. for downstream merging with
+        text-line detection).
+        """
+        if not images:
+            return []
+        if full_page is None:
+            full_page = layout_results is None
+        if full_page:
+            if layout_results is not None:
+                logger.info(
+                    "RecognitionPredictor called with full_page=True and "
+                    "layout_results; layout will be used as fallback if the "
+                    "full-page output devolves into a repetition loop."
+                )
+            return self._full_page_ocr(images, fallback_layout=layout_results)
+        if layout_results is None:
+            raise ValueError("layout_results required when full_page=False")
+        if len(images) != len(layout_results):
+            raise ValueError(
+                f"images and layout_results must be same length "
+                f"({len(images)} vs {len(layout_results)})"
+            )
+        manager = self.manager or get_default_manager()
+
+        # Build a flat batch across all pages for max concurrency
+        batch: List[BatchInputItem] = []
+        block_index_map: List[tuple[int, int]] = []  # (page_idx, block_idx)
+        skipped_flags: List[bool] = []
+
+        with timing_span("recognition_build_block_batch", page_count=len(images)):
+            for page_idx, (img, layout) in enumerate(zip(images, layout_results)):
+                page_boxes = sorted(layout.bboxes, key=lambda b: (b.position, b.bbox[1], b.bbox[0]))
+                page_boxes = page_boxes[: settings.SURYA_MAX_BLOCKS_PER_PAGE]
+                if len(layout.bboxes) > len(page_boxes):
+                    logger.info(
+                        f"capped OCR blocks for page {page_idx}: "
+                        f"{len(layout.bboxes)} -> {len(page_boxes)}"
+                    )
+                for block_idx, box in enumerate(page_boxes):
+                    skip = box.label in SKIP_CANON_LABELS
+                    skipped_flags.append(skip)
+                    if skip:
+                        continue
+                    crop = _crop_block(img, box.polygon)
+                    max_tokens = image_token_budget(
+                        box.count, ceiling=settings.SURYA_MAX_TOKENS_BLOCK_CEILING
+                    )
+                    batch.append(
+                        BatchInputItem(
+                            image=crop,
+                            prompt_type=PROMPT_TYPE_BLOCK,
+                            max_tokens=max_tokens,
+                            metadata={"page_idx": page_idx, "block_idx": block_idx},
+                        )
+                    )
+                    block_index_map.append((page_idx, block_idx))
+
+        with timing_span(
+            "recognition_manager_generate",
+            item_count=len(batch),
+            skipped_count=sum(1 for flag in skipped_flags if flag),
+            max_tokens_sum=sum(item.max_tokens or 0 for item in batch),
+        ):
+            outputs = manager.generate(batch) if batch else []
+
+        # Index outputs by (page_idx, block_idx)
+        with timing_span(
+            "recognition_assemble_pages",
+            output_count=len(outputs),
+            token_count=sum(out.token_count or 0 for out in outputs),
+        ):
+            out_by_key = {}
+            for out in outputs:
+                key = (out.metadata["page_idx"], out.metadata["block_idx"])
+                out_by_key[key] = out
+
+            # Assemble PageOCRResult per page
+            results: List[PageOCRResult] = []
+            for page_idx, (img, layout) in enumerate(zip(images, layout_results)):
+                w, h = img.size
+                blocks: List[BlockOCRResult] = []
+                page_boxes = sorted(layout.bboxes, key=lambda b: (b.position, b.bbox[1], b.bbox[0]))
+                page_boxes = page_boxes[: settings.SURYA_MAX_BLOCKS_PER_PAGE]
+                for block_idx, box in enumerate(page_boxes):
+                    skip = box.label in SKIP_CANON_LABELS
+                    if skip:
+                        blocks.append(
+                            BlockOCRResult(
+                                polygon=box.polygon,
+                                label=box.label,
+                                raw_label=box.raw_label,
+                                reading_order=box.position,
+                                html="",
+                                skipped=True,
+                                confidence=1.0,
+                            )
+                        )
+                        continue
+                    out = out_by_key.get((page_idx, block_idx))
+                    if out is None or out.error:
+                        blocks.append(
+                            BlockOCRResult(
+                                polygon=box.polygon,
+                                label=box.label,
+                                raw_label=box.raw_label,
+                                reading_order=box.position,
+                                html="",
+                                skipped=False,
+                                error=True,
+                                confidence=0.0,
+                            )
+                        )
+                        continue
+                    html = clean_block_html(out.raw)
+                    conf = out.mean_token_prob if out.mean_token_prob is not None else 1.0
+                    blocks.append(
+                        BlockOCRResult(
+                            polygon=box.polygon,
+                            label=box.label,
+                            raw_label=box.raw_label,
+                            reading_order=box.position,
+                            html=html,
+                            skipped=False,
+                            error=False,
+                            confidence=conf,
+                            raw_logprobs=out.logprobs,
+                        )
+                    )
+                results.append(
+                    PageOCRResult(blocks=blocks, image_bbox=[0, 0, float(w), float(h)])
+                )
+        return results
+
+    def _full_page_ocr(
+        self,
+        images: List[Image.Image],
+        fallback_layout: Optional[List[LayoutResult]] = None,
+    ) -> List[PageOCRResult]:
+        """One HIGH_ACCURACY_BBOX_PROMPT request per page; parses divs into blocks.
+
+        On per-page failure (parse error, empty output, or a detected
+        repetition loop in the decoder output), falls back to layout +
+        block-mode OCR for that page only. ``fallback_layout``, if given,
+        provides per-page LayoutResults to use on fallback; otherwise the
+        LayoutPredictor is invoked lazily for just the affected pages.
+        """
+        manager = self.manager or get_default_manager()
+        with timing_span("recognition_build_full_page_batch", page_count=len(images)):
+            batch = [
+                BatchInputItem(
+                    image=img,
+                    prompt_type=PROMPT_TYPE_HIGH_ACCURACY_BBOX,
+                    max_tokens=settings.SURYA_MAX_TOKENS_FULL_PAGE,
+                    metadata={"page_idx": i},
+                )
+                for i, img in enumerate(images)
+            ]
+        with timing_span(
+            "recognition_full_page_generate",
+            item_count=len(batch),
+            max_tokens=settings.SURYA_MAX_TOKENS_FULL_PAGE,
+        ):
+            outputs = manager.generate(batch)
+        out_by_page = {o.metadata["page_idx"]: o for o in outputs}
+
+        results: List[Optional[PageOCRResult]] = [None] * len(images)
+        needs_fallback: List[int] = []
+        for page_idx, img in enumerate(images):
+            w, h = img.size
+            page_bbox = [0, 0, float(w), float(h)]
+            out = out_by_page.get(page_idx)
+            if out is None or out.error:
+                # Hard failure (request lost / server error). Always fallback.
+                needs_fallback.append(page_idx)
+                continue
+            if not out.raw:
+                # Empty model output. If the page is genuinely blank, the
+                # model is correct — return an empty result. Only fall back
+                # when the page has content the model failed to emit.
+                if is_blank_region(img):
+                    results[page_idx] = PageOCRResult(blocks=[], image_bbox=page_bbox)
+                else:
+                    logger.info(
+                        f"empty full-page output for non-blank page {page_idx}; "
+                        f"falling back to layout + block OCR"
+                    )
+                    needs_fallback.append(page_idx)
+                continue
+            if _detect_repeat_loop(out.raw):
+                logger.info(
+                    f"full-page output for page {page_idx} appears to loop; "
+                    f"falling back to layout + block OCR"
+                )
+                needs_fallback.append(page_idx)
+                continue
+            try:
+                parsed = parse_full_page_html(out.raw)
+            except Exception as e:
+                logger.warning(
+                    f"Full-page parse failed for page {page_idx}: {e}; "
+                    f"falling back to layout + block OCR"
+                )
+                needs_fallback.append(page_idx)
+                continue
+            confidence = out.mean_token_prob if out.mean_token_prob is not None else 1.0
+            blocks: List[BlockOCRResult] = []
+            for idx, item in enumerate(parsed):
+                x0 = item.bbox[0] / settings.BBOX_SCALE * w
+                y0 = item.bbox[1] / settings.BBOX_SCALE * h
+                x1 = item.bbox[2] / settings.BBOX_SCALE * w
+                y1 = item.bbox[3] / settings.BBOX_SCALE * h
+                polygon = [[x0, y0], [x1, y0], [x1, y1], [x0, y1]]
+                canon = LAYOUT_PRED_RELABEL.get(item.label, item.label)
+                skipped = canon in SKIP_CANON_LABELS
+                blocks.append(
+                    BlockOCRResult(
+                        polygon=polygon,
+                        label=canon,
+                        raw_label=item.label,
+                        reading_order=idx,
+                        html="" if skipped else item.html,
+                        skipped=skipped,
+                        error=False,
+                        confidence=confidence,
+                    )
+                )
+            blocks = _drop_blank_text_blocks(img, blocks)
+            results[page_idx] = PageOCRResult(blocks=blocks, image_bbox=page_bbox)
+
+        # Block-mode fallback for any pages whose full-page output failed or looped.
+        if needs_fallback:
+            fb_images = [images[i] for i in needs_fallback]
+            if fallback_layout is not None:
+                fb_layouts = [fallback_layout[i] for i in needs_fallback]
+            else:
+                # Lazy import to avoid the surya.layout ↔ surya.recognition cycle.
+                from surya.layout import LayoutPredictor
+
+                logger.info(
+                    f"running layout for {len(fb_images)} page(s) requiring "
+                    f"block-mode fallback"
+                )
+                fb_layouts = LayoutPredictor(self.manager)(fb_images)
+            fb_results = self.__call__(fb_images, fb_layouts, full_page=False)
+            for fb_idx, page_idx in enumerate(needs_fallback):
+                results[page_idx] = fb_results[fb_idx]
+
+        # Backfill any still-None pages with empty results (defensive — shouldn't happen).
+        out_results: List[PageOCRResult] = []
+        for page_idx, img in enumerate(images):
+            r = results[page_idx]
+            if r is None:
+                w, h = img.size
+                r = PageOCRResult(blocks=[], image_bbox=[0, 0, float(w), float(h)])
+            out_results.append(r)
+        return out_results

surya/recognition/schema.py

@@ -0,0 +1,19 @@
+from typing import List
+
+from pydantic import BaseModel
+
+from surya.common.polygon import PolygonBox
+
+
+class BlockOCRResult(PolygonBox):
+    label: str  # canonicalized layout label (Picture, Text, ...)
+    raw_label: str = ""  # original model label
+    reading_order: int  # 0-indexed position in layout output
+    html: str = ""  # block HTML (BLOCK_PROMPT output, "" if skipped)
+    skipped: bool = False  # True if label was in SKIP_OCR_LABELS
+    error: bool = False
+
+
+class PageOCRResult(BaseModel):
+    blocks: List[BlockOCRResult]
+    image_bbox: List[float]

surya/scripts/build_gguf.py

@@ -0,0 +1,311 @@
+"""Build a GGUF artifact from a surya VLM checkpoint, suitable for stock llama.cpp.
+
+Three patches are applied to llama.cpp's convert_hf_to_gguf.py before
+running the HF→GGUF conversion. All fixes are baked into the output
+artifact, so the resulting GGUF runs on unpatched llama.cpp.
+
+Two checkpoint-side patches are also applied to a working copy:
+  - config.json:           architectures → ["Qwen3_5ForConditionalGeneration"]
+  - tokenizer_config.json: tokenizer_class → "PreTrainedTokenizerFast",
+                           strip backend / extra_special_tokens / is_local
+
+The original checkpoint is never modified — patches land in a sibling
+working dir under --out-dir.
+
+Usage:
+    python -m surya.scripts.build_gguf \\
+        --checkpoint datalab-to/surya-2.1.6 \\
+        --out-dir ./gguf-build
+"""
+
+import argparse
+import json
+import os
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+LLAMA_CPP_REPO = "https://github.com/ggerganov/llama.cpp.git"
+# Pinned commit the converter patches were authored against. Bump when
+# the upstream converter drifts; re-validate the anchor strings below.
+LLAMA_CPP_REV = "bbeb89d76c41bc250f16e4a6fefcc9b530d6e3f3"
+
+
+# ---- llama.cpp converter patches -----------------------------------------
+# Each patch is (sentinel, anchor, replacement). Idempotent: if `sentinel`
+# is already in the file, the patch is skipped.
+
+_PATCH_REGISTER_HASH_ANCHOR = """        if chkhsh == "862f827721df956049dff5ca81a57f29e575280bc622e290d3bf4e35eca29015":
+            # ref: https://huggingface.co/codefuse-ai/F2LLM-v2-4B
+            res = "f2llmv2"
+"""
+
+_PATCH_REGISTER_HASH_REPLACEMENT = (
+    _PATCH_REGISTER_HASH_ANCHOR
+    + """        if chkhsh == "11865354be60ff9206694aed04242190f1807029bbd750bfbced09b4d26f1ad2":
+            # surya-2.1.0 char-level WordLevel tokenizer (Split regex=".")
+            res = "default"
+"""
+)
+
+_PATCH_VOCAB_ANCHOR = """    def _set_vocab_gpt2(self) -> None:
+        tokens, toktypes, tokpre = self.get_vocab_base()
+        self.gguf_writer.add_tokenizer_model("gpt2")
+        self.gguf_writer.add_tokenizer_pre(tokpre)
+        self.gguf_writer.add_token_list(tokens)
+        self.gguf_writer.add_token_types(toktypes)
+
+        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)
+        special_vocab.add_to_gguf(self.gguf_writer)
+"""
+
+_PATCH_VOCAB_REPLACEMENT = '''    def _set_vocab_gpt2(self) -> None:
+        tokens, toktypes, tokpre = self.get_vocab_base()
+        # surya: char-level / WordLevel vocabs contain raw bytes (e.g. " ",
+        # "\\n"). llama.cpp's gpt2 vocab decoder applies bytes_to_unicode
+        # when emitting NORMAL tokens, so encode bytes here for round-trip.
+        # Idempotent on already-encoded vocabs.
+        tokens = self._maybe_encode_gpt2_bytes(tokens, toktypes)
+        self.gguf_writer.add_tokenizer_model("gpt2")
+        self.gguf_writer.add_tokenizer_pre(tokpre)
+        self.gguf_writer.add_token_list(tokens)
+        self.gguf_writer.add_token_types(toktypes)
+
+        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)
+        special_vocab.add_to_gguf(self.gguf_writer)
+        # surya: char-level / WordLevel vocabs have no merges, but the gpt2
+        # vocab loader requires the field. Write a single dummy entry.
+        if not special_vocab.merges:
+            self.gguf_writer.add_token_merges(["a a"])
+
+    @staticmethod
+    def _maybe_encode_gpt2_bytes(tokens: list[str], toktypes: list[int]) -> list[str]:
+        """Apply GPT-2 bytes_to_unicode to NORMAL tokens iff any contain raw
+        whitespace/control bytes. Idempotent on already-encoded vocabs."""
+        bs = (list(range(ord("!"), ord("~") + 1))
+              + list(range(ord("¡"), ord("¬") + 1))
+              + list(range(ord("®"), ord("ÿ") + 1)))
+        cs = bs[:]
+        n = 0
+        for b in range(256):
+            if b not in bs:
+                bs.append(b)
+                cs.append(2 ** 8 + n)
+                n += 1
+        byte_to_unicode = {b: chr(c) for b, c in zip(bs, cs)}
+        printable = set(range(ord("!"), ord("~") + 1))
+        needs = False
+        for tok, ttype in zip(tokens, toktypes):
+            if ttype != gguf.TokenType.NORMAL:
+                continue
+            for ch in tok:
+                cb = ord(ch)
+                if cb < 0x80 and cb not in printable:
+                    needs = True
+                    break
+            if needs:
+                break
+        if not needs:
+            return tokens
+        out: list[str] = []
+        for tok, ttype in zip(tokens, toktypes):
+            if ttype == gguf.TokenType.NORMAL:
+                out.append("".join(byte_to_unicode[b] for b in tok.encode("utf-8")))
+            else:
+                out.append(tok)
+        return out
+'''
+
+CONVERTER_PATCHES = [
+    {
+        "name": "register surya-2.1.0 pre-tokenizer hash",
+        "sentinel": '"11865354be60ff9206694aed04242190f1807029bbd750bfbced09b4d26f1ad2"',
+        "anchor": _PATCH_REGISTER_HASH_ANCHOR,
+        "replacement": _PATCH_REGISTER_HASH_REPLACEMENT,
+    },
+    {
+        "name": "byte-encode NORMAL tokens + dummy merges in _set_vocab_gpt2",
+        "sentinel": "_maybe_encode_gpt2_bytes",
+        "anchor": _PATCH_VOCAB_ANCHOR,
+        "replacement": _PATCH_VOCAB_REPLACEMENT,
+    },
+]
+
+
+def patch_converter(convert_py: Path) -> None:
+    text = convert_py.read_text()
+    changed = False
+    for p in CONVERTER_PATCHES:
+        if p["sentinel"] in text:
+            print(f"  [skip]  {p['name']} (already applied)")
+            continue
+        if p["anchor"] not in text:
+            raise RuntimeError(
+                f"converter patch {p['name']!r} could not find its anchor. "
+                f"llama.cpp upstream may have drifted; re-pin LLAMA_CPP_REV "
+                f"and re-validate anchors."
+            )
+        text = text.replace(p["anchor"], p["replacement"], 1)
+        print(f"  [apply] {p['name']}")
+        changed = True
+    if changed:
+        convert_py.write_text(text)
+
+
+# ---- checkpoint-side patches ----------------------------------------------
+
+
+def patch_checkpoint(src: Path, dst: Path) -> None:
+    """Symlink-clone src into dst, with config.json + tokenizer_config.json
+    rewritten for stock transformers/llama.cpp compatibility."""
+    if dst.exists():
+        shutil.rmtree(dst)
+    dst.mkdir(parents=True)
+    overrides = {"config.json", "tokenizer_config.json"}
+    for entry in src.iterdir():
+        if entry.name in overrides:
+            continue
+        os.symlink(entry.resolve(), dst / entry.name)
+
+    cfg = json.loads((src / "config.json").read_text())
+    cfg["architectures"] = ["Qwen3_5ForConditionalGeneration"]
+    (dst / "config.json").write_text(json.dumps(cfg, indent=2))
+
+    tk = json.loads((src / "tokenizer_config.json").read_text())
+    tk["tokenizer_class"] = "PreTrainedTokenizerFast"
+    for k in ("backend", "extra_special_tokens", "is_local"):
+        tk.pop(k, None)
+    (dst / "tokenizer_config.json").write_text(json.dumps(tk, indent=2))
+
+
+# ---- llama.cpp resolution -------------------------------------------------
+
+
+def ensure_llama_cpp(repo_dir: Path, rev: str) -> Path:
+    if not repo_dir.exists():
+        repo_dir.parent.mkdir(parents=True, exist_ok=True)
+        print(f"[clone] {LLAMA_CPP_REPO} → {repo_dir}")
+        subprocess.check_call(["git", "clone", LLAMA_CPP_REPO, str(repo_dir)])
+    head = subprocess.check_output(
+        ["git", "-C", str(repo_dir), "rev-parse", "HEAD"], text=True
+    ).strip()
+    if head != rev:
+        # Discard any prior patches so the checkout is clean.
+        subprocess.check_call(
+            ["git", "-C", str(repo_dir), "reset", "--hard", "--quiet", "HEAD"]
+        )
+        subprocess.check_call(
+            ["git", "-C", str(repo_dir), "fetch", "--quiet", "origin"]
+        )
+        print(f"[checkout] llama.cpp @ {rev}")
+        subprocess.check_call(["git", "-C", str(repo_dir), "checkout", "--quiet", rev])
+    return repo_dir
+
+
+# ---- checkpoint resolution ------------------------------------------------
+
+
+def resolve_checkpoint(checkpoint: str) -> Path:
+    p = Path(checkpoint)
+    if p.exists():
+        return p.resolve()
+    from huggingface_hub import snapshot_download
+
+    print(f"[download] {checkpoint}")
+    return Path(snapshot_download(checkpoint))
+
+
+# ---- main -----------------------------------------------------------------
+
+
+def main() -> int:
+    from surya.settings import settings
+
+    ap = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    ap.add_argument(
+        "--checkpoint",
+        default=settings.SURYA_MODEL_CHECKPOINT,
+        help="HF repo id or local checkpoint dir",
+    )
+    ap.add_argument("--out-dir", type=Path, default=Path("./gguf-build"))
+    ap.add_argument(
+        "--name",
+        default="surya-2",
+        help="Output basename. Produces <name>.gguf and <name>-mmproj.gguf",
+    )
+    ap.add_argument(
+        "--llama-cpp-dir",
+        type=Path,
+        default=Path.home() / ".cache" / "datalab" / "llama.cpp",
+    )
+    ap.add_argument("--llama-cpp-rev", default=LLAMA_CPP_REV)
+    ap.add_argument(
+        "--outtype",
+        default="f16",
+        help="convert_hf_to_gguf --outtype (f16, bf16, q8_0, ...)",
+    )
+    ap.add_argument(
+        "--keep-work",
+        action="store_true",
+        help="Keep the patched-checkpoint working dir on success",
+    )
+    args = ap.parse_args()
+
+    args.out_dir.mkdir(parents=True, exist_ok=True)
+    work = args.out_dir / "_patched_ckpt"
+
+    src = resolve_checkpoint(args.checkpoint)
+    print(f"[checkpoint] {src}")
+
+    print("[patch] checkpoint config + tokenizer_config")
+    patch_checkpoint(src, work)
+
+    repo = ensure_llama_cpp(args.llama_cpp_dir, args.llama_cpp_rev)
+    convert_py = repo / "convert_hf_to_gguf.py"
+    print("[patch] llama.cpp convert_hf_to_gguf.py")
+    patch_converter(convert_py)
+
+    out_llm = (args.out_dir / f"{args.name}.gguf").resolve()
+    out_mmproj = (args.out_dir / f"{args.name}-mmproj.gguf").resolve()
+
+    print(f"[convert] LLM → {out_llm}")
+    subprocess.check_call(
+        [
+            sys.executable,
+            str(convert_py),
+            str(work),
+            "--outfile",
+            str(out_llm),
+            "--outtype",
+            args.outtype,
+        ]
+    )
+    print(f"[convert] mmproj → {out_mmproj}")
+    subprocess.check_call(
+        [
+            sys.executable,
+            str(convert_py),
+            str(work),
+            "--mmproj",
+            "--outfile",
+            str(out_mmproj),
+            "--outtype",
+            args.outtype,
+        ]
+    )
+
+    if not args.keep_work:
+        shutil.rmtree(work)
+
+    print()
+    print(f"  LLM:    {out_llm}")
+    print(f"  mmproj: {out_mmproj}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

surya/scripts/config.py

@@ -0,0 +1,98 @@
+from typing import List
+
+import click
+import os
+from surya.input.load import load_from_folder, load_from_file
+from surya.settings import settings
+
+
+class CLILoader:
+    def __init__(self, filepath: str, cli_options: dict, highres: bool = False):
+        self.page_range = cli_options.get("page_range")
+        if self.page_range:
+            self.page_range = self.parse_range_str(self.page_range)
+        self.filepath = filepath
+        self.config = cli_options
+        self.save_images = cli_options.get("images", False)
+        self.debug = cli_options.get("debug", False)
+        self.output_dir = cli_options.get("output_dir")
+
+        # Opt in to leaving the inference server up so later commands reuse it.
+        if cli_options.get("keep_server"):
+            settings.SURYA_INFERENCE_KEEP_ALIVE = True
+
+        self.load(highres)
+
+    @staticmethod
+    def common_options(fn):
+        fn = click.argument("input_path", type=click.Path(exists=True), required=True)(
+            fn
+        )
+        fn = click.option(
+            "--output_dir",
+            type=click.Path(exists=False),
+            required=False,
+            default=os.path.join(settings.RESULT_DIR, "surya"),
+            help="Directory to save output.",
+        )(fn)
+        fn = click.option(
+            "--page_range",
+            type=str,
+            default=None,
+            help="Page range to convert, specify comma separated page numbers or ranges.  Example: 0,5-10,20",
+        )(fn)
+        fn = click.option(
+            "--images",
+            is_flag=True,
+            help="Save images of detected bboxes.",
+            default=False,
+        )(fn)
+        fn = click.option(
+            "--debug", "-d", is_flag=True, help="Enable debug mode.", default=False
+        )(fn)
+        fn = click.option(
+            "--keep_server",
+            is_flag=True,
+            default=False,
+            help="Keep the inference server (vllm/llama.cpp) running after this command exits so later commands reuse it instead of re-spawning.",
+        )(fn)
+        return fn
+
+    def load(self, highres: bool = False):
+        highres_images = None
+        if os.path.isdir(self.filepath):
+            images, names = load_from_folder(self.filepath, self.page_range)
+            folder_name = os.path.basename(self.filepath)
+            if highres:
+                highres_images, _ = load_from_folder(
+                    self.filepath, self.page_range, settings.IMAGE_DPI_HIGHRES
+                )
+        else:
+            images, names = load_from_file(self.filepath, self.page_range)
+            folder_name = os.path.basename(self.filepath).split(".")[0]
+            if highres:
+                highres_images, _ = load_from_file(
+                    self.filepath, self.page_range, settings.IMAGE_DPI_HIGHRES
+                )
+
+        self.images = images
+        self.highres_images = highres_images
+        self.names = names
+
+        self.result_path = os.path.abspath(os.path.join(self.output_dir, folder_name))
+        os.makedirs(self.result_path, exist_ok=True)
+
+    @staticmethod
+    def parse_range_str(range_str: str) -> List[int]:
+        range_lst = range_str.split(",")
+        page_lst = []
+        for i in range_lst:
+            if "-" in i:
+                start, end = i.split("-")
+                page_lst += list(range(int(start), int(end) + 1))
+            else:
+                page_lst.append(int(i))
+        page_lst = sorted(
+            list(set(page_lst))
+        )  # Deduplicate page numbers and sort in order
+        return page_lst

surya/scripts/detect_layout.py

@@ -0,0 +1,62 @@
+import time
+import click
+import copy
+import json
+from collections import defaultdict
+
+from surya.inference import SuryaInferenceManager
+from surya.layout import LayoutPredictor
+from surya.debug.draw import draw_polys_on_image
+from surya.logging import configure_logging, get_logger
+from surya.scripts.config import CLILoader
+import os
+
+configure_logging()
+logger = get_logger()
+
+
+@click.command(help="Detect layout of an input file or folder (PDFs or image).")
+@CLILoader.common_options
+def detect_layout_cli(input_path: str, **kwargs):
+    loader = CLILoader(input_path, kwargs)
+
+    manager = SuryaInferenceManager()
+    layout_predictor = LayoutPredictor(manager)
+
+    start = time.time()
+    layout_predictions = layout_predictor(loader.images)
+
+    if loader.debug:
+        logger.debug(f"Layout took {time.time() - start} seconds")
+
+    if loader.save_images:
+        for idx, (image, layout_pred, name) in enumerate(
+            zip(loader.images, layout_predictions, loader.names)
+        ):
+            polygons = [p.polygon for p in layout_pred.bboxes]
+            labels = [f"{p.label}-{p.position}" for p in layout_pred.bboxes]
+            bbox_image = draw_polys_on_image(
+                polygons, copy.deepcopy(image), labels=labels
+            )
+            bbox_image.save(
+                os.path.join(loader.result_path, f"{name}_{idx}_layout.png")
+            )
+
+    predictions_by_page = defaultdict(list)
+    for idx, (pred, name, image) in enumerate(
+        zip(layout_predictions, loader.names, loader.images)
+    ):
+        out_pred = pred.model_dump()
+        out_pred["page"] = len(predictions_by_page[name]) + 1
+        predictions_by_page[name].append(out_pred)
+
+    with open(
+        os.path.join(loader.result_path, "results.json"), "w+", encoding="utf-8"
+    ) as f:
+        json.dump(predictions_by_page, f, ensure_ascii=False)
+
+    logger.info(f"Wrote results to {loader.result_path}")
+
+
+if __name__ == "__main__":
+    detect_layout_cli()

surya/scripts/detect_text.py

@@ -0,0 +1,59 @@
+import click
+import copy
+import json
+import time
+from collections import defaultdict
+
+from surya.detection import DetectionPredictor
+from surya.debug.draw import draw_polys_on_image
+from surya.logging import configure_logging, get_logger
+from surya.scripts.config import CLILoader
+import os
+
+configure_logging()
+logger = get_logger()
+
+
+@click.command(help="Detect bboxes in an input file or folder (PDFs or image).")
+@CLILoader.common_options
+def detect_text_cli(input_path: str, **kwargs):
+    loader = CLILoader(input_path, kwargs)
+
+    det_predictor = DetectionPredictor()
+
+    start = time.time()
+    predictions = det_predictor(loader.images, include_maps=loader.debug)
+    end = time.time()
+    if loader.debug:
+        logger.debug(f"Detection took {end - start} seconds")
+
+    if loader.save_images:
+        for idx, (image, pred, name) in enumerate(
+            zip(loader.images, predictions, loader.names)
+        ):
+            polygons = [p.polygon for p in pred.bboxes]
+            bbox_image = draw_polys_on_image(polygons, copy.deepcopy(image))
+            bbox_image.save(os.path.join(loader.result_path, f"{name}_{idx}_bbox.png"))
+
+            if loader.debug:
+                heatmap = pred.heatmap
+                heatmap.save(os.path.join(loader.result_path, f"{name}_{idx}_heat.png"))
+
+    predictions_by_page = defaultdict(list)
+    for idx, (pred, name, image) in enumerate(
+        zip(predictions, loader.names, loader.images)
+    ):
+        out_pred = pred.model_dump(exclude=["heatmap", "affinity_map"])
+        out_pred["page"] = len(predictions_by_page[name]) + 1
+        predictions_by_page[name].append(out_pred)
+
+    with open(
+        os.path.join(loader.result_path, "results.json"), "w+", encoding="utf-8"
+    ) as f:
+        json.dump(predictions_by_page, f, ensure_ascii=False)
+
+    logger.info(f"Wrote results to {loader.result_path}")
+
+
+if __name__ == "__main__":
+    detect_text_cli()

surya/scripts/ocr_text.py

@@ -0,0 +1,48 @@
+import os
+import click
+import json
+import time
+from collections import defaultdict
+
+from surya.inference import SuryaInferenceManager
+from surya.logging import configure_logging, get_logger
+from surya.recognition import RecognitionPredictor
+from surya.scripts.config import CLILoader
+
+configure_logging()
+logger = get_logger()
+
+
+@click.command(help="OCR text — full-page OCR (one VLM call per page).")
+@CLILoader.common_options
+def ocr_text_cli(input_path: str, **kwargs):
+    # Full-page OCR is the default path: one VLM call per page returns layout
+    # + content together. Pages whose full-page output fails to parse fall
+    # back to layout + per-block OCR automatically (see RecognitionPredictor).
+    loader = CLILoader(input_path, kwargs, highres=True)
+
+    manager = SuryaInferenceManager()
+    rec_predictor = RecognitionPredictor(manager)
+
+    start = time.time()
+    page_results = rec_predictor(loader.highres_images, full_page=True)
+
+    if loader.debug:
+        logger.debug(f"OCR took {time.time() - start:.2f} seconds")
+
+    out_preds = defaultdict(list)
+    for name, page in zip(loader.names, page_results):
+        out_pred = page.model_dump()
+        out_pred["page"] = len(out_preds[name]) + 1
+        out_preds[name].append(out_pred)
+
+    with open(
+        os.path.join(loader.result_path, "results.json"), "w+", encoding="utf-8"
+    ) as f:
+        json.dump(out_preds, f, ensure_ascii=False)
+
+    logger.info(f"Wrote results to {loader.result_path}")
+
+
+if __name__ == "__main__":
+    ocr_text_cli()

surya/scripts/run_streamlit_app.py

@@ -0,0 +1,9 @@
+import subprocess
+import os
+
+
+def streamlit_app_cli():
+    cur_dir = os.path.dirname(os.path.abspath(__file__))
+    ocr_app_path = os.path.join(cur_dir, "streamlit_app.py")
+    cmd = ["streamlit", "run", ocr_app_path, "--server.fileWatcherType", "none", "--server.headless", "true"]
+    subprocess.run(cmd, env={**os.environ, "IN_STREAMLIT": "true"})

surya/scripts/screenshot_app.py

@@ -0,0 +1,226 @@
+"""Screenshot-friendly Surya viewer.
+
+Shows a PDF/image page on the left and full-page OCR output on the right, side
+by side, for clean screenshots. You can scroll through pages and preview them
+before running OCR, then export the side-by-side view as a PNG.
+
+Run with `surya_screenshot`, then open http://localhost:8504.
+"""
+
+from __future__ import annotations
+
+import base64
+import io
+import os
+import tempfile
+import uuid
+from typing import List, Optional
+
+import pypdfium2
+from flask import Flask, jsonify, render_template, request
+from PIL import Image
+from werkzeug.utils import secure_filename
+
+from surya.inference import SuryaInferenceManager
+from surya.logging import configure_logging, get_logger
+from surya.recognition import RecognitionPredictor
+from surya.recognition.schema import PageOCRResult
+from surya.settings import settings
+
+configure_logging()
+logger = get_logger()
+
+app = Flask(__name__)
+
+ALLOWED_EXT = {".pdf", ".png", ".jpg", ".jpeg", ".gif", ".webp"}
+UPLOAD_DIR = os.path.join(tempfile.gettempdir(), "surya_screenshot")
+os.makedirs(UPLOAD_DIR, exist_ok=True)
+
+_rec: Optional[RecognitionPredictor] = None
+
+
+def get_rec() -> RecognitionPredictor:
+    """Lazily build the recognition predictor (shared inference manager)."""
+    global _rec
+    if _rec is None:
+        _rec = RecognitionPredictor(SuryaInferenceManager())
+    return _rec
+
+
+# Datalab-flavored palette for layout block overlays, keyed by canonical label.
+LABEL_COLORS = {
+    "Text": "#2563eb",
+    "SectionHeader": "#0ea5e9",
+    "PageHeader": "#7c3aed",
+    "PageFooter": "#7c3aed",
+    "Caption": "#c026d3",
+    "Footnote": "#64748b",
+    "Equation": "#9333ea",
+    "Table": "#f59e0b",
+    "TableOfContents": "#f59e0b",
+    "Form": "#ea580c",
+    "ListGroup": "#10b981",
+    "Picture": "#db2777",
+    "Figure": "#db2777",
+    "Diagram": "#db2777",
+    "Code": "#0d9488",
+    "default": "#ef4444",
+}
+
+
+def _logo_data_url() -> str:
+    path = os.path.join(settings.BASE_DIR, "static", "datalab-logo.png")
+    try:
+        with open(path, "rb") as f:
+            return "data:image/png;base64," + base64.b64encode(f.read()).decode()
+    except Exception:
+        return ""
+
+
+def _pil_to_data_url(img: Image.Image, fmt: str = "PNG") -> str:
+    buf = io.BytesIO()
+    img.save(buf, format=fmt)
+    return (
+        f"data:image/{fmt.lower()};base64," + base64.b64encode(buf.getvalue()).decode()
+    )
+
+
+def _is_pdf(path: str) -> bool:
+    return path.lower().endswith(".pdf")
+
+
+def _page_count(path: str) -> int:
+    if _is_pdf(path):
+        doc = pypdfium2.PdfDocument(path)
+        n = len(doc)
+        doc.close()
+        return n
+    return 1
+
+
+def _render_page(path: str, page: int, dpi: int) -> Image.Image:
+    """Render a 0-indexed page of a PDF (or load an image file) as RGB."""
+    if _is_pdf(path):
+        doc = pypdfium2.PdfDocument(path)
+        try:
+            pil = doc[page].render(scale=dpi / 72).to_pil().convert("RGB")
+        finally:
+            doc.close()
+        return pil
+    return Image.open(path).convert("RGB")
+
+
+def _assemble_page_html(page: PageOCRResult) -> str:
+    """Whole-page HTML from a PageOCRResult (math stays in <math> tags)."""
+    parts: List[str] = []
+    for blk in page.blocks:
+        if blk.skipped:
+            continue
+        x0, y0, x1, y1 = (int(c) for c in blk.bbox)
+        parts.append(
+            f'<div data-bbox="{x0} {y0} {x1} {y1}" '
+            f'data-label="{blk.label}">{blk.html or ""}</div>'
+        )
+    return "\n".join(parts)
+
+
+@app.route("/")
+def index():
+    return render_template("surya_screenshot.html", logo=_logo_data_url())
+
+
+@app.route("/info", methods=["POST"])
+def info():
+    path = (request.json or {}).get("file_path", "").strip()
+    if not path:
+        return jsonify({"error": "file_path is required"}), 400
+    if not os.path.exists(path):
+        return jsonify({"error": f"File not found: {path}"}), 400
+    try:
+        return jsonify({"page_count": _page_count(path)})
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+
+@app.route("/upload", methods=["POST"])
+def upload():
+    """Accept a drag/drop (or browsed) file, save to a temp path, return it."""
+    f = request.files.get("file")
+    if f is None or not f.filename:
+        return jsonify({"error": "no file uploaded"}), 400
+    ext = os.path.splitext(f.filename)[1].lower()
+    if ext not in ALLOWED_EXT:
+        return jsonify({"error": f"unsupported file type: {ext or '(none)'}"}), 400
+    safe = secure_filename(f.filename) or f"upload{ext}"
+    dest = os.path.join(UPLOAD_DIR, f"{uuid.uuid4().hex}_{safe}")
+    f.save(dest)
+    try:
+        return jsonify(
+            {"file_path": dest, "page_count": _page_count(dest), "name": f.filename}
+        )
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+
+@app.route("/page", methods=["POST"])
+def page():
+    """Render a single page for preview (no OCR)."""
+    data = request.json or {}
+    path = data.get("file_path", "").strip()
+    page_num = int(data.get("page", 0))
+    if not path or not os.path.exists(path):
+        return jsonify({"error": "valid file_path is required"}), 400
+    try:
+        img = _render_page(path, page_num, settings.IMAGE_DPI_HIGHRES)
+        return jsonify(
+            {
+                "image_base64": _pil_to_data_url(img),
+                "width": img.size[0],
+                "height": img.size[1],
+            }
+        )
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+
+@app.route("/process", methods=["POST"])
+def process():
+    """Run full-page OCR on one page; return the page image + OCR HTML + blocks."""
+    data = request.json or {}
+    path = data.get("file_path", "").strip()
+    page_num = int(data.get("page", 0))
+    if not path or not os.path.exists(path):
+        return jsonify({"error": "valid file_path is required"}), 400
+    try:
+        img = _render_page(path, page_num, settings.IMAGE_DPI_HIGHRES)
+        page_result = get_rec()([img], full_page=True)[0]
+        blocks = [
+            {
+                "bbox": [int(c) for c in blk.bbox],
+                "label": blk.label,
+                "color": LABEL_COLORS.get(blk.label, LABEL_COLORS["default"]),
+            }
+            for blk in page_result.blocks
+            if not blk.skipped
+        ]
+        return jsonify(
+            {
+                "image_base64": _pil_to_data_url(img),
+                "width": img.size[0],
+                "height": img.size[1],
+                "html": _assemble_page_html(page_result),
+                "blocks": blocks,
+                "n_blocks": len(page_result.blocks),
+            }
+        )
+    except Exception as e:
+        logger.exception("Full-page OCR failed")
+        return jsonify({"error": str(e)}), 500
+
+
+def main():
+    app.run(host="0.0.0.0", port=8504)
+
+
+if __name__ == "__main__":
+    main()

surya/scripts/streamlit_app.py

@@ -0,0 +1,490 @@
+"""Surya2 streamlit app — exercise layout, recognition, table_rec via the
+inference manager. Detection + OCR-error stay in their own torch paths."""
+
+from __future__ import annotations
+
+import io
+import re
+import tempfile
+import time
+from typing import List
+
+import pypdfium2
+import streamlit as st
+import streamlit.components.v1 as components
+from PIL import Image, ImageDraw
+
+from surya.debug.draw import draw_polys_on_image, draw_bboxes_on_image
+from surya.detection import TextDetectionResult
+from surya.inference import SuryaInferenceManager
+from surya.layout import LayoutPredictor
+from surya.layout.schema import LayoutResult
+from surya.recognition import RecognitionPredictor
+from surya.recognition.schema import PageOCRResult
+from surya.settings import settings
+from surya.table_rec import TableRecPredictor
+from surya.table_rec.schema import TableResult
+
+
+# KaTeX-enabled HTML wrapper. The OCR HTML wraps math in <math>...</math>
+# (KaTeX-compatible LaTeX inside), which a browser would otherwise show as
+# raw text. We convert those tags to \( \) / \[ \] delimiters and let KaTeX
+# auto-render typeset them inside an iframe component.
+_KATEX_HEAD = r"""<!doctype html><html><head>
+<meta charset="utf-8">
+<link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.11/dist/katex.min.css">
+<script src="https://cdn.jsdelivr.net/npm/katex@0.16.11/dist/katex.min.js"></script>
+<script src="https://cdn.jsdelivr.net/npm/katex@0.16.11/dist/contrib/auto-render.min.js"></script>
+<style>
+/* White "paper" card so the text stays readable in both light and dark
+   Streamlit themes (the iframe is otherwise transparent and our text is dark). */
+html,body{background:#ffffff;}
+body{font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",sans-serif;font-size:15px;line-height:1.55;color:#111111;margin:0;padding:14px;}
+table{border-collapse:collapse;margin:6px 0;} td,th{border:1px solid #bbb;padding:3px 6px;color:#111111;}
+[data-label="SectionHeader"],[data-label="PageHeader"]{font-weight:600;}
+</style></head><body>
+"""
+
+_KATEX_TAIL = r"""
+<script>
+renderMathInElement(document.body, {
+  delimiters: [
+    {left: "\\[", right: "\\]", display: true},
+    {left: "\\(", right: "\\)", display: false}
+  ],
+  throwOnError: false
+});
+</script></body></html>
+"""
+
+_MATH_RE = re.compile(r"<math\b([^>]*)>(.*?)</math>", re.DOTALL | re.IGNORECASE)
+
+
+def _math_to_katex(html_str: str) -> str:
+    """Rewrite <math>...</math> tags into KaTeX \\( \\) / \\[ \\] delimiters."""
+
+    def repl(m: "re.Match") -> str:
+        attrs, inner = m.group(1), m.group(2)
+        if re.search(r"""display\s*=\s*["']block["']""", attrs):
+            return "\\[" + inner + "\\]"
+        return "\\(" + inner + "\\)"
+
+    return _MATH_RE.sub(repl, html_str or "")
+
+
+def render_ocr_html(html_str: str, height: int = 400) -> None:
+    """Render OCR HTML with math typeset by KaTeX (iframe component)."""
+    components.html(
+        _KATEX_HEAD + _math_to_katex(html_str) + _KATEX_TAIL,
+        height=height,
+        scrolling=True,
+    )
+
+
+def _assemble_page_html(page: PageOCRResult) -> str:
+    """Reconstruct a div-block whole-page HTML from a PageOCRResult."""
+    parts: List[str] = []
+    for blk in page.blocks:
+        if blk.skipped:
+            continue
+        x0, y0, x1, y1 = (int(c) for c in blk.bbox)
+        body = blk.html or ""
+        parts.append(
+            f'<div data-bbox="{x0} {y0} {x1} {y1}" data-label="{blk.label}">{body}</div>'
+        )
+    return "\n".join(parts)
+
+
+def _show_timing(label: str, elapsed_s: float, extra: str = "") -> None:
+    """Render a small caption with wall-clock + optional extra detail."""
+    detail = f" — {extra}" if extra else ""
+    st.caption(f"⏱ {label}: {elapsed_s * 1000:.0f} ms ({elapsed_s:.2f}s){detail}")
+
+
+@st.cache_resource()
+def load_predictors_cached():
+    manager = SuryaInferenceManager()
+    layout_predictor = LayoutPredictor(manager)
+    rec_predictor = RecognitionPredictor(manager)
+    table_rec_predictor = TableRecPredictor(manager)
+
+    # Lazy-import detection / ocr_error to keep startup snappy when the user
+    # only wants VLM modes
+    from surya.detection import DetectionPredictor
+    from surya.ocr_error import OCRErrorPredictor
+
+    return {
+        "manager": manager,
+        "layout": layout_predictor,
+        "recognition": rec_predictor,
+        "table_rec": table_rec_predictor,
+        "detection": DetectionPredictor(),
+        "ocr_error": OCRErrorPredictor(),
+    }
+
+
+def text_detection(img) -> tuple[Image.Image, TextDetectionResult, float]:
+    t = time.perf_counter()
+    text_pred = predictors["detection"]([img])[0]
+    elapsed = time.perf_counter() - t
+    text_polygons = [p.polygon for p in text_pred.bboxes]
+    det_img = draw_polys_on_image(text_polygons, img.copy())
+    return det_img, text_pred, elapsed
+
+
+def layout_detection(img) -> tuple[Image.Image, LayoutResult, float]:
+    t = time.perf_counter()
+    pred = predictors["layout"]([img])[0]
+    elapsed = time.perf_counter() - t
+    polygons = [p.polygon for p in pred.bboxes]
+    labels = [
+        f"{p.label}-{p.position}-c{p.count}-{round(p.confidence or 0, 2)}"
+        for p in pred.bboxes
+    ]
+    annotated = draw_polys_on_image(
+        polygons, img.copy(), labels=labels, label_font_size=14
+    )
+    return annotated, pred, elapsed
+
+
+def block_ocr(img) -> tuple[Image.Image, PageOCRResult, LayoutResult, float, float]:
+    """Layout → block crops → BLOCK_PROMPT. Returns layout + block-OCR timings."""
+    t_layout = time.perf_counter()
+    layout = predictors["layout"]([img])[0]
+    layout_elapsed = time.perf_counter() - t_layout
+
+    t_blocks = time.perf_counter()
+    page_results = predictors["recognition"]([img], [layout])
+    blocks_elapsed = time.perf_counter() - t_blocks
+    page = page_results[0]
+
+    annotated = img.copy()
+    draw = ImageDraw.Draw(annotated)
+    for blk in page.blocks:
+        x0, y0, x1, y1 = blk.bbox
+        color = "red" if blk.error else ("orange" if blk.skipped else "green")
+        draw.rectangle((x0, y0, x1, y1), outline=color, width=3)
+        draw.text((x0 + 4, y0 + 4), f"{blk.reading_order} {blk.label}", fill=color)
+    return annotated, page, layout, layout_elapsed, blocks_elapsed
+
+
+def full_page_ocr(img) -> tuple[Image.Image, PageOCRResult, float]:
+    """Single HIGH_ACCURACY_BBOX_PROMPT call on the whole page."""
+    t = time.perf_counter()
+    page_results = predictors["recognition"]([img], full_page=True)
+    elapsed = time.perf_counter() - t
+    page = page_results[0]
+    annotated = img.copy()
+    draw = ImageDraw.Draw(annotated)
+    for blk in page.blocks:
+        x0, y0, x1, y1 = blk.bbox
+        color = "red" if blk.error else ("orange" if blk.skipped else "green")
+        draw.rectangle((x0, y0, x1, y1), outline=color, width=3)
+        draw.text((x0 + 4, y0 + 4), f"{blk.reading_order} {blk.label}", fill=color)
+    return annotated, page, elapsed
+
+
+def table_recognition(
+    img: Image.Image,
+    mode: str,
+    skip_table_detection: bool,
+) -> tuple[Image.Image, List[TableResult], float, float]:
+    """Returns (annotated_img, table_preds, layout_elapsed, table_rec_elapsed)."""
+    layout_elapsed = 0.0
+    if skip_table_detection:
+        table_imgs = [img]
+        table_counts = [0]
+        table_bboxes = [(0, 0, img.size[0], img.size[1])]
+    else:
+        t = time.perf_counter()
+        layout = predictors["layout"]([img])[0]
+        layout_elapsed = time.perf_counter() - t
+        tables = [b for b in layout.bboxes if b.label in ("Table", "TableOfContents")]
+        if not tables:
+            return img.copy(), [], layout_elapsed, 0.0
+        table_bboxes = [tuple(int(c) for c in b.bbox) for b in tables]
+        table_imgs = [img.crop(b) for b in table_bboxes]
+        table_counts = [b.count for b in tables]
+
+    t = time.perf_counter()
+    if mode == "full":
+        table_preds = predictors["table_rec"].predict_full(
+            table_imgs, counts=table_counts
+        )
+    else:
+        table_preds = predictors["table_rec"].predict_simple(table_imgs)
+    table_rec_elapsed = time.perf_counter() - t
+
+    out_img = img.copy()
+    for pred, table_img, tbbox in zip(table_preds, table_imgs, table_bboxes):
+        if pred.error or pred.mode != "simple" or not pred.rows:
+            continue
+        row_bboxes = [r.bbox for r in pred.rows]
+        col_bboxes = [c.bbox for c in pred.cols]
+        row_labels = [r.label for r in pred.rows]
+        col_labels = [c.label for c in pred.cols]
+        annot = table_img.copy()
+        annot = draw_bboxes_on_image(
+            row_bboxes, annot, labels=row_labels, label_font_size=14, color="blue"
+        )
+        annot = draw_bboxes_on_image(
+            col_bboxes, annot, labels=col_labels, label_font_size=14, color="red"
+        )
+        # Paste annotated crop back at the table's position in the page.
+        out_img.paste(annot, (tbbox[0], tbbox[1]))
+    return out_img, table_preds, layout_elapsed, table_rec_elapsed
+
+
+def ocr_errors(pdf_file, page_count, sample_len=512, max_samples=10, max_pages=15):
+    from pdftext.extraction import plain_text_output
+
+    with tempfile.NamedTemporaryFile(suffix=".pdf") as f:
+        f.write(pdf_file.getvalue())
+        f.seek(0)
+
+        page_middle = page_count // 2
+        page_range = range(
+            max(page_middle - max_pages, 0), min(page_middle + max_pages, page_count)
+        )
+        text = plain_text_output(f.name, page_range=page_range)
+
+    sample_gap = len(text) // max_samples
+    if len(text) == 0 or sample_gap == 0:
+        return "This PDF has no text or very little text", ["no text"]
+
+    if sample_gap < sample_len:
+        sample_gap = sample_len
+
+    samples = []
+    for i in range(0, len(text), sample_gap):
+        samples.append(text[i : i + sample_len])
+
+    results = predictors["ocr_error"](samples)
+    label = "This PDF has good text."
+    if results.labels.count("bad") / len(results.labels) > 0.2:
+        label = "This PDF may have garbled or bad OCR text."
+    return label, results.labels
+
+
+def open_pdf(pdf_file):
+    stream = io.BytesIO(pdf_file.getvalue())
+    return pypdfium2.PdfDocument(stream)
+
+
+@st.cache_data()
+def get_page_image(pdf_file, page_num, dpi=settings.IMAGE_DPI):
+    doc = open_pdf(pdf_file)
+    renderer = doc.render(
+        pypdfium2.PdfBitmap.to_pil,
+        page_indices=[page_num - 1],
+        scale=dpi / 72,
+    )
+    png = list(renderer)[0]
+    png_image = png.convert("RGB")
+    doc.close()
+    return png_image
+
+
+@st.cache_data()
+def page_counter(pdf_file):
+    doc = open_pdf(pdf_file)
+    doc_len = len(doc)
+    doc.close()
+    return doc_len
+
+
+st.set_page_config(layout="wide")
+col1, col2 = st.columns([0.55, 0.45])
+
+predictors = load_predictors_cached()
+
+st.markdown(
+    """
+# Surya 2 Demo
+
+VLM-backed layout, OCR, and table recognition. The model runs in a local
+`llama-server` (or vllm) process, started on first use.
+
+Modes:
+- **Layout**: page → list of blocks with label + bbox + token count
+- **Block OCR**: layout + per-block HTML
+- **Table Rec (simple)**: row + column bboxes only
+- **Table Rec (full)**: full HTML for each detected table
+"""
+)
+
+in_file = st.sidebar.file_uploader(
+    "PDF file or image:", type=["pdf", "png", "jpg", "jpeg", "gif", "webp"]
+)
+
+if in_file is None:
+    st.stop()
+
+filetype = in_file.type
+page_count = None
+if "pdf" in filetype:
+    page_count = page_counter(in_file)
+    page_number = st.sidebar.number_input(
+        f"Page number out of {page_count}:", min_value=1, value=1, max_value=page_count
+    )
+    # Render at high DPI so the OCR / table-rec demos see fine glyphs.
+    # Layout + detection internally downsample (or accept the small perf hit
+    # at demo scale); we always render and display the high-DPI page here.
+    pil_image = get_page_image(in_file, page_number, settings.IMAGE_DPI_HIGHRES)
+else:
+    pil_image = Image.open(in_file).convert("RGB")
+    page_number = None
+
+run_full_page_ocr = st.sidebar.button("Run Full-Page OCR")
+run_text_det = st.sidebar.button("Run Text Detection")
+run_layout = st.sidebar.button("Run Layout Analysis")
+run_table_rec = st.sidebar.button("Run Table Rec")
+run_block_ocr = st.sidebar.button("Run Block OCR")
+run_ocr_errors = st.sidebar.button("Run bad-PDF-text detection")
+
+table_mode = st.sidebar.radio(
+    "Table mode",
+    options=["simple", "full"],
+    index=0,
+    help="simple: rows+cols only. full: full HTML.",
+)
+skip_table_detection = st.sidebar.checkbox(
+    "Skip table detection",
+    value=False,
+    help="Treat the entire page/image as a single table.",
+)
+
+if pil_image is None:
+    st.stop()
+
+
+if run_text_det:
+    det_img, text_pred, elapsed = text_detection(pil_image)
+    with col1:
+        _show_timing("Text detection", elapsed, f"{len(text_pred.bboxes)} polys")
+        st.image(det_img, caption="Detected Text", use_container_width=True)
+        st.json(
+            text_pred.model_dump(exclude=["heatmap", "affinity_map"]), expanded=False
+        )
+
+
+if run_layout:
+    annotated, pred, elapsed = layout_detection(pil_image)
+    with col1:
+        _show_timing("Layout", elapsed, f"{len(pred.bboxes)} blocks")
+        st.image(annotated, caption="Detected Layout", use_container_width=True)
+        st.json(pred.model_dump(), expanded=False)
+
+
+if run_block_ocr:
+    annotated, page, layout, t_layout, t_blocks = block_ocr(pil_image)
+    with col1:
+        n_blocks = len(page.blocks)
+        n_ok = sum(1 for b in page.blocks if not b.skipped and not b.error)
+        _show_timing("Block OCR — layout", t_layout, f"{n_blocks} blocks")
+        _show_timing("Block OCR — per-block OCR", t_blocks, f"{n_ok} OCR'd")
+        _show_timing("Block OCR — total", t_layout + t_blocks)
+        st.image(
+            annotated,
+            caption="Block OCR (green=ok, orange=skipped, red=error)",
+            use_container_width=True,
+        )
+        full_html = _assemble_page_html(page)
+        with st.expander("Full page HTML (rendered)", expanded=False):
+            render_ocr_html(full_html, height=600)
+        with st.expander("Full page HTML (source)", expanded=False):
+            st.code(full_html, language="html")
+        for blk in page.blocks:
+            with st.expander(
+                f"#{blk.reading_order} {blk.label} (conf {blk.confidence:.2f})"
+            ):
+                # Diagnostics: show numeric bbox + polygon + a thumbnail with the
+                # drawn rectangle highlighted, then the actual crop fed to OCR.
+                xs = [p[0] for p in blk.polygon]
+                ys = [p[1] for p in blk.polygon]
+                bbox_drawn = [int(min(xs)), int(min(ys)), int(max(xs)), int(max(ys))]
+                cx0 = max(0, int(min(xs)) - 4)
+                cy0 = max(0, int(min(ys)) - 4)
+                cx1 = min(pil_image.size[0], int(max(xs)) + 4)
+                cy1 = min(pil_image.size[1], int(max(ys)) + 4)
+                st.text(
+                    f"bbox(drawn) = {bbox_drawn}\n"
+                    f"crop(ocr)  = {(cx0, cy0, cx1, cy1)}  (= bbox ± 4px pad)"
+                )
+                # Thumbnail with this block's rectangle highlighted in red.
+                thumb = pil_image.copy()
+                ImageDraw.Draw(thumb).rectangle(bbox_drawn, outline="red", width=4)
+                st.image(thumb, caption="this block's drawn rect (red)", width=300)
+                # The actual crop fed to OCR
+                if cx1 > cx0 and cy1 > cy0:
+                    st.image(pil_image.crop((cx0, cy0, cx1, cy1)), caption="OCR crop")
+                if blk.skipped:
+                    st.info("Block skipped (visual label)")
+                elif blk.error:
+                    st.error("Block OCR errored")
+                else:
+                    render_ocr_html(blk.html, height=160)
+                    st.code(blk.html, language="html")
+
+
+if run_full_page_ocr:
+    annotated, page, elapsed = full_page_ocr(pil_image)
+    with col1:
+        n_blocks = len(page.blocks)
+        n_ok = sum(1 for b in page.blocks if not b.skipped and not b.error)
+        _show_timing("Full-Page OCR", elapsed, f"{n_blocks} blocks parsed, {n_ok} OK")
+        st.image(
+            annotated,
+            caption="Full-Page OCR (green=ok, orange=skipped, red=error)",
+            use_container_width=True,
+        )
+        full_html = _assemble_page_html(page)
+        with st.expander("Full page HTML (rendered)", expanded=False):
+            render_ocr_html(full_html, height=600)
+        with st.expander("Full page HTML (source)", expanded=False):
+            st.code(full_html, language="html")
+        for blk in page.blocks:
+            with st.expander(
+                f"#{blk.reading_order} {blk.label} (conf {blk.confidence:.2f})"
+            ):
+                if blk.skipped:
+                    st.info("Block skipped (visual label)")
+                elif blk.error:
+                    st.error("Block OCR errored")
+                else:
+                    render_ocr_html(blk.html, height=160)
+                    st.code(blk.html, language="html")
+
+
+if run_table_rec:
+    table_img, preds, t_layout, t_table = table_recognition(
+        pil_image, table_mode, skip_table_detection
+    )
+    with col1:
+        if not skip_table_detection:
+            _show_timing("Table Rec — layout", t_layout, f"{len(preds)} tables found")
+        _show_timing(f"Table Rec — {table_mode}", t_table)
+        if not skip_table_detection:
+            _show_timing("Table Rec — total", t_layout + t_table)
+        st.image(table_img, caption="Table Recognition", use_container_width=True)
+        for pred in preds:
+            if pred.mode == "full" and pred.html:
+                with st.expander("Table HTML"):
+                    render_ocr_html(pred.html, height=400)
+                    st.code(pred.html, language="html")
+            else:
+                st.json(pred.model_dump(), expanded=False)
+
+
+if run_ocr_errors:
+    if "pdf" not in filetype:
+        st.error("This feature only works with PDFs.")
+    else:
+        label, results = ocr_errors(in_file, page_count)
+        with col1:
+            st.write(label)
+            st.json(results)
+
+
+with col2:
+    st.image(pil_image, caption="Uploaded Image", use_container_width=True)

surya/scripts/table_recognition.py

@@ -0,0 +1,139 @@
+import os
+import click
+import copy
+import json
+from collections import defaultdict
+
+from surya.common.util import expand_bbox
+from surya.debug.draw import draw_bboxes_on_image
+from surya.inference import SuryaInferenceManager
+from surya.layout import LayoutPredictor
+from surya.logging import configure_logging, get_logger
+from surya.scripts.config import CLILoader
+from surya.table_rec import TableRecPredictor
+
+configure_logging()
+logger = get_logger()
+
+
+@click.command(help="Run table recognition on an input file or folder.")
+@CLILoader.common_options
+@click.option(
+    "--skip_table_detection",
+    is_flag=True,
+    help="Tables are already cropped, so don't re-detect tables.",
+    default=False,
+)
+@click.option(
+    "--mode",
+    type=click.Choice(["simple", "full"]),
+    default="simple",
+    help="simple: rows+cols only (geometric cells). full: full HTML (BLOCK_PROMPT).",
+)
+def table_recognition_cli(
+    input_path: str, skip_table_detection: bool, mode: str, **kwargs
+):
+    # Layout runs on the low-DPI render; table crops come from the high-DPI
+    # image so the table_rec model sees readable cell content.
+    loader = CLILoader(input_path, kwargs, highres=True)
+
+    manager = SuryaInferenceManager()
+    layout_predictor = LayoutPredictor(manager)
+    table_rec_predictor = TableRecPredictor(manager)
+
+    pnums = []
+    prev_name = None
+    for name in loader.names:
+        if prev_name is None or prev_name != name:
+            pnums.append(0)
+        else:
+            pnums.append(pnums[-1] + 1)
+        prev_name = name
+
+    table_imgs = []
+    table_counts = []
+    table_counts_per_img = []
+
+    if skip_table_detection:
+        for img in loader.highres_images:
+            table_imgs.append(img)
+            table_counts.append(1)
+            table_counts_per_img.append(0)
+    else:
+        layout_predictions = layout_predictor(
+            loader.images,
+            target_image_sizes=[img.size for img in loader.highres_images],
+        )
+        for layout_pred, img in zip(layout_predictions, loader.highres_images):
+            tables_on_page = [
+                line
+                for line in layout_pred.bboxes
+                if line.label in ("Table", "TableOfContents")
+            ]
+            table_counts.append(len(tables_on_page))
+            for line in tables_on_page:
+                bbox = expand_bbox(line.bbox)
+                table_imgs.append(img.crop(bbox))
+                table_counts_per_img.append(line.count)
+
+    table_preds = table_rec_predictor(table_imgs, mode=mode)
+
+    img_idx = 0
+    prev_count = 0
+    table_predictions = defaultdict(list)
+    for i in range(sum(table_counts)):
+        while i >= prev_count + table_counts[img_idx]:
+            prev_count += table_counts[img_idx]
+            img_idx += 1
+
+        pred = table_preds[i]
+        orig_name = loader.names[img_idx]
+        pnum = pnums[img_idx]
+        table_img = table_imgs[i]
+
+        out_pred = pred.model_dump()
+        out_pred["page"] = pnum + 1
+        table_idx = i - prev_count
+        out_pred["table_idx"] = table_idx
+        table_predictions[orig_name].append(out_pred)
+
+        if loader.save_images and pred.rows:
+            rows = [line.bbox for line in pred.rows]
+            cols = [line.bbox for line in pred.cols]
+            row_labels = [f"Row {line.row_id}" for line in pred.rows]
+            col_labels = [f"Col {line.col_id}" for line in pred.cols]
+            cells = [line.bbox for line in pred.cells]
+
+            rc_image = copy.deepcopy(table_img)
+            rc_image = draw_bboxes_on_image(
+                rows, rc_image, labels=row_labels, label_font_size=20, color="blue"
+            )
+            rc_image = draw_bboxes_on_image(
+                cols, rc_image, labels=col_labels, label_font_size=20, color="red"
+            )
+            rc_image.save(
+                os.path.join(
+                    loader.result_path,
+                    f"{orig_name}_page{pnum + 1}_table{table_idx}_rc.png",
+                )
+            )
+
+            cell_image = copy.deepcopy(table_img)
+            cell_image = draw_bboxes_on_image(cells, cell_image, color="green")
+            cell_image.save(
+                os.path.join(
+                    loader.result_path,
+                    f"{orig_name}_page{pnum + 1}_table{table_idx}_cells.png",
+                )
+            )
+
+    with open(
+        os.path.join(loader.result_path, "results.json"), "w+", encoding="utf-8"
+    ) as f:
+        json.dump(table_predictions, f, ensure_ascii=False)
+
+    logger.info(f"Wrote results to {loader.result_path}")
+
+
+if __name__ == "__main__":
+    table_recognition_cli()

surya/scripts/templates/surya_screenshot.html

@@ -0,0 +1,331 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+  <meta charset="UTF-8">
+  <meta name="viewport" content="width=device-width, initial-scale=1.0">
+  <title>Surya · Full-Page OCR</title>
+  <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.11/dist/katex.min.css">
+  <script defer src="https://cdn.jsdelivr.net/npm/katex@0.16.11/dist/katex.min.js"></script>
+  <script src="https://cdn.jsdelivr.net/npm/html2canvas@1.4.1/dist/html2canvas.min.js"></script>
+  <style>
+    :root {
+      --ink: #0f1115;
+      --bg: #f4f6f9;
+      --panel: #ffffff;
+      --border: #e4e7ec;
+      --accent: #2563eb;
+      --accent-hover: #1d4ed8;
+      --muted: #667085;
+    }
+    * { margin: 0; padding: 0; box-sizing: border-box; }
+    body {
+      font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
+      background: var(--bg);
+      color: var(--ink);
+      height: 100vh;
+      display: flex;
+      flex-direction: column;
+      overflow: hidden;
+    }
+    header {
+      display: flex; align-items: center; gap: 16px;
+      padding: 12px 20px;
+      background: var(--panel);
+      border-bottom: 1px solid var(--border);
+      box-shadow: 0 1px 3px rgba(16,24,40,.04);
+      z-index: 10;
+    }
+    .brand { display: flex; align-items: center; gap: 10px; }
+    .brand img { height: 30px; width: auto; }
+    .brand .title { font-size: 18px; font-weight: 700; letter-spacing: -.01em; }
+    .brand .sub { font-size: 13px; color: var(--muted); font-weight: 500;
+      padding-left: 10px; margin-left: 4px; border-left: 1px solid var(--border); }
+    .controls { display: flex; align-items: center; gap: 10px; flex: 1; flex-wrap: wrap; }
+    input[type=text] {
+      flex: 1; min-width: 260px; max-width: 560px;
+      padding: 9px 12px; font-size: 14px;
+      border: 1px solid var(--border); border-radius: 8px; background: #fff;
+    }
+    input[type=text]:focus { outline: none; border-color: var(--accent); box-shadow: 0 0 0 3px rgba(37,99,235,.12); }
+    button {
+      padding: 9px 14px; font-size: 14px; font-weight: 600;
+      border: none; border-radius: 8px; cursor: pointer;
+      background: var(--accent); color: #fff; transition: background .15s;
+    }
+    button:hover:not(:disabled) { background: var(--accent-hover); }
+    button.secondary { background: #fff; color: var(--ink); border: 1px solid var(--border); }
+    button.secondary:hover:not(:disabled) { background: #f2f4f7; }
+    button:disabled { opacity: .45; cursor: not-allowed; }
+    .pager { display: flex; align-items: center; gap: 6px; }
+    .pager .pageind { font-size: 13px; color: var(--muted); min-width: 96px; text-align: center; }
+    .pager button { padding: 7px 11px; }
+    .toggle { display: flex; align-items: center; gap: 7px; font-size: 13px; color: var(--muted); cursor: pointer; user-select: none; }
+    .status { font-size: 13px; font-weight: 600; min-width: 80px; }
+    .status.loading { color: #b45309; }
+    .status.error { color: #d92d20; }
+    .status.ok { color: #079455; }
+
+    .stage {
+      flex: 1; display: flex; gap: 16px; padding: 16px; overflow: hidden;
+    }
+    .panel {
+      flex: 1; min-width: 0; display: flex; flex-direction: column;
+      background: var(--panel); border: 1px solid var(--border);
+      border-radius: 12px; overflow: hidden;
+      box-shadow: 0 1px 2px rgba(16,24,40,.05);
+    }
+    .panel-head {
+      padding: 11px 16px; font-size: 13px; font-weight: 700;
+      letter-spacing: .02em; text-transform: uppercase; color: var(--muted);
+      border-bottom: 1px solid var(--border); background: #fcfcfd;
+    }
+    .panel-body { flex: 1; overflow: auto; }
+    .img-wrap {
+      display: flex; align-items: flex-start; justify-content: center;
+      padding: 16px; background: #f0f2f5; min-height: 100%;
+    }
+    #pageCanvas { max-width: 100%; height: auto; border-radius: 6px;
+      box-shadow: 0 2px 10px rgba(16,24,40,.12); background: #fff; }
+    .ocr {
+      padding: 28px 32px; line-height: 1.6; font-size: 16px; color: #1d2433;
+    }
+    .ocr [data-label="SectionHeader"], .ocr [data-label="Title"] { font-weight: 700; font-size: 1.15em; margin: .5em 0 .3em; }
+    .ocr [data-label="PageHeader"], .ocr [data-label="PageFooter"] { color: var(--muted); font-size: .9em; }
+    .ocr table { border-collapse: collapse; margin: 14px 0; width: 100%; }
+    .ocr th, .ocr td { border: 1px solid #d0d5dd; padding: 6px 10px; text-align: left; }
+    .ocr th { background: #f2f4f7; font-weight: 600; }
+    .ocr img { max-width: 100%; height: auto; }
+    .placeholder { padding: 48px 32px; color: var(--muted); font-size: 15px; text-align: center; }
+    #dropOverlay { position: fixed; inset: 0; z-index: 100; display: none;
+      align-items: center; justify-content: center; background: rgba(15,17,21,.55); }
+    #dropOverlay.active { display: flex; }
+    #dropOverlay .drop-card { padding: 40px 64px; border: 3px dashed #fff;
+      border-radius: 16px; color: #fff; font-size: 22px; font-weight: 700;
+      background: rgba(37,99,235,.30); }
+    /* While screenshotting, expand panels to full content height so
+       html2canvas captures everything, not just the visible scroll area. */
+    body.capturing { height: auto !important; overflow: visible !important; }
+    body.capturing .stage { height: auto !important; overflow: visible !important; }
+    body.capturing .panel-body { height: auto !important; overflow: visible !important; }
+  </style>
+</head>
+<body>
+  <header>
+    <div class="brand">
+      {% if logo %}<img src="{{ logo }}" alt="Datalab">{% endif %}
+      <span class="title">Surya</span>
+      <span class="sub">Full-Page OCR</span>
+    </div>
+    <div class="controls">
+      <input type="text" id="filePath" placeholder="Drop a file, Browse, or type a server path…">
+      <input type="file" id="fileInput" accept=".pdf,.png,.jpg,.jpeg,.gif,.webp" style="display:none" onchange="onPick(this)">
+      <button class="secondary" onclick="document.getElementById('fileInput').click()">Browse</button>
+      <button class="secondary" id="loadBtn" onclick="loadFile()">Load</button>
+      <div class="pager">
+        <button class="secondary" id="prevBtn" onclick="changePage(-1)" disabled>◀</button>
+        <span class="pageind" id="pageInd">—</span>
+        <button class="secondary" id="nextBtn" onclick="changePage(1)" disabled>▶</button>
+      </div>
+      <button id="runBtn" onclick="runOCR()" disabled>Run Full-Page OCR</button>
+      <label class="toggle"><input type="checkbox" id="showBoxes" checked onchange="drawLeft()"> Layout boxes</label>
+      <button class="secondary" id="copyBtn" onclick="copyHtml()" disabled>Copy HTML</button>
+      <button class="secondary" id="shotBtn" onclick="saveScreenshot()" disabled>Save Screenshot</button>
+      <span class="status" id="status"></span>
+    </div>
+  </header>
+
+  <div class="stage">
+    <div class="panel">
+      <div class="panel-head">PDF Page</div>
+      <div class="panel-body"><div class="img-wrap"><canvas id="pageCanvas"></canvas></div></div>
+    </div>
+    <div class="panel">
+      <div class="panel-head">Full-Page OCR</div>
+      <div class="panel-body"><div class="ocr" id="ocr"><div class="placeholder">Load a file, scroll to a page, then run full-page OCR.</div></div></div>
+    </div>
+  </div>
+
+  <div id="dropOverlay"><div class="drop-card">Drop a PDF or image to load</div></div>
+
+  <script>
+    const S = { path: "", name: "", page: 0, count: 0, img: null, blocks: null, html: null, ocrPage: null };
+
+    const $ = (id) => document.getElementById(id);
+    function setStatus(msg, kind) { const s = $("status"); s.textContent = msg || ""; s.className = "status " + (kind || ""); }
+
+    async function post(url, body) {
+      const r = await fetch(url, { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(body) });
+      const data = await r.json();
+      if (!r.ok) throw new Error(data.error || "Request failed");
+      return data;
+    }
+
+    async function loadFile() {
+      const val = $("filePath").value.trim();
+      if (!val) { setStatus("Drop a file or enter a path", "error"); return; }
+      // If the box still shows the name of an uploaded file, just reload it.
+      if (S.path && val === S.name) { S.page = 0; await showPage(); return; }
+      setStatus("Loading…", "loading");
+      try {
+        const info = await post("/info", { file_path: val });
+        S.path = val; S.name = val; S.count = info.page_count; S.page = 0;
+        await showPage();
+        setStatus("");
+      } catch (e) { setStatus(e.message, "error"); }
+    }
+
+    function onPick(input) { if (input.files && input.files[0]) uploadFile(input.files[0]); }
+
+    async function uploadFile(file) {
+      setStatus("Uploading…", "loading");
+      const fd = new FormData(); fd.append("file", file);
+      try {
+        const r = await fetch("/upload", { method: "POST", body: fd });
+        const data = await r.json();
+        if (!r.ok) throw new Error(data.error || "Upload failed");
+        S.path = data.file_path; S.name = data.name; S.count = data.page_count; S.page = 0;
+        $("filePath").value = data.name;
+        await showPage();
+        setStatus("");
+      } catch (e) { setStatus(e.message, "error"); }
+    }
+
+    async function showPage() {
+      setStatus("Rendering…", "loading");
+      try {
+        const data = await post("/page", { file_path: S.path, page: S.page });
+        // New page → clear any previous OCR output.
+        S.blocks = null; S.html = null; S.ocrPage = null;
+        $("ocr").innerHTML = '<div class="placeholder">Run full-page OCR to see the extracted content.</div>';
+        $("shotBtn").disabled = true;
+        $("copyBtn").disabled = true;
+        loadImage(data.image_base64, () => drawLeft());
+        updatePager();
+        $("runBtn").disabled = false;
+        setStatus("");
+      } catch (e) { setStatus(e.message, "error"); }
+    }
+
+    function loadImage(src, cb) {
+      const im = new Image();
+      im.onload = () => { S.img = im; cb && cb(); };
+      im.src = src;
+    }
+
+    function updatePager() {
+      $("pageInd").textContent = S.count ? `Page ${S.page + 1} of ${S.count}` : "—";
+      $("prevBtn").disabled = S.page <= 0;
+      $("nextBtn").disabled = S.page >= S.count - 1;
+    }
+
+    function changePage(delta) {
+      const next = S.page + delta;
+      if (next < 0 || next >= S.count) return;
+      S.page = next;
+      showPage();
+    }
+
+    function drawLeft() {
+      if (!S.img) return;
+      const c = $("pageCanvas"), ctx = c.getContext("2d");
+      c.width = S.img.naturalWidth; c.height = S.img.naturalHeight;
+      ctx.drawImage(S.img, 0, 0);
+      if (!S.blocks || !$("showBoxes").checked) return;
+      ctx.lineWidth = 3;
+      ctx.font = 'bold 15px -apple-system, "Segoe UI", sans-serif';
+      ctx.textBaseline = "top";
+      S.blocks.forEach((b) => {
+        const [x1, y1, x2, y2] = b.bbox;
+        ctx.strokeStyle = b.color; ctx.fillStyle = b.color + "26";
+        ctx.fillRect(x1, y1, x2 - x1, y2 - y1);
+        ctx.strokeRect(x1, y1, x2 - x1, y2 - y1);
+        const tw = ctx.measureText(b.label).width, ly = Math.max(y1 - 22, 0);
+        ctx.fillStyle = b.color; ctx.fillRect(x1, ly, tw + 12, 21);
+        ctx.fillStyle = "#fff"; ctx.fillText(b.label, x1 + 6, ly + 3);
+      });
+    }
+
+    async function runOCR() {
+      if (!S.path) return;
+      setStatus("Running OCR…", "loading");
+      $("runBtn").disabled = true;
+      try {
+        const data = await post("/process", { file_path: S.path, page: S.page });
+        S.blocks = data.blocks; S.html = data.html || ""; S.ocrPage = S.page;
+        loadImage(data.image_base64, () => drawLeft());
+        const ocr = $("ocr");
+        ocr.innerHTML = data.html || '<div class="placeholder">No content detected on this page.</div>';
+        renderMath(ocr);
+        $("shotBtn").disabled = false;
+        $("copyBtn").disabled = !S.html;
+        setStatus(`${data.n_blocks} blocks`, "ok");
+      } catch (e) { setStatus(e.message, "error"); }
+      finally { $("runBtn").disabled = false; }
+    }
+
+    function renderMath(root) {
+      root.querySelectorAll("math").forEach((el) => {
+        const block = el.getAttribute("display") === "block";
+        try {
+          const span = document.createElement(block ? "div" : "span");
+          span.innerHTML = katex.renderToString(el.textContent, { displayMode: block, throwOnError: false });
+          el.replaceWith(span);
+        } catch (e) { /* leave raw on failure */ }
+      });
+    }
+
+    function copyHtml() {
+      if (!S.html) return;
+      const done = () => setStatus("HTML copied", "ok");
+      if (navigator.clipboard && navigator.clipboard.writeText) {
+        navigator.clipboard.writeText(S.html).then(done).catch(() => fallbackCopy(S.html, done));
+      } else {
+        fallbackCopy(S.html, done);
+      }
+    }
+
+    function fallbackCopy(text, done) {
+      const ta = document.createElement("textarea");
+      ta.value = text; ta.style.position = "fixed"; ta.style.opacity = "0";
+      document.body.appendChild(ta); ta.select();
+      try { document.execCommand("copy"); done(); }
+      catch (e) { setStatus("Copy failed", "error"); }
+      document.body.removeChild(ta);
+    }
+
+    async function saveScreenshot() {
+      setStatus("Capturing…", "loading");
+      const stage = document.querySelector(".stage");
+      document.body.classList.add("capturing");
+      try {
+        // Let the expanded layout settle before measuring full size.
+        await new Promise((r) => requestAnimationFrame(() => requestAnimationFrame(r)));
+        const w = stage.scrollWidth, h = stage.scrollHeight;
+        const canvas = await html2canvas(stage, {
+          backgroundColor: "#f4f6f9", scale: 2, useCORS: true, logging: false,
+          width: w, height: h, windowWidth: w, windowHeight: h, scrollX: 0, scrollY: 0,
+        });
+        const a = document.createElement("a");
+        a.href = canvas.toDataURL("image/png");
+        a.download = `surya_ocr_page_${(S.ocrPage ?? S.page) + 1}.png`;
+        a.click();
+        setStatus("Saved", "ok");
+      } catch (e) { setStatus("Screenshot failed: " + e.message, "error"); }
+      finally { document.body.classList.remove("capturing"); }
+    }
+
+    $("filePath").addEventListener("keypress", (e) => { if (e.key === "Enter") loadFile(); });
+
+    // Drag & drop anywhere in the window.
+    let dragDepth = 0;
+    window.addEventListener("dragenter", (e) => { e.preventDefault(); dragDepth++; $("dropOverlay").classList.add("active"); });
+    window.addEventListener("dragover", (e) => { e.preventDefault(); });
+    window.addEventListener("dragleave", (e) => { e.preventDefault(); if (--dragDepth <= 0) { dragDepth = 0; $("dropOverlay").classList.remove("active"); } });
+    window.addEventListener("drop", (e) => {
+      e.preventDefault(); dragDepth = 0; $("dropOverlay").classList.remove("active");
+      const f = e.dataTransfer.files && e.dataTransfer.files[0];
+      if (f) uploadFile(f);
+    });
+  </script>
+</body>
+</html>

surya/settings.py

@@ -0,0 +1,165 @@
+import os
+from typing import Callable, Dict, Optional
+
+import torch
+from dotenv import find_dotenv
+from pydantic import computed_field
+from pydantic_settings import BaseSettings
+from pathlib import Path
+from platformdirs import user_cache_dir
+
+
+class Settings(BaseSettings):
+    # General
+    TORCH_DEVICE: Optional[str] = None
+    IMAGE_DPI: int = 96  # used for layout + text detection (coarse structure)
+    IMAGE_DPI_HIGHRES: int = 192  # used for recognition + table rec (fine glyphs)
+    IN_STREAMLIT: bool = False
+    DISABLE_TQDM: bool = False
+    S3_BASE_URL: str = "https://models.datalab.to"
+    PARALLEL_DOWNLOAD_WORKERS: int = 10
+    MODEL_CACHE_DIR: str = str(Path(user_cache_dir("datalab")) / "models")
+    LOGLEVEL: str = "INFO"
+
+    # Paths
+    RESULT_DIR: str = "results"
+    BASE_DIR: str = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    FONT_DIR: str = os.path.join(BASE_DIR, "static", "fonts")
+
+    @computed_field
+    def TORCH_DEVICE_MODEL(self) -> str:
+        if self.TORCH_DEVICE is not None:
+            return self.TORCH_DEVICE
+        if torch.cuda.is_available():
+            return "cuda"
+        if torch.backends.mps.is_available():
+            return "mps"
+        return "cpu"
+
+    # ---- Surya2 inference (VLM-backed: vllm | llamacpp) ---------------------
+    SURYA_MODEL_CHECKPOINT: str = "datalab-to/surya-ocr-2"
+    SURYA_GGUF_REPO: str = "datalab-to/surya-ocr-2-gguf"
+    SURYA_GGUF_MODEL_FILE: str = "surya-2.gguf"
+    SURYA_GGUF_MMPROJ_FILE: str = "surya-2-mmproj.gguf"
+    # If set, used directly instead of HF download (handy for local-conversion testing)
+    SURYA_GGUF_LOCAL_MODEL_PATH: Optional[str] = None
+    SURYA_GGUF_LOCAL_MMPROJ_PATH: Optional[str] = None
+
+    # Backend selection
+    SURYA_INFERENCE_BACKEND: Optional[str] = None  # "vllm" | "llamacpp" | None (auto)
+    SURYA_INFERENCE_URL: Optional[str] = None  # external server, skip spawn
+    SURYA_INFERENCE_AUTOSTART: bool = True
+    # Leave an auto-spawned server running after the process exits so later
+    # commands attach to it instead of re-spawning (avoids repeated startup /
+    # model-load cost). Stop it manually when done — see `surya/inference`.
+    SURYA_INFERENCE_KEEP_ALIVE: bool = False
+    SURYA_INFERENCE_HOST: str = "127.0.0.1"
+    SURYA_INFERENCE_PORT: Optional[int] = None  # None = pick a free port
+    SURYA_INFERENCE_PARALLEL: int = 8
+    # Max concurrent in-flight chat-completion requests to the inference server
+    # per batch. Cap to roughly VLLM_MAX_NUM_SEQS so block fan-out keeps the
+    # GPU's sequence slots full without flooding the queue. Tuned empirically.
+    SURYA_INFERENCE_MAX_INFLIGHT: int = 16
+    # Per-parallel-slot KV-cache budget for the llama.cpp backend. Worst-case
+    # one OCR request: ~2k for image prefill + SURYA_MAX_TOKENS_FULL_PAGE
+    # (8192) generation + ~2k prompt/chat-template overhead ≈ 12k. Below this
+    # llama-server silently truncates outputs once a slot fills.
+    SURYA_INFERENCE_CTX_PER_SLOT: int = 12288
+    # Optional override for the *total* ctx passed to llama-server. When None
+    # (default), total = max(16384, PARALLEL * CTX_PER_SLOT). Set this only
+    # if you've hand-tuned for a specific machine.
+    SURYA_INFERENCE_CTX_SIZE: Optional[int] = None
+    SURYA_INFERENCE_TIMEOUT_SECONDS: float = 600.0
+    SURYA_INFERENCE_STARTUP_TIMEOUT: float = 600.0
+    SURYA_INFERENCE_LOGPROBS: bool = True
+    SURYA_INFERENCE_MAX_RETRIES: int = 1
+    # Force layout/table_rec output through a JSON schema via guided decoding.
+    # Eliminates malformed-JSON failures at small decode-throughput cost.
+    SURYA_GUIDED_LAYOUT: bool = True
+    # Disabled: with no minItems in TABLE_REC_JSON_SCHEMA, the constrained
+    # decoder closes the array after one element at temperature=0. The model
+    # produces well-formed JSON without the schema.
+    SURYA_GUIDED_TABLE_REC: bool = False
+
+    # Token budgets
+    SURYA_MAX_TOKENS_LAYOUT: int = 3072
+    SURYA_MAX_TOKENS_TABLE_REC: int = 3072
+    SURYA_MAX_TOKENS_BLOCK_CEILING: int = 8192
+    SURYA_MAX_TOKENS_FULL_PAGE: int = 6144
+    SURYA_MAX_BLOCKS_PER_PAGE: int = 80
+
+    BBOX_SCALE: int = 1000
+
+    # vllm
+    VLLM_DOCKER_IMAGE: str = "vllm/vllm-openai:v0.20.1"
+    VLLM_API_KEY: str = "EMPTY"
+    VLLM_GPUS: str = "0"
+    VLLM_GPU_TYPE: str = "4090"
+    # bfloat16 needs an Ampere+ GPU (compute capability >= 8.0). On older cards
+    # (e.g. T4 / Turing) vllm refuses to start with bf16 — set float16 there.
+    VLLM_DTYPE: str = "bfloat16"
+    VLLM_MAX_MODEL_LEN: int = 18000
+    VLLM_GPU_MEMORY_UTILIZATION: float = 0.85
+    # MTP speculative decoding only feeds the nested-Docker spawn path (vllm.py),
+    # which is disabled in production; the real launch (start_single_container.sh)
+    # never passes --speculative-config. Benchmarked +27-36% SLOWER at OCR's 16-wide
+    # block concurrency (GPU already compute-bound) — see
+    # docs/quantization_benchmark_results.md §5b. Do not wire MTP into the launch.
+    VLLM_ENABLE_MTP: bool = True
+    VLLM_MTP_TOKENS: int = 2
+    VLLM_EXTRA_ARGS: Optional[str] = None
+    DOCKER_HF_CACHE_PATH: str = "~/.cache/huggingface"
+
+    # llama.cpp
+    LLAMA_CPP_BINARY: str = "llama-server"
+    LLAMA_CPP_NGL: int = 99  # all layers on GPU (Metal on macOS, CUDA on Linux GPU); harmless no-op on pure-CPU builds
+    LLAMA_CPP_NO_MMPROJ_OFFLOAD: bool = False
+    LLAMA_CPP_EXTRA_ARGS: Optional[str] = None
+
+    # ---- Detection (kept) ---------------------------------------------------
+    DETECTOR_BATCH_SIZE: Optional[int] = None
+    DETECTOR_MODEL_CHECKPOINT: str = "s3://text_detection/2025_05_07"
+    DETECTOR_IMAGE_CHUNK_HEIGHT: int = 1400
+    DETECTOR_TEXT_THRESHOLD: float = 0.6
+    DETECTOR_BLANK_THRESHOLD: float = 0.35
+    DETECTOR_POSTPROCESSING_CPU_WORKERS: int = min(8, os.cpu_count())
+    DETECTOR_MIN_PARALLEL_THRESH: int = 3
+    DETECTOR_BOX_Y_EXPAND_MARGIN: float = 0.05
+
+    # ---- OCR Error (kept) ---------------------------------------------------
+    OCR_ERROR_MODEL_CHECKPOINT: str = "s3://ocr_error_detection/2025_02_18"
+    OCR_ERROR_BATCH_SIZE: Optional[int] = None
+
+    # ---- Debug / draw fonts (label rendering on annotated images) ----------
+    RECOGNITION_RENDER_FONTS: Dict[str, str] = {
+        "all": os.path.join(FONT_DIR, "GoNotoCurrent-Regular.ttf"),
+        "zh": os.path.join(FONT_DIR, "GoNotoCJKCore.ttf"),
+        "ja": os.path.join(FONT_DIR, "GoNotoCJKCore.ttf"),
+        "ko": os.path.join(FONT_DIR, "GoNotoCJKCore.ttf"),
+    }
+    RECOGNITION_FONT_DL_BASE: str = (
+        "https://github.com/satbyy/go-noto-universal/releases/download/v7.0"
+    )
+
+    @computed_field
+    def MODEL_DTYPE(self) -> torch.dtype:
+        if self.TORCH_DEVICE_MODEL == "cpu":
+            return torch.float32
+        return torch.float16
+
+    @computed_field
+    def MODEL_DTYPE_BFLOAT(self) -> torch.dtype:
+        if self.TORCH_DEVICE_MODEL == "cpu":
+            return torch.float32
+        return torch.bfloat16
+
+    @computed_field
+    def INFERENCE_MODE(self) -> Callable:
+        return torch.inference_mode
+
+    class Config:
+        env_file = find_dotenv("local.env")
+        extra = "ignore"
+
+
+settings = Settings()

surya/table_rec/__init__.py

@@ -0,0 +1,217 @@
+"""TableRecPredictor: dual-path table structure recognition.
+
+- predict_simple: TABLE_REC_PROMPT → rows + columns only, cells derived
+  geometrically (row × column intersections).
+- predict_full: BLOCK_PROMPT on the table crop → full <table> HTML with
+  colspan / rowspan / <th>. The HTML lives on TableResult.html for marker to
+  consume directly.
+"""
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+from PIL import Image
+
+from surya.inference import SuryaInferenceManager, get_default_manager
+from surya.inference.parsers import clean_block_html, denorm_bbox, parse_table_rec
+from surya.inference.prompts import (
+    PROMPT_TYPE_BLOCK,
+    PROMPT_TYPE_TABLE_REC,
+    TABLE_REC_JSON_SCHEMA,
+)
+from surya.inference.schema import BatchInputItem
+from surya.inference.util import image_token_budget
+from surya.logging import get_logger
+from surya.settings import settings
+from surya.table_rec.schema import TableCell, TableCol, TableResult, TableRow
+
+logger = get_logger()
+
+
+def _polygon_from_bbox(bbox):
+    x0, y0, x1, y1 = bbox
+    return [[x0, y0], [x1, y0], [x1, y1], [x0, y1]]
+
+
+def _intersect_bbox(a, b):
+    x0 = max(a[0], b[0])
+    y0 = max(a[1], b[1])
+    x1 = min(a[2], b[2])
+    y1 = min(a[3], b[3])
+    if x1 <= x0 or y1 <= y0:
+        return None
+    return (x0, y0, x1, y1)
+
+
+class TableRecPredictor:
+    def __init__(self, manager: Optional[SuryaInferenceManager] = None):
+        self.manager = manager
+        self._disable_tqdm = settings.DISABLE_TQDM
+
+    @property
+    def disable_tqdm(self) -> bool:
+        return self._disable_tqdm
+
+    @disable_tqdm.setter
+    def disable_tqdm(self, value: bool) -> None:
+        self._disable_tqdm = bool(value)
+
+    def to(self, *args, **kwargs):
+        return
+
+    def __call__(
+        self, images: List[Image.Image], mode: str = "simple"
+    ) -> List[TableResult]:
+        if mode == "full":
+            return self.predict_full(images)
+        return self.predict_simple(images)
+
+    def predict_simple(self, images: List[Image.Image]) -> List[TableResult]:
+        if not images:
+            return []
+        manager = self.manager or get_default_manager()
+        guided = TABLE_REC_JSON_SCHEMA if settings.SURYA_GUIDED_TABLE_REC else None
+        batch = [
+            BatchInputItem(
+                image=img,
+                prompt_type=PROMPT_TYPE_TABLE_REC,
+                max_tokens=settings.SURYA_MAX_TOKENS_TABLE_REC,
+                guided_json=guided,
+            )
+            for img in images
+        ]
+        outputs = manager.generate(batch)
+
+        results: List[TableResult] = []
+        for img, out in zip(images, outputs):
+            w, h = img.size
+            page_bbox = [0, 0, float(w), float(h)]
+            if out.error or not out.raw:
+                results.append(
+                    TableResult(
+                        rows=[],
+                        cols=[],
+                        cells=[],
+                        image_bbox=page_bbox,
+                        raw=out.raw,
+                        mode="simple",
+                        error=True,
+                    )
+                )
+                continue
+            try:
+                elements = parse_table_rec(out.raw)
+            except Exception as e:
+                logger.warning(
+                    f"Table rec parse failed: {e}; raw[:200]={out.raw[:200]!r}"
+                )
+                results.append(
+                    TableResult(
+                        rows=[],
+                        cols=[],
+                        cells=[],
+                        image_bbox=page_bbox,
+                        raw=out.raw,
+                        mode="simple",
+                        error=True,
+                    )
+                )
+                continue
+
+            rows: List[TableRow] = []
+            cols: List[TableCol] = []
+            for el in elements:
+                pixel_bbox = denorm_bbox(el.bbox, w, h, scale=settings.BBOX_SCALE)
+                poly = _polygon_from_bbox(pixel_bbox)
+                if el.label == "Row":
+                    rows.append(TableRow(polygon=poly, row_id=len(rows)))
+                else:
+                    cols.append(TableCol(polygon=poly, col_id=len(cols)))
+
+            # Derive cells geometrically (row × column intersections)
+            cells: List[TableCell] = []
+            cell_id = 0
+            for row in rows:
+                for col in cols:
+                    inter = _intersect_bbox(row.bbox, col.bbox)
+                    if inter is None:
+                        continue
+                    cells.append(
+                        TableCell(
+                            polygon=_polygon_from_bbox(inter),
+                            row_id=row.row_id,
+                            col_id=col.col_id,
+                            cell_id=cell_id,
+                        )
+                    )
+                    cell_id += 1
+            results.append(
+                TableResult(
+                    rows=rows,
+                    cols=cols,
+                    cells=cells,
+                    image_bbox=page_bbox,
+                    raw=out.raw,
+                    mode="simple",
+                    error=False,
+                )
+            )
+        return results
+
+    def predict_full(
+        self, images: List[Image.Image], counts: Optional[List[int]] = None
+    ) -> List[TableResult]:
+        """Full-HTML path: BLOCK_PROMPT on table crops. Use when complex
+        structure (spanning cells, headers) matters and ground-truth-style
+        HTML is preferred. `counts` (one per image) shapes max_tokens."""
+        if not images:
+            return []
+        manager = self.manager or get_default_manager()
+        if counts is None:
+            counts = [0] * len(images)
+        batch = []
+        for img, count in zip(images, counts):
+            batch.append(
+                BatchInputItem(
+                    image=img,
+                    prompt_type=PROMPT_TYPE_BLOCK,
+                    max_tokens=image_token_budget(
+                        count,
+                        ceiling=settings.SURYA_MAX_TOKENS_BLOCK_CEILING,
+                        floor=1024,
+                    ),
+                )
+            )
+        outputs = manager.generate(batch)
+        results: List[TableResult] = []
+        for img, out in zip(images, outputs):
+            w, h = img.size
+            page_bbox = [0, 0, float(w), float(h)]
+            if out.error:
+                results.append(
+                    TableResult(
+                        rows=[],
+                        cols=[],
+                        cells=[],
+                        image_bbox=page_bbox,
+                        raw=out.raw,
+                        mode="full",
+                        error=True,
+                    )
+                )
+                continue
+            html = clean_block_html(out.raw)
+            results.append(
+                TableResult(
+                    rows=[],
+                    cols=[],
+                    cells=[],
+                    image_bbox=page_bbox,
+                    raw=out.raw,
+                    html=html,
+                    mode="full",
+                    error=False,
+                )
+            )
+        return results

surya/table_rec/schema.py

@@ -0,0 +1,48 @@
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+from surya.common.polygon import PolygonBox
+
+
+class TableRow(PolygonBox):
+    row_id: int
+
+    @property
+    def label(self) -> str:
+        return f"Row {self.row_id}"
+
+
+class TableCol(PolygonBox):
+    col_id: int
+
+    @property
+    def label(self) -> str:
+        return f"Column {self.col_id}"
+
+
+class TableCell(PolygonBox):
+    """Geometric cell derived from row × column intersection.
+
+    The simple-path TableRecPredictor doesn't return spanning info from the
+    model — colspan/rowspan/header come from the full-path HTML output if
+    needed."""
+
+    row_id: int
+    col_id: int
+    cell_id: int
+
+    @property
+    def label(self) -> str:
+        return f"Cell {self.cell_id}"
+
+
+class TableResult(BaseModel):
+    rows: List[TableRow]
+    cols: List[TableCol]
+    cells: List[TableCell]
+    image_bbox: List[float]
+    raw: Optional[str] = None  # raw model output
+    html: Optional[str] = None  # populated when full-path was used
+    mode: str = "simple"  # "simple" | "full"
+    error: bool = False

surya/timing.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+import contextvars
+import logging
+import os
+import threading
+import time
+from contextlib import contextmanager
+from typing import Any, Iterator
+
+
+_current_timing: contextvars.ContextVar["TimingCollector | None"] = contextvars.ContextVar(
+    "surya_current_timing",
+    default=None,
+)
+
+
+def timing_enabled() -> bool:
+    return os.getenv("SUYA_TIMING_ENABLED", "true").strip().lower() not in {
+        "0",
+        "false",
+        "no",
+        "off",
+    }
+
+
+class TimingCollector:
+    def __init__(self, *, request_id: str = "-", batch_size: int = 0) -> None:
+        self.request_id = request_id
+        self.batch_size = batch_size
+        self._lock = threading.Lock()
+        self._events: list[dict[str, Any]] = []
+
+    def record(self, name: str, duration_ms: float, **metadata: Any) -> None:
+        if not timing_enabled():
+            return
+        event = {
+            "name": name,
+            "duration_ms": round(duration_ms, 2),
+        }
+        if metadata:
+            event["metadata"] = {
+                key: value
+                for key, value in metadata.items()
+                if value is not None
+            }
+        with self._lock:
+            self._events.append(event)
+
+    def summary(self) -> list[dict[str, Any]]:
+        with self._lock:
+            return list(self._events)
+
+
+def get_current_timing() -> TimingCollector | None:
+    return _current_timing.get()
+
+
+def set_current_timing(collector: TimingCollector | None):
+    return _current_timing.set(collector)
+
+
+def reset_current_timing(token) -> None:
+    _current_timing.reset(token)
+
+
+@contextmanager
+def timing_span(name: str, **metadata: Any) -> Iterator[None]:
+    collector = get_current_timing()
+    if collector is None or not timing_enabled():
+        yield
+        return
+    start = time.perf_counter()
+    try:
+        yield
+    finally:
+        collector.record(name, (time.perf_counter() - start) * 1000, **metadata)
+
+
+def log_timing_summary(
+    logger: logging.Logger,
+    collector: TimingCollector,
+    *,
+    message: str = "surya_timing_summary",
+) -> None:
+    if not timing_enabled():
+        return
+    logger.info(
+        "%s request_id=%s batch_size=%s events=%s",
+        message,
+        collector.request_id,
+        collector.batch_size,
+        collector.summary(),
+    )