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>

tests/test_cer_divergence.py

@@ -0,0 +1,18 @@
+from scripts.cer_divergence import cer
+
+
+def test_identical_text_is_zero():
+    assert cer("hello world", "hello world") == 0.0
+
+
+def test_single_substitution():
+    # 1 edit over 5 reference chars
+    assert cer("hello", "hallo") == 0.2
+
+
+def test_empty_reference_with_output_is_one():
+    assert cer("", "abc") == 1.0
+
+
+def test_empty_both_is_zero():
+    assert cer("", "") == 0.0

tests/test_endpoint_routes.py

@@ -0,0 +1,21 @@
+def test_app_exposes_all_legacy_and_openai_routes():
+    from surya.endpoint.app import app
+
+    paths = {getattr(r, "path", None) for r in app.routes}
+    expected = {
+        "/home",
+        "/v1/api/ai/suya_ocr",
+        "/v1/api/ai/suya_ocr/",
+        "/v1/api/ai/suya_ocr_vllm",
+        "/v1/api/ai/suya_ocr_vllm/",
+        "/v1/api/ai/suya_ocr_vllm/health",
+        "/v1/api/ai/suya_text_det",
+        "/v1/api/ai/suya_text_det/",
+        "/v1/api/ai/suya_layout_det",
+        "/v1/api/ai/suya_layout_det/",
+        "/v1/api/ai/suya_table_rec",
+        "/v1/api/ai/suya_table_rec/",
+        "/image2text",
+        "/v1/chat/completions",
+    }
+    assert expected.issubset(paths), f"missing routes: {expected - paths}"

tests/test_openai_endpoint.py

@@ -0,0 +1,96 @@
+import base64
+import io
+from unittest.mock import patch
+
+from fastapi.testclient import TestClient
+from PIL import Image
+
+
+def _png_data_url() -> str:
+    buf = io.BytesIO()
+    Image.new("RGB", (8, 8), "white").save(buf, format="PNG")
+    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
+    return f"data:image/png;base64,{b64}"
+
+
+def _client() -> TestClient:
+    from surya.endpoint.app import app
+
+    return TestClient(app)
+
+
+def _image_message():
+    return [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image_url", "image_url": {"url": _png_data_url()}},
+                {"type": "text", "text": "extract the text"},
+            ],
+        }
+    ]
+
+
+def test_returns_openai_chat_completion_shape():
+    body = {"model": "surya-ocr", "messages": _image_message()}
+    fake = {"text_lines": "hello\nworld", "ocr_text_json": {"blocks": []}, "elapsed_seconds": 1.23}
+    with patch("surya.endpoint.service.ocr_via_batcher", return_value=fake):
+        r = _client().post("/v1/chat/completions", json=body)
+    assert r.status_code == 200
+    d = r.json()
+    assert d["object"] == "chat.completion"
+    assert d["model"] == "surya-ocr"
+    assert d["choices"][0]["message"]["content"] == "hello\nworld"
+    assert d["choices"][0]["finish_reason"] == "stop"
+    assert d["surya"]["ocr_text_json"] == {"blocks": []}
+    assert d["surya"]["elapsed_seconds"] == 1.23
+    assert "usage" in d
+
+
+def test_mode_full_page_skips_text_detection():
+    body = {"model": "m", "mode": "full_page", "messages": _image_message()}
+    fake = {"text_lines": "x", "ocr_text_json": {}, "elapsed_seconds": None}
+    with patch("surya.endpoint.service.ocr_via_batcher", return_value=fake) as m:
+        r = _client().post("/v1/chat/completions", json=body)
+    assert r.status_code == 200
+    assert m.call_args.kwargs["skip_text_detection"] is True
+
+
+def test_mode_table_routes_to_table_image():
+    body = {"model": "m", "mode": "table", "messages": _image_message()}
+    fake = {"text_lines": "tbl", "ocr_text_json": {"x": 1}, "elapsed_seconds": None}
+    with patch("surya.endpoint.service.table_image", return_value=fake) as m:
+        r = _client().post("/v1/chat/completions", json=body)
+    assert r.status_code == 200
+    assert m.called
+    assert r.json()["choices"][0]["message"]["content"] == "tbl"
+
+
+def test_ocr_with_boxes_false_omits_structured_json():
+    body = {"model": "m", "ocr_with_boxes": False, "messages": _image_message()}
+    fake = {"text_lines": "x", "ocr_text_json": {"a": 1}, "elapsed_seconds": None}
+    with patch("surya.endpoint.service.ocr_via_batcher", return_value=fake):
+        r = _client().post("/v1/chat/completions", json=body)
+    assert r.json()["surya"]["ocr_text_json"] is None
+
+
+def test_stream_true_returns_400():
+    body = {"model": "m", "stream": True, "messages": _image_message()}
+    r = _client().post("/v1/chat/completions", json=body)
+    assert r.status_code == 400
+    assert r.json()["error"]["type"] == "invalid_request_error"
+
+
+def test_missing_image_returns_400():
+    body = {"model": "m", "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]}
+    r = _client().post("/v1/chat/completions", json=body)
+    assert r.status_code == 400
+    assert r.json()["error"]["type"] == "invalid_request_error"
+
+
+def test_pipeline_error_returns_500():
+    body = {"model": "m", "messages": _image_message()}
+    with patch("surya.endpoint.service.ocr_via_batcher", side_effect=RuntimeError("boom")):
+        r = _client().post("/v1/chat/completions", json=body)
+    assert r.status_code == 500
+    assert r.json()["error"]["type"] == "internal_error"

tests/test_parse_timing.py

@@ -0,0 +1,30 @@
+from scripts.parse_timing import parse_line, aggregate
+
+SAMPLE = (
+    "2026-06-11 14:35:00 - vllm_batcher - _process_jobs - line:119 - INFO - "
+    "surya_timing_summary request_id=a,b batch_size=2 events=["
+    "{'name': 'openai_chat_completion', 'duration_ms': 100.0, 'metadata': {'token_count': 40}}, "
+    "{'name': 'openai_chat_completion', 'duration_ms': 300.0, 'metadata': {'token_count': 60}}, "
+    "{'name': 'recognition_manager_generate', 'duration_ms': 450.0}]"
+)
+
+
+def test_parse_line_extracts_events_and_batch_size():
+    rec = parse_line(SAMPLE)
+    assert rec["batch_size"] == 2
+    assert len(rec["events"]) == 3
+    assert rec["events"][0]["metadata"]["token_count"] == 40
+
+
+def test_parse_line_returns_none_for_unrelated_line():
+    assert parse_line("2026-06-11 - foo - bar - line:1 - INFO - request_start id=x") is None
+
+
+def test_aggregate_counts_and_sums_by_span_name():
+    rec = parse_line(SAMPLE)
+    agg = aggregate([rec])
+    chat = agg["openai_chat_completion"]
+    assert chat["count"] == 2
+    assert chat["total_ms"] == 400.0
+    assert chat["mean_ms"] == 200.0
+    assert chat["total_tokens"] == 100

tests/test_quant_aggregate.py

@@ -0,0 +1,34 @@
+from scripts.quant.aggregate import SUMMARY_FIELDS, build_row, rows_to_csv, rows_to_markdown
+
+
+def test_build_row_defaults_unset_fields_to_none():
+    row = build_row(method="awq", status="ok", t4_deployable=True, mean_cer=0.01)
+    assert row["method"] == "awq"
+    assert row["status"] == "ok"
+    assert row["t4_deployable"] is True
+    assert row["mean_cer"] == 0.01
+    assert row["throughput_rps"] is None
+    assert set(row.keys()) == set(SUMMARY_FIELDS)
+
+
+def test_build_row_rejects_unknown_field():
+    try:
+        build_row(method="awq", bogus=1)
+        assert False, "expected KeyError"
+    except KeyError as exc:
+        assert "bogus" in str(exc)
+
+
+def test_rows_to_csv_has_header_and_order():
+    rows = [build_row(method="bf16", status="ok")]
+    csv_text = rows_to_csv(rows)
+    assert csv_text.splitlines()[0] == ",".join(SUMMARY_FIELDS)
+    assert "bf16" in csv_text.splitlines()[1]
+
+
+def test_rows_to_markdown_renders_failed_row():
+    rows = [build_row(method="gptq", status="failed", error="OOM at load")]
+    md = rows_to_markdown(rows)
+    assert "| gptq |" in md
+    assert "failed" in md
+    assert "OOM at load" in md

tests/test_quant_bbox_iou.py

@@ -0,0 +1,48 @@
+import json
+
+from scripts.quant.bbox_iou import iou, match_boxes, bbox_iou_over_dirs
+
+
+def test_iou_identical_is_one():
+    assert iou([0, 0, 10, 10], [0, 0, 10, 10]) == 1.0
+
+
+def test_iou_disjoint_is_zero():
+    assert iou([0, 0, 10, 10], [20, 20, 30, 30]) == 0.0
+
+
+def test_iou_half_overlap():
+    # two 10x10 boxes overlapping in a 10x5 region -> 50/150
+    assert abs(iou([0, 0, 10, 10], [0, 5, 10, 15]) - (50 / 150)) < 1e-9
+
+
+def test_match_boxes_counts_missed_and_extra():
+    ref = [[0, 0, 10, 10], [100, 100, 110, 110]]
+    cand = [[0, 0, 10, 10], [200, 200, 210, 210], [300, 300, 310, 310]]
+
+    result = match_boxes(ref, cand, iou_threshold=0.5)
+
+    assert result["matched"] == 1
+    assert result["missed"] == 1  # ref box at 100,100 unmatched
+    assert result["extra"] == 2   # two cand boxes unmatched
+    assert abs(result["mean_matched_iou"] - 1.0) < 1e-9
+
+
+def test_match_boxes_empty():
+    result = match_boxes([], [], iou_threshold=0.5)
+    assert result == {"matched": 0, "missed": 0, "extra": 0, "mean_matched_iou": 0.0}
+
+
+def test_bbox_iou_over_dirs(tmp_path):
+    ref_dir = tmp_path / "ref"
+    cand_dir = tmp_path / "cand"
+    ref_dir.mkdir()
+    cand_dir.mkdir()
+    (ref_dir / "p1.json").write_text(json.dumps({"boxes": [[0, 0, 10, 10]]}))
+    (cand_dir / "p1.json").write_text(json.dumps({"boxes": [[0, 0, 10, 10]]}))
+
+    result = bbox_iou_over_dirs(ref_dir, cand_dir, iou_threshold=0.5)
+
+    assert abs(result["mean_bbox_iou"] - 1.0) < 1e-9
+    assert result["mean_missed_lines"] == 0.0
+    assert result["mean_extra_lines"] == 0.0

tests/test_quant_build_model.py

@@ -0,0 +1,39 @@
+from pathlib import Path
+
+import scripts.quant.build_model as bm
+
+
+def test_baseline_and_online_need_no_build(tmp_path):
+    assert bm.needs_build("bf16") is False
+    assert bm.needs_build("fp8") is False
+    assert bm.needs_build("awq") is True
+    assert bm.needs_build("bnb4") is True
+
+
+def test_build_is_idempotent_when_output_exists(tmp_path, monkeypatch):
+    out = tmp_path / "awq"
+    out.mkdir()
+    (out / "config.json").write_text("{}", encoding="utf-8")
+    called = {"n": 0}
+
+    def fake_compressor(*a, **k):
+        called["n"] += 1
+
+    monkeypatch.setattr(bm, "_build_compressor", fake_compressor)
+    result = bm.build_model("awq", base_model="datalab-to/surya-ocr-2", out_dir=out, calib_images=[])
+    assert result == out
+    assert called["n"] == 0  # skipped because config.json already present
+
+
+def test_build_dispatches_to_compressor(tmp_path, monkeypatch):
+    out = tmp_path / "gptq"
+    seen = {}
+
+    def fake_compressor(method, base_model, out_dir, calib_images):
+        seen["method"] = method
+        out_dir.mkdir(parents=True, exist_ok=True)
+        (out_dir / "config.json").write_text("{}", encoding="utf-8")
+
+    monkeypatch.setattr(bm, "_build_compressor", fake_compressor)
+    bm.build_model("gptq", base_model="b", out_dir=out, calib_images=[])
+    assert seen["method"] == "gptq"

tests/test_quant_capture.py

@@ -0,0 +1,41 @@
+import json
+
+from scripts.quant.capture import extract_capture, write_capture
+
+
+def test_extract_capture_pulls_text_boxes_and_latency():
+    body = {
+        "data": {
+            "text_lines": "line one\nline two",
+            "elapsed_seconds": 5.4,
+            "ocr_text_json": {
+                "blocks": [
+                    {"bbox": [1, 2, 3, 4], "html": "line one"},
+                    {"bbox": [5, 6, 7, 8], "html": "line two"},
+                ]
+            },
+        },
+        "message": "success",
+        "code": 200,
+    }
+
+    result = extract_capture(body)
+
+    assert result["text"] == "line one\nline two"
+    assert result["boxes"] == [[1, 2, 3, 4], [5, 6, 7, 8]]
+    assert result["elapsed_seconds"] == 5.4
+
+
+def test_extract_capture_tolerates_missing_fields():
+    result = extract_capture({"data": {}})
+    assert result == {"text": "", "boxes": [], "elapsed_seconds": None}
+
+
+def test_write_capture_writes_txt_and_json(tmp_path):
+    cap = {"text": "hello", "boxes": [[0, 0, 1, 1]], "elapsed_seconds": 2.0}
+
+    write_capture(cap, tmp_path, "page1")
+
+    assert (tmp_path / "page1.txt").read_text(encoding="utf-8") == "hello"
+    loaded = json.loads((tmp_path / "page1.json").read_text(encoding="utf-8"))
+    assert loaded == {"boxes": [[0, 0, 1, 1]], "elapsed_seconds": 2.0}

tests/test_quant_manifest.py

@@ -0,0 +1,29 @@
+from pathlib import Path
+
+from scripts.quant.manifest import load_manifest
+
+
+def test_load_manifest_skips_blanks_and_comments(tmp_path):
+    repo = tmp_path / "repo"
+    repo.mkdir()
+    (repo / "a.png").write_bytes(b"x")
+    (repo / "b.jpg").write_bytes(b"y")
+    manifest = tmp_path / "manifest.txt"
+    manifest.write_text("# header\n\na.png\nb.jpg\n", encoding="utf-8")
+
+    result = load_manifest(manifest, repo)
+
+    assert result == [repo / "a.png", repo / "b.jpg"]
+
+
+def test_load_manifest_raises_on_missing_image(tmp_path):
+    repo = tmp_path / "repo"
+    repo.mkdir()
+    manifest = tmp_path / "manifest.txt"
+    manifest.write_text("missing.png\n", encoding="utf-8")
+
+    try:
+        load_manifest(manifest, repo)
+        assert False, "expected FileNotFoundError"
+    except FileNotFoundError as exc:
+        assert "missing.png" in str(exc)

tests/test_quant_plot.py

@@ -0,0 +1,28 @@
+from scripts.quant.aggregate import build_row
+from scripts.quant.plot import pareto_points, render_all
+
+
+def test_pareto_points_skips_failed_and_missing():
+    rows = [
+        build_row(method="bf16", status="ok", mean_latency_s=6.0, mean_cer=0.0, t4_deployable=True),
+        build_row(method="awq", status="ok", mean_latency_s=4.0, mean_cer=0.01, t4_deployable=True),
+        build_row(method="gptq", status="failed", t4_deployable=True),
+        build_row(method="fp8", status="ok", mean_latency_s=3.0, mean_cer=0.02, t4_deployable=False),
+    ]
+    points = pareto_points(rows, x_key="mean_latency_s", y_key="mean_cer")
+    methods = {p["method"] for p in points}
+    assert methods == {"bf16", "awq", "fp8"}  # gptq failed -> skipped
+    awq = next(p for p in points if p["method"] == "awq")
+    assert awq["x"] == 4.0 and awq["y"] == 0.01 and awq["t4_deployable"] is True
+
+
+def test_render_all_writes_png_files(tmp_path):
+    rows = [
+        build_row(method="bf16", status="ok", mean_latency_s=6.0, mean_cer=0.0,
+                  mean_bbox_iou=1.0, t4_deployable=True),
+        build_row(method="awq", status="ok", mean_latency_s=4.0, mean_cer=0.01,
+                  mean_bbox_iou=0.98, t4_deployable=True),
+    ]
+    written = render_all(rows, tmp_path)
+    assert all(p.exists() for p in written)
+    assert any(p.name == "pareto_latency_cer.png" for p in written)

tests/test_quant_recipes.py

@@ -0,0 +1,37 @@
+from scripts.quant.recipes import METHOD_SPECS, method_names, vllm_serve_args
+
+
+def test_all_seven_methods_present():
+    assert method_names() == ["bf16", "fp8", "int8", "awq", "gptq", "bnb8", "bnb4"]
+
+
+def test_t4_flags():
+    assert METHOD_SPECS["fp8"]["t4_deployable"] is False
+    assert METHOD_SPECS["awq"]["t4_deployable"] is True
+    assert METHOD_SPECS["bnb8"]["t4_deployable"] is True
+
+
+def test_baseline_serves_base_model():
+    args = vllm_serve_args("bf16", model_path="/m/base", base_model="/m/base", port=8001)
+    assert "--model" in args and "/m/base" in args
+    assert "--quantization" not in args
+    assert args[args.index("--port") + 1] == "8001"
+
+
+def test_fp8_uses_online_quantization_on_base():
+    args = vllm_serve_args("fp8", model_path="/m/base", base_model="/m/base", port=8001)
+    assert args[args.index("--quantization") + 1] == "fp8"
+    assert "/m/base" in args
+
+
+def test_compressor_serves_built_path_without_quant_flag():
+    args = vllm_serve_args("awq", model_path="/m/awq", base_model="/m/base", port=8001)
+    assert "/m/awq" in args
+    assert "--quantization" not in args  # quant config travels with the compressed checkpoint
+
+
+def test_bnb_uses_bitsandbytes_flags():
+    args = vllm_serve_args("bnb4", model_path="/m/bnb4", base_model="/m/base", port=8001)
+    assert args[args.index("--quantization") + 1] == "bitsandbytes"
+    assert args[args.index("--load-format") + 1] == "bitsandbytes"
+    assert "/m/bnb4" in args

tests/test_quant_run_all.py

@@ -0,0 +1,29 @@
+import scripts.quant.run_all as ra
+from scripts.quant.aggregate import SUMMARY_FIELDS
+
+
+def test_run_method_failure_becomes_failed_row(monkeypatch):
+    def boom(*a, **k):
+        raise RuntimeError("OOM at load")
+
+    monkeypatch.setattr(ra, "_measure_method", boom)
+    row = ra.run_method("gptq", base_model="b", work_dir=ra.Path("/tmp/x"),
+                        eval_images=[], reference_dir=ra.Path("/tmp/ref"))
+    assert row["method"] == "gptq"
+    assert row["status"] == "failed"
+    assert "OOM at load" in row["error"]
+    assert set(row.keys()) == set(SUMMARY_FIELDS)
+
+
+def test_run_method_success_passes_through_metrics(monkeypatch):
+    def fake_measure(method, base_model, work_dir, eval_images, reference_dir):
+        return {"mean_cer": 0.01, "mean_bbox_iou": 0.98, "mean_latency_s": 4.2,
+                "model_size_mb": 512.0}
+
+    monkeypatch.setattr(ra, "_measure_method", fake_measure)
+    row = ra.run_method("awq", base_model="b", work_dir=ra.Path("/tmp/x"),
+                        eval_images=[], reference_dir=ra.Path("/tmp/ref"))
+    assert row["status"] == "ok"
+    assert row["mean_cer"] == 0.01
+    assert row["t4_deployable"] is True
+    assert row["mean_latency_s"] == 4.2

tests/test_quant_serve.py

@@ -0,0 +1,16 @@
+from scripts.quant.serve import health_url, parse_listening_pid
+
+
+def test_health_url_from_port():
+    assert health_url(8001) == "http://127.0.0.1:8001/health"
+
+
+def test_parse_listening_pid_extracts_from_ss_line():
+    ss_output = (
+        'LISTEN 0 4096 127.0.0.1:8001 0.0.0.0:* users:(("python",pid=12345,fd=7))\n'
+    )
+    assert parse_listening_pid(ss_output, 8001) == 12345
+
+
+def test_parse_listening_pid_returns_none_when_absent():
+    assert parse_listening_pid("LISTEN 0 4096 127.0.0.1:9999 0.0.0.0:*\n", 8001) is None

tests/test_resolve_max_workers.py

@@ -0,0 +1,13 @@
+from surya.inference.backends.openai_client import resolve_max_workers
+
+
+def test_caps_at_max_inflight():
+    assert resolve_max_workers(batch_len=160, max_inflight=16) == 16
+
+
+def test_uses_batch_len_when_smaller():
+    assert resolve_max_workers(batch_len=4, max_inflight=16) == 4
+
+
+def test_never_below_one():
+    assert resolve_max_workers(batch_len=0, max_inflight=16) == 1