File size: 8,598 Bytes
6b4a31f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 | """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,), {})
)
|