"""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(""): break retry = self._ocr_once(png, rep) if retry.rstrip().endswith(""): 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,), {}) )