code/pbext/__init__.py

"""ParseBench provider for a Surya checkpoint served by vLLM (OpenAI API).

Import this package before the parse-bench CLI (e.g. via sitecustomize) to
register the provider and the `surya_vllm` pipeline. Configuration comes from
the environment so one pipeline serves any checkpoint:

  SURYA_VLLM_URL    e.g. http://127.0.0.1:8000/v1
  SURYA_VLLM_MODEL  the served model name
  SURYA_PROMPT      optional prompt override (default: the short OCR prompt)
  SURYA_FURNITURE   "1" to keep page headers/footers in markdown (default 1)
  SURYA_FIGTEXT     "1" to keep text transcribed inside figure blocks (default 0)
"""

from __future__ import annotations

import base64
import io
import os
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Any

import requests

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import surya_convert as sc  # noqa: E402

from parse_bench.extensions import register_pipeline, register_provider  # noqa: E402
from parse_bench.inference.providers.base import (  # noqa: E402
    Provider,
    ProviderPermanentError,
    ProviderTransientError,
)
from parse_bench.schemas.parse_output import (  # noqa: E402
    LayoutItemIR,
    LayoutSegmentIR,
    ParseLayoutPageIR,
    ParseOutput,
)
from parse_bench.schemas.pipeline import PipelineSpec  # noqa: E402
from parse_bench.schemas.pipeline_io import (  # noqa: E402
    InferenceRequest,
    InferenceResult,
    RawInferenceResult,
)
from parse_bench.schemas.product import ProductType  # noqa: E402

SHORT_PROMPT = (
    "OCR this image to HTML. Each block is a div with data-label and data-bbox "
    "(x0 y0 x1 y1, normalized 0-1000)."
)


def _page_pngs(path: Path, dpi: int) -> list[bytes]:
    if path.suffix.lower() in (".png", ".jpg", ".jpeg", ".jfif"):
        return [path.read_bytes()]
    import pymupdf

    doc = pymupdf.open(str(path))
    out = []
    for pg in doc:
        out.append(pg.get_pixmap(dpi=dpi).tobytes("png"))
    doc.close()
    return out


def build_output(raw_output: dict, example_id: str, pipeline_name: str) -> ParseOutput:
    """ParseOutput (markdown + normalised layout pages) from the raw page HTML."""
    pages = raw_output.get("pages_html") or []
    cfg = raw_output.get("_config", {})
    furniture = cfg.get("furniture", True)
    sc.figure_text = bool(cfg.get("figtext", sc.figure_text))
    md = "\n\n".join(sc.to_markdown(h, keep_furniture=furniture) for h in pages)
    layout_pages = []
    for i, h in enumerate(pages, 1):
        items = []
        for it in sc.to_layout_items(h):
            seg = LayoutSegmentIR(x=it["x"], y=it["y"], w=it["w"], h=it["h"],
                                  confidence=1.0, label=it["label"])
            items.append(LayoutItemIR(type=it["type"], value=it["text"], bbox=seg,
                                      layout_segments=[seg]))
        # ParseBench's is_header / is_footer rules read these per-page fields
        # (LlamaParse's furniture contract); without them every such rule fails.
        head, foot = sc.furniture_md(h)
        if items or head or foot:
            layout_pages.append(ParseLayoutPageIR(page_number=i, width=1000.0, height=1000.0,
                                                  items=items, page_header_markdown=head,
                                                  page_footer_markdown=foot))
    return ParseOutput(task_type="parse", example_id=example_id, pipeline_name=pipeline_name,
                       pages=[], layout_pages=layout_pages, markdown=md)


@register_provider("surya_vllm")
class SuryaVllmProvider(Provider):
    def __init__(self, provider_name: str, base_config: dict[str, Any] | None = None):
        super().__init__(provider_name, base_config)
        self._url = os.environ.get("SURYA_VLLM_URL", "http://127.0.0.1:8000/v1").rstrip("/")
        self._model = os.environ.get("SURYA_VLLM_MODEL", "surya")
        p = os.environ.get("SURYA_PROMPT") or "short"
        self._prompt = {"short": SHORT_PROMPT}.get(p, p)
        self._furniture = os.environ.get("SURYA_FURNITURE", "1") == "1"
        self._figtext = os.environ.get("SURYA_FIGTEXT", "0") == "1"
        self._dpi = int(self.base_config.get("dpi", 192))
        self._max_tokens = int(os.environ.get("SURYA_MAX_TOKENS") or self.base_config.get("max_tokens", 12288))

    def _ocr(self, png: bytes) -> str:
        """Greedy OCR. If the page does not close (a loop that hit the token cap), retry with each
        repetition penalty in SURYA_RETRY_REP_PEN (comma list, default "1.1,1.25"; "" = no retry)
        and keep the first output that closes."""
        out = self._ocr_once(png, 1.0)
        for rep in [float(x) for x in os.environ.get("SURYA_RETRY_REP_PEN", "1.1,1.25").split(",") if x]:
            if out.rstrip().endswith("</div>"):
                break
            retry = self._ocr_once(png, rep)
            if retry.rstrip().endswith("</div>"):
                return retry
        return out

    def _ocr_once(self, png: bytes, repetition_penalty: float) -> str:
        b64 = base64.b64encode(png).decode()
        body = {
            "model": self._model,
            "temperature": 0.0,
            "repetition_penalty": repetition_penalty,
            "max_tokens": self._max_tokens,
            "messages": [{"role": "user", "content": [
                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}},
                {"type": "text", "text": self._prompt},
            ]}],
        }
        for attempt in range(4):
            try:
                r = requests.post(f"{self._url}/chat/completions", json=body, timeout=900)
                if r.status_code in (408, 429, 500, 502, 503, 504):
                    time.sleep(2 * (attempt + 1))
                    continue
                if r.status_code != 200:
                    raise ProviderPermanentError(f"HTTP {r.status_code}: {r.text[:200]}")
                return r.json()["choices"][0]["message"]["content"] or ""
            except requests.RequestException as e:
                if attempt == 3:
                    raise ProviderTransientError(str(e)) from e
                time.sleep(2 * (attempt + 1))
        raise ProviderTransientError("vLLM unavailable")

    def run_inference(self, pipeline: PipelineSpec, request: InferenceRequest) -> RawInferenceResult:
        started = datetime.now()
        path = Path(request.source_file_path)
        if not path.exists():
            raise ProviderPermanentError(f"missing {path}")
        pages = [self._ocr(p) for p in _page_pngs(path, self._dpi)]
        done = datetime.now()
        return RawInferenceResult(
            request=request, pipeline=pipeline, pipeline_name=pipeline.pipeline_name,
            product_type=request.product_type,
            raw_output={"pages_html": pages, "_config": {"model": self._model, "dpi": self._dpi,
                                                          "furniture": self._furniture,
                                                          "figtext": self._figtext,
                                                          "prompt": self._prompt}},
            started_at=started, completed_at=done,
            latency_in_ms=int((done - started).total_seconds() * 1000),
        )

    def normalize(self, raw_result: RawInferenceResult) -> InferenceResult:
        out = build_output(raw_result.raw_output, raw_result.request.example_id,
                           raw_result.pipeline_name)
        return InferenceResult(
            request=raw_result.request, pipeline_name=raw_result.pipeline_name,
            product_type=raw_result.product_type, raw_output=raw_result.raw_output,
            output=out, started_at=raw_result.started_at,
            completed_at=raw_result.completed_at, latency_in_ms=raw_result.latency_in_ms,
        )


register_pipeline(PipelineSpec(pipeline_name="surya_vllm", provider_name="surya_vllm",
                               product_type=ProductType.PARSE, config={"dpi": 192}))


# Without an adapter keyed to this provider, LlamaParse's priority-100 adapter
# (whose matches() accepts any ParseOutput) captures the result and hard-fails
# on our label vocabulary. Reuse the Datalab adapter: it reads normalised
# layout_pages, which is what normalize() emits.
from parse_bench.evaluation.layout_adapters.adapters import DatalabLayoutAdapter  # noqa: E402
from parse_bench.evaluation.layout_adapters.registry import register_layout_adapter  # noqa: E402

register_layout_adapter("surya_vllm", priority=95)(
    type("SuryaVllmLayoutAdapter", (DatalabLayoutAdapter,), {})
)