Add Accuracy Benchmark for OCR models (#25364)
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
Chetan Kumar Verma
co-authored by
Ma Mingfei
Ma Mingfei
parent
e2b55bdbab
commit
b3ab56545b
@@ -0,0 +1,727 @@
|
||||
"""
|
||||
Benchmark DeepSeek-OCR-2 (and similar OCR VLMs) on olmOCR-bench via a running sglang server.
|
||||
|
||||
Usage:
|
||||
# 0. Download the dataset (one-time, ~2 GB with PDFs via Git LFS)
|
||||
hf download --repo-type dataset \\
|
||||
allenai/olmOCR-bench --local-dir ./olmOCR-bench
|
||||
|
||||
# 1. Start the sglang server (matches run.sh)
|
||||
python -m sglang.launch_server \\
|
||||
--model-path deepseek-ai/DeepSeek-OCR-2 --host 127.0.0.1 --port 30000
|
||||
|
||||
# 2. Run the full benchmark (all 7 splits, ~7,010 tests)
|
||||
python -m benchmark.ocr.bench_sglang --port 30000 --split all --concurrency 8
|
||||
|
||||
# 3. Quick run on a single split
|
||||
python -m benchmark.ocr.bench_sglang --port 30000 --split arxiv_math --concurrency 16
|
||||
|
||||
# 4. Limit pages for a fast smoke-test
|
||||
python -m benchmark.ocr.bench_sglang --port 30000 --split old_scans --max-samples 10
|
||||
|
||||
# 5. Custom dataset location
|
||||
python -m benchmark.ocr.bench_sglang --bench-dir /data/olmOCR-bench/bench_data
|
||||
|
||||
Dataset:
|
||||
allenai/olmOCR-bench – 7 splits, 1,403 PDFs, 7,010 unit tests
|
||||
Splits: arxiv_math | old_scans_math | table_tests | old_scans |
|
||||
headers_footers | multi_column | long_tiny_text
|
||||
PDFs are stored via Git LFS; hf download (step 0) is required.
|
||||
|
||||
Reference scores (olmOCR-bench):
|
||||
DeepSeek-OCR v1 : 75.7 ± 1.0
|
||||
DeepSeek-OCR-2 : 76.3 (reported on HF model card)
|
||||
olmOCR v0.4.0 : 82.4 ± 1.1
|
||||
PaddleOCR-VL : 80.0 ± 1.0
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
import aiohttp
|
||||
from tqdm.asyncio import tqdm as atqdm
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Paths: allow running from the repo root or from benchmark/ocr/
|
||||
# ---------------------------------------------------------------------------
|
||||
_SCRIPT_DIR = Path(__file__).resolve().parent
|
||||
if str(_SCRIPT_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_SCRIPT_DIR))
|
||||
|
||||
from eval_utils import (
|
||||
aggregate_results,
|
||||
evaluate_olmocr_tests,
|
||||
print_results_table,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Constants
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
OLMOCR_BENCH_SPLITS = [
|
||||
"arxiv_math",
|
||||
"old_scans_math",
|
||||
"table_tests",
|
||||
"old_scans",
|
||||
"headers_footers",
|
||||
"multi_column",
|
||||
"long_tiny_text",
|
||||
]
|
||||
|
||||
# DeepSeek-OCR-2 prompt formats (https://github.com/deepseek-ai/DeepSeek-OCR-2)
|
||||
_PROMPT_MARKDOWN = "<|grounding|>Convert the document to markdown."
|
||||
_PROMPT_FREE_OCR = "Free OCR."
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Argument dataclass
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchArgs:
|
||||
port: int = 30000
|
||||
host: str = "127.0.0.1"
|
||||
model: str = "deepseek-ai/DeepSeek-OCR-2"
|
||||
split: str = "all"
|
||||
concurrency: int = 8
|
||||
output_dir: str = "./ocr_bench_results"
|
||||
max_samples: int = -1
|
||||
prompt_mode: str = "markdown"
|
||||
bench_dir: str = "./olmOCR-bench/bench_data"
|
||||
request_timeout: int = 300
|
||||
save_raw_outputs: bool = False
|
||||
render_dpi: int = 150
|
||||
debug: bool = False
|
||||
debug_accuracy: bool = False
|
||||
|
||||
@staticmethod
|
||||
def add_cli_args(parser: argparse.ArgumentParser) -> None:
|
||||
parser.add_argument(
|
||||
"--port", type=int, default=BenchArgs.port, help="sglang server port"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--host", type=str, default=BenchArgs.host, help="sglang server host"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default=BenchArgs.model,
|
||||
help="Model identifier (must match the running server)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--split",
|
||||
type=str,
|
||||
default=BenchArgs.split,
|
||||
choices=OLMOCR_BENCH_SPLITS + ["all"],
|
||||
help="Dataset split to evaluate. Use 'all' for all splits.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--concurrency",
|
||||
type=int,
|
||||
default=BenchArgs.concurrency,
|
||||
help="Max concurrent requests to the sglang server",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=str,
|
||||
default=BenchArgs.output_dir,
|
||||
help="Directory for result JSON files",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-samples",
|
||||
type=int,
|
||||
default=BenchArgs.max_samples,
|
||||
help="Max samples per split (-1 = all)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt-mode",
|
||||
type=str,
|
||||
default=BenchArgs.prompt_mode,
|
||||
choices=["markdown", "free_ocr"],
|
||||
help=(
|
||||
"Prompt mode for the OCR model: "
|
||||
"'markdown' → '<|grounding|>Convert the document to markdown.'; "
|
||||
"'free_ocr' → 'Free OCR.'"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--bench-dir",
|
||||
type=str,
|
||||
default=BenchArgs.bench_dir,
|
||||
help=(
|
||||
"Local directory containing the olmOCR-bench bench_data/ folder "
|
||||
"(JSONL files + pdfs/ sub-directory). Download first with: "
|
||||
"hf download --repo-type dataset allenai/olmOCR-bench "
|
||||
"--local-dir ./olmOCR-bench"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--request-timeout",
|
||||
type=int,
|
||||
default=BenchArgs.request_timeout,
|
||||
help="Per-request timeout in seconds",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--save-raw-outputs",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Include raw OCR text in result JSON (useful for debugging)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--render-dpi",
|
||||
type=int,
|
||||
default=BenchArgs.render_dpi,
|
||||
help="DPI for rendering PDF pages to images",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--debug",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help=(
|
||||
"Enable debug logging: print per-sample errors immediately, "
|
||||
"show full tracebacks, and abort on the first server connection failure."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--debug-accuracy",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help=(
|
||||
"Print per-sample accuracy details: input PDF path, expected "
|
||||
"expressions/text, OCR output, and pass/fail per test."
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_cli_args(cls, args: argparse.Namespace) -> "BenchArgs":
|
||||
import dataclasses
|
||||
|
||||
attrs = [f.name for f in dataclasses.fields(cls)]
|
||||
return cls(**{attr: getattr(args, attr) for attr in attrs})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Server preflight check
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def preflight_check(api_url: str, model: str, debug: bool) -> None:
|
||||
"""
|
||||
Send a minimal request to the server before the benchmark starts.
|
||||
Raises SystemExit with a clear message if the server is unreachable or
|
||||
returns an unexpected error.
|
||||
"""
|
||||
print(f" Preflight check → {api_url} … ", end="", flush=True)
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": "ping"}],
|
||||
"max_tokens": 1,
|
||||
}
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(
|
||||
api_url,
|
||||
json=payload,
|
||||
timeout=aiohttp.ClientTimeout(total=10),
|
||||
) as resp:
|
||||
if resp.status in (200, 400): # 400 = bad request but server is alive
|
||||
print("OK")
|
||||
return
|
||||
body = await resp.text()
|
||||
print(f"FAILED (HTTP {resp.status})")
|
||||
raise SystemExit(
|
||||
f"Server returned HTTP {resp.status}.\nResponse: {body[:400]}\n"
|
||||
f"Ensure the server is running: python -m sglang.launch_server "
|
||||
f"--model-path {model} --host 127.0.0.1 --port <PORT>"
|
||||
)
|
||||
except (aiohttp.ClientConnectorError, asyncio.TimeoutError) as exc:
|
||||
print("FAILED")
|
||||
msg = (
|
||||
f"Cannot reach sglang server at {api_url}\n"
|
||||
f"Error: {exc}\n"
|
||||
"Check:\n"
|
||||
" 1. Is the server running? (python -m sglang.launch_server ...)\n"
|
||||
" 2. Is --port correct?\n"
|
||||
" 3. Are you running inside the same docker container as the server?"
|
||||
)
|
||||
if debug:
|
||||
msg += f"\n\nFull traceback:\n{traceback.format_exc()}"
|
||||
raise SystemExit(msg)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PDF → base64 PNG
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def pdf_page_to_base64_png(pdf_bytes: bytes, page_num: int = 0, dpi: int = 150) -> str:
|
||||
"""
|
||||
Render a single PDF page to a base64-encoded PNG string.
|
||||
|
||||
Tries PyMuPDF (fitz) first; falls back to pdf2image / poppler.
|
||||
page_num is 0-indexed.
|
||||
"""
|
||||
try:
|
||||
import fitz # PyMuPDF
|
||||
|
||||
doc = fitz.open(stream=pdf_bytes, filetype="pdf")
|
||||
if page_num >= len(doc):
|
||||
page_num = len(doc) - 1
|
||||
page = doc[page_num]
|
||||
mat = fitz.Matrix(dpi / 72.0, dpi / 72.0)
|
||||
pix = page.get_pixmap(matrix=mat)
|
||||
img_bytes = pix.tobytes("png")
|
||||
doc.close()
|
||||
return base64.b64encode(img_bytes).decode("utf-8")
|
||||
|
||||
except ImportError:
|
||||
pass # Try pdf2image below
|
||||
|
||||
try:
|
||||
from pdf2image import convert_from_bytes
|
||||
|
||||
images = convert_from_bytes(
|
||||
pdf_bytes, dpi=dpi, first_page=page_num + 1, last_page=page_num + 1
|
||||
)
|
||||
if not images:
|
||||
raise ValueError(f"pdf2image returned no images for page {page_num}")
|
||||
buf = io.BytesIO()
|
||||
images[0].save(buf, format="PNG")
|
||||
return base64.b64encode(buf.getvalue()).decode("utf-8")
|
||||
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"No PDF rendering library found. "
|
||||
"Install PyMuPDF (pip install pymupdf) or pdf2image (pip install pdf2image)."
|
||||
) from exc
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# OCR request via sglang OpenAI-compatible API
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def run_ocr_request(
|
||||
session: aiohttp.ClientSession,
|
||||
api_url: str,
|
||||
model: str,
|
||||
image_b64: str,
|
||||
text_prompt: str,
|
||||
timeout: int,
|
||||
) -> Tuple[str, float]:
|
||||
"""
|
||||
Send an image + text prompt to the sglang /v1/chat/completions endpoint.
|
||||
|
||||
Returns (ocr_text, latency_seconds).
|
||||
On error returns ("ERROR: ...", -1.0).
|
||||
"""
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": f"data:image/png;base64,{image_b64}"},
|
||||
},
|
||||
{"type": "text", "text": text_prompt},
|
||||
],
|
||||
}
|
||||
],
|
||||
"max_tokens": 4096,
|
||||
"temperature": 0.0,
|
||||
}
|
||||
|
||||
t0 = time.perf_counter()
|
||||
try:
|
||||
async with session.post(
|
||||
api_url,
|
||||
json=payload,
|
||||
timeout=aiohttp.ClientTimeout(total=timeout),
|
||||
) as resp:
|
||||
latency = time.perf_counter() - t0
|
||||
if resp.status != 200:
|
||||
body = await resp.text()
|
||||
return f"ERROR: HTTP {resp.status} – {body[:200]}", -1.0
|
||||
data = await resp.json()
|
||||
text = data["choices"][0]["message"].get("content") or ""
|
||||
return text, latency
|
||||
except asyncio.TimeoutError:
|
||||
return "ERROR: request timed out", -1.0
|
||||
except Exception:
|
||||
return f"ERROR: {traceback.format_exc(limit=3)}", -1.0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Per-sample processing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Local dataset loading (olmOCR-bench flat JSONL)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def load_jsonl_split(bench_dir: Path, split_name: str) -> List[dict]:
|
||||
"""
|
||||
Load test cases from a local olmOCR-bench JSONL file and group them by
|
||||
(pdf, page) so each unique PDF page becomes one benchmark "sample".
|
||||
|
||||
Requires the dataset to have been downloaded first::
|
||||
|
||||
hf download --repo-type dataset \\
|
||||
allenai/olmOCR-bench --local-dir ./olmOCR-bench
|
||||
"""
|
||||
jsonl_path = bench_dir / f"{split_name}.jsonl"
|
||||
if not jsonl_path.exists():
|
||||
raise FileNotFoundError(
|
||||
f"JSONL not found: {jsonl_path}\n"
|
||||
"Download the dataset first:\n"
|
||||
" hf download --repo-type dataset "
|
||||
"--resume-download allenai/olmOCR-bench --local-dir ./olmOCR-bench"
|
||||
)
|
||||
|
||||
test_cases: List[dict] = []
|
||||
with open(jsonl_path, encoding="utf-8") as fh:
|
||||
for line in fh:
|
||||
line = line.strip()
|
||||
if line:
|
||||
test_cases.append(json.loads(line))
|
||||
|
||||
# Group test cases by (pdf_relative_path, page)
|
||||
pdf_dir = bench_dir / "pdfs"
|
||||
groups: Dict[Tuple[str, int], dict] = {}
|
||||
for tc in test_cases:
|
||||
pdf_rel: str = tc["pdf"]
|
||||
page: int = tc.get("page", 1)
|
||||
key = (pdf_rel, page)
|
||||
if key not in groups:
|
||||
groups[key] = {
|
||||
"pdf_path": str(pdf_dir / pdf_rel),
|
||||
"pdf_rel": pdf_rel,
|
||||
"page": page,
|
||||
"tests": [],
|
||||
}
|
||||
groups[key]["tests"].append(tc)
|
||||
|
||||
return list(groups.values())
|
||||
|
||||
|
||||
def _log_sample_error(label: str, error: str, debug: bool) -> None:
|
||||
"""Print a sample error. Always shows a one-liner; full detail only in debug mode."""
|
||||
short = error.splitlines()[0] if error else "unknown error"
|
||||
print(f" [ERROR] {label}: {short}", file=sys.stderr)
|
||||
if debug and len(error.splitlines()) > 1:
|
||||
print(error, file=sys.stderr)
|
||||
|
||||
|
||||
async def process_sample(
|
||||
semaphore: asyncio.Semaphore,
|
||||
session: aiohttp.ClientSession,
|
||||
api_url: str,
|
||||
args: BenchArgs,
|
||||
text_prompt: str,
|
||||
sample: dict,
|
||||
) -> dict:
|
||||
"""Process one benchmark sample (a single PDF page) and return its result dict."""
|
||||
label = f"{sample.get('pdf_rel', '?')} page {sample.get('page', '?')}"
|
||||
async with semaphore:
|
||||
pdf_path = Path(sample["pdf_path"])
|
||||
if not pdf_path.exists():
|
||||
err = f"PDF not found: {pdf_path}"
|
||||
_log_sample_error(label, err, args.debug)
|
||||
return {"error": err, "test_results": [], "passed": 0, "total": 0}
|
||||
|
||||
try:
|
||||
pdf_bytes = pdf_path.read_bytes()
|
||||
except Exception as exc:
|
||||
err = f"PDF read error: {exc}"
|
||||
_log_sample_error(label, err, args.debug)
|
||||
return {"error": err, "test_results": [], "passed": 0, "total": 0}
|
||||
|
||||
page_num = sample["page"] - 1 # convert 1-indexed → 0-indexed
|
||||
|
||||
try:
|
||||
image_b64 = pdf_page_to_base64_png(
|
||||
pdf_bytes, page_num=page_num, dpi=args.render_dpi
|
||||
)
|
||||
except Exception as exc:
|
||||
err = f"PDF render error: {exc}"
|
||||
if args.debug:
|
||||
err += "\n" + traceback.format_exc()
|
||||
_log_sample_error(label, err, args.debug)
|
||||
return {"error": err, "test_results": [], "passed": 0, "total": 0}
|
||||
|
||||
ocr_text, latency = await run_ocr_request(
|
||||
session, api_url, args.model, image_b64, text_prompt, args.request_timeout
|
||||
)
|
||||
|
||||
if ocr_text.startswith("ERROR:"):
|
||||
_log_sample_error(label, ocr_text, args.debug)
|
||||
return {"error": ocr_text, "test_results": [], "passed": 0, "total": 0}
|
||||
|
||||
tests = sample["tests"]
|
||||
test_results = evaluate_olmocr_tests(tests, ocr_text)
|
||||
|
||||
# Accuracy debug: show input expectations and full OCR output
|
||||
if args.debug_accuracy:
|
||||
sep = "-" * 72
|
||||
print(f"\n{sep}", flush=True)
|
||||
print(f"[INPUT ] PDF : {pdf_path}", flush=True)
|
||||
print(
|
||||
f"[INPUT ] Page : {sample['page']} | {len(tests)} test(s)", flush=True
|
||||
)
|
||||
for i, t in enumerate(tests):
|
||||
ttype = t.get("type", "?")
|
||||
if ttype in ("present", "absent", "text_presence", "text_absence"):
|
||||
expected = t.get("text", "")
|
||||
elif ttype in ("math", "math_formula_accuracy"):
|
||||
expected = t.get("math") or t.get("latex", "")
|
||||
elif ttype in ("order", "natural_reading_order"):
|
||||
expected = (
|
||||
f"before={t.get('before', '')!r} after={t.get('after', '')!r}"
|
||||
)
|
||||
else:
|
||||
expected = str(
|
||||
{
|
||||
k: v
|
||||
for k, v in t.items()
|
||||
if k not in ("pdf", "page", "id", "type", "url", "checked")
|
||||
}
|
||||
)
|
||||
print(
|
||||
f"[INPUT ] [{i+1}] type={ttype!r:12s} expected: {expected[:120]}",
|
||||
flush=True,
|
||||
)
|
||||
# Print OCR output (truncate long outputs)
|
||||
ocr_preview = (
|
||||
ocr_text
|
||||
if len(ocr_text) <= 800
|
||||
else ocr_text[:800] + f"\n... [{len(ocr_text)} chars total, truncated]"
|
||||
)
|
||||
print(
|
||||
f"[OUTPUT] OCR text ({len(ocr_text)} chars, latency={latency:.2f}s):",
|
||||
flush=True,
|
||||
)
|
||||
print(ocr_preview, flush=True)
|
||||
|
||||
# Log individual test failures in debug mode
|
||||
if args.debug or args.debug_accuracy:
|
||||
for tr in test_results:
|
||||
status = "PASS" if tr.get("passed") else "FAIL"
|
||||
detail = tr.get("error", "") if not tr.get("passed") else ""
|
||||
suffix = f" | {detail}" if detail else ""
|
||||
print(
|
||||
f"[RESULT] [{status}] type={tr.get('type')!r}{suffix}",
|
||||
flush=True,
|
||||
)
|
||||
if args.debug_accuracy:
|
||||
print(sep, flush=True)
|
||||
|
||||
result: dict = {
|
||||
"pdf": sample["pdf_rel"],
|
||||
"page": sample["page"],
|
||||
"latency": round(latency, 3),
|
||||
"test_results": test_results,
|
||||
"passed": sum(1 for r in test_results if r.get("passed")),
|
||||
"total": len(test_results),
|
||||
# Store expected values so the HTML report can render them
|
||||
"test_inputs": [
|
||||
{
|
||||
"type": t.get("type"),
|
||||
"math": t.get("math") or t.get("latex", ""),
|
||||
"text": t.get("text", ""),
|
||||
"before": t.get("before", ""),
|
||||
"after": t.get("after", ""),
|
||||
# table-specific fields
|
||||
"cell": t.get("cell", ""),
|
||||
"up": t.get("up"),
|
||||
"down": t.get("down"),
|
||||
"left": t.get("left"),
|
||||
"right": t.get("right"),
|
||||
"top_heading": t.get("top_heading"),
|
||||
"left_heading": t.get("left_heading"),
|
||||
}
|
||||
for t in tests
|
||||
],
|
||||
}
|
||||
if args.save_raw_outputs:
|
||||
result["ocr_output"] = ocr_text
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Split-level runner
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def run_split(
|
||||
split_name: str,
|
||||
dataset,
|
||||
args: BenchArgs,
|
||||
api_url: str,
|
||||
text_prompt: str,
|
||||
) -> dict:
|
||||
"""Evaluate one olmOCR-bench split; return aggregated results dict."""
|
||||
samples = list(dataset)
|
||||
if args.max_samples > 0:
|
||||
samples = samples[: args.max_samples]
|
||||
|
||||
semaphore = asyncio.Semaphore(args.concurrency)
|
||||
sample_results: List[dict] = []
|
||||
_first_error: List[str] = [] # capture first error for summary
|
||||
|
||||
connector = aiohttp.TCPConnector(limit=args.concurrency + 4)
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
tasks = [
|
||||
process_sample(semaphore, session, api_url, args, text_prompt, sample)
|
||||
for sample in samples
|
||||
]
|
||||
for future in atqdm(
|
||||
asyncio.as_completed(tasks),
|
||||
total=len(tasks),
|
||||
desc=f" [{split_name}]",
|
||||
leave=True,
|
||||
):
|
||||
result = await future
|
||||
if "error" in result and not _first_error:
|
||||
_first_error.append(result["error"])
|
||||
sample_results.append(result)
|
||||
|
||||
agg = aggregate_results(split_name, sample_results)
|
||||
agg["samples"] = sample_results # include per-sample data for report generation
|
||||
|
||||
# Always surface error summary so silent 0/0 can't happen
|
||||
error_count = agg.get("error_samples", 0)
|
||||
if error_count > 0:
|
||||
first = _first_error[0] if _first_error else "(unknown)"
|
||||
short = first.splitlines()[0]
|
||||
print(
|
||||
f" WARNING: {error_count}/{len(samples)} samples errored and were skipped.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
print(f" First error: {short}", file=sys.stderr)
|
||||
if not args.debug:
|
||||
print(
|
||||
" Re-run with --debug for full per-sample error details.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
return agg
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def parse_args() -> BenchArgs:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Benchmark OCR VLMs on olmOCR-bench via sglang",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||||
)
|
||||
BenchArgs.add_cli_args(parser)
|
||||
ns = parser.parse_args()
|
||||
return BenchArgs.from_cli_args(ns)
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
api_url = f"http://{args.host}:{args.port}/v1/chat/completions"
|
||||
text_prompt = (
|
||||
_PROMPT_MARKDOWN if args.prompt_mode == "markdown" else _PROMPT_FREE_OCR
|
||||
)
|
||||
splits_to_run = OLMOCR_BENCH_SPLITS if args.split == "all" else [args.split]
|
||||
|
||||
print("=" * 60)
|
||||
print(f" OCR Accuracy Benchmark – olmOCR-bench")
|
||||
print("=" * 60)
|
||||
print(f" Model : {args.model}")
|
||||
print(f" Server : {api_url}")
|
||||
print(f" Prompt mode : {args.prompt_mode}")
|
||||
print(f" Splits : {', '.join(splits_to_run)}")
|
||||
print(f" Concurrency : {args.concurrency}")
|
||||
print(f" Bench dir : {args.bench_dir}")
|
||||
print(f" Output dir : {args.output_dir}")
|
||||
if args.debug:
|
||||
print(f" Debug mode : ON (errors)")
|
||||
if args.debug_accuracy:
|
||||
print(f" Debug mode : ON (accuracy — input/output per sample)")
|
||||
print("=" * 60)
|
||||
|
||||
await preflight_check(api_url, args.model, args.debug)
|
||||
|
||||
bench_dir = Path(args.bench_dir)
|
||||
if not bench_dir.exists():
|
||||
raise SystemExit(
|
||||
f"Benchmark directory not found: {bench_dir}\n"
|
||||
"Download the dataset first:\n"
|
||||
" hf download --repo-type dataset "
|
||||
"allenai/olmOCR-bench --local-dir ./olmOCR-bench"
|
||||
)
|
||||
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
all_results: Dict[str, dict] = {}
|
||||
|
||||
for split in splits_to_run:
|
||||
print(f"\nLoading split '{split}' from {bench_dir} …")
|
||||
try:
|
||||
samples = load_jsonl_split(bench_dir, split)
|
||||
except FileNotFoundError as exc:
|
||||
print(f" WARNING: {exc}")
|
||||
continue
|
||||
except Exception as exc:
|
||||
print(f" WARNING: could not load split '{split}': {exc}")
|
||||
continue
|
||||
|
||||
n = (
|
||||
len(samples)
|
||||
if args.max_samples <= 0
|
||||
else min(len(samples), args.max_samples)
|
||||
)
|
||||
print(
|
||||
f" {n} PDF pages to evaluate ({sum(len(s['tests']) for s in samples[:n])} tests) …"
|
||||
)
|
||||
|
||||
split_result = await run_split(split, samples, args, api_url, text_prompt)
|
||||
all_results[split] = split_result
|
||||
|
||||
# Save per-split JSON
|
||||
out_path = os.path.join(args.output_dir, f"{split}.json")
|
||||
with open(out_path, "w", encoding="utf-8") as f:
|
||||
json.dump(split_result, f, indent=2, ensure_ascii=False)
|
||||
print(
|
||||
f" Score: {split_result['overall_score']:.1f}% "
|
||||
f"({split_result['total_passed']}/{split_result['total_tests']} tests passed)"
|
||||
)
|
||||
print(f" Saved → {out_path}")
|
||||
|
||||
if not all_results:
|
||||
print("No results collected – exiting.")
|
||||
return
|
||||
|
||||
# Print final table
|
||||
print_results_table(all_results)
|
||||
|
||||
# Save summary
|
||||
summary_path = os.path.join(args.output_dir, "summary.json")
|
||||
with open(summary_path, "w", encoding="utf-8") as f:
|
||||
json.dump(all_results, f, indent=2, ensure_ascii=False)
|
||||
print(f"\nFull summary saved → {summary_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user