surya-ocr/scripts/quant/finalize.py

"""Assemble the quantization benchmark summary from captured per-page outputs.

Unlike run_all's inline scoring (which globs the reference and penalizes any page
a method did not capture), this scores **candidate-driven**: for each method it
scores only the pages that method actually produced, against the matching
reference page. That lets fast methods run the full eval set while slow/eager
methods (e.g. bitsandbytes) run a representative subset, without the subset being
unfairly penalized for missing pages.

Per-method status/reason and t4-deployability are passed in via a small spec so
methods that could not be built or served are recorded as explicit failed rows.

Run inside the bench container:  python3 -m scripts.quant.finalize
"""
from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Dict, List

from scripts.cer_divergence import cer
from scripts.quant.aggregate import build_row, rows_to_csv, rows_to_markdown
from scripts.quant.bbox_iou import match_boxes

WORK = Path("results/quant")
REF = WORK / "reference"
RESULTS = WORK / "results"

# Per-method outcome metadata for this run (A100, vLLM v0.22.1 container).
# status "ok" methods are scored from their captured pages — if a method has no
# captures on disk it falls back to a failed row with `fallback_error`. "failed"
# methods record why outright. bf16 is the reference (CER 0 / IoU 1 by construction).
#
# int8/awq/gptq were built with llm-compressor main (0.12.x): the PyPI 0.11 line
# pins transformers<5 and cannot load qwen3_5, but main supports transformers 5.x.
# int8 is plain GPTQ W8A8 (SmoothQuant default mappings unresolvable for qwen3_5).
SPECS: List[Dict[str, Any]] = [
    {"method": "bf16", "status": "ok", "t4_deployable": True, "is_reference": True},
    {"method": "fp8", "status": "failed", "t4_deployable": False,
     "error": "vLLM FP8 Marlin kernel: size_n=32 not divisible by tile_n_size=64 "
              "(model layer shape incompatible with fp8 dynamic)"},
    {"method": "int8", "status": "ok", "t4_deployable": True, "is_reference": False,
     "fallback_error": "checkpoint built (llm-compressor main, GPTQ W8A8) but vLLM "
                       "failed to serve it — see /tmp/serve_int8.log"},
    {"method": "awq", "status": "ok", "t4_deployable": True, "is_reference": False,
     "fallback_error": "checkpoint built (llm-compressor main, AWQ W4A16) but vLLM "
                       "failed to serve it — see /tmp/serve_awq.log"},
    {"method": "gptq", "status": "ok", "t4_deployable": True, "is_reference": False,
     "fallback_error": "checkpoint built (llm-compressor main, GPTQ W4A16) but vLLM "
                       "failed to serve it — see /tmp/serve_gptq.log"},
    {"method": "bnb8", "status": "ok", "t4_deployable": True, "is_reference": False},
    {"method": "bnb4", "status": "ok", "t4_deployable": True, "is_reference": False},
]


def _load_page(d: Path, stem: str):
    j = json.loads((d / f"{stem}.json").read_text(encoding="utf-8"))
    txt = (d / f"{stem}.txt").read_text(encoding="utf-8") if (d / f"{stem}.txt").exists() else ""
    return j, txt


def _score(method: str, is_reference: bool) -> Dict[str, Any]:
    cand_dir = REF if is_reference else RESULTS / method
    stems = sorted(p.stem for p in cand_dir.glob("*.json"))
    latencies, cers, ious, missed, extra = [], [], [], [], []
    for stem in stems:
        cj, ctxt = _load_page(cand_dir, stem)
        if cj.get("elapsed_seconds") is not None:
            latencies.append(cj["elapsed_seconds"])
        rj, rtxt = _load_page(REF, stem)
        cers.append(cer(rtxt, ctxt))
        m = match_boxes(rj.get("boxes", []), cj.get("boxes", []), iou_threshold=0.5)
        ious.append(m["mean_matched_iou"])
        missed.append(m["missed"])
        extra.append(m["extra"])
    latencies.sort()
    n = len(latencies)
    model_dir = WORK / "models" / method
    size = None
    if model_dir.exists():
        size = round(sum(f.stat().st_size for f in model_dir.rglob("*") if f.is_file()) / (1024 * 1024), 1)
    avg = lambda xs: round(sum(xs) / len(xs), 4) if xs else None
    return {
        "n_pages": len(stems),
        "mean_latency_s": round(sum(latencies) / n, 2) if n else None,
        "p50_latency_s": round(latencies[n // 2], 2) if n else None,
        "p95_latency_s": round(latencies[min(n - 1, int(n * 0.95))], 2) if n else None,
        "mean_cer": avg(cers),
        "max_cer": round(max(cers), 4) if cers else None,
        "mean_bbox_iou": avg(ious),
        "mean_missed_lines": avg(missed),
        "mean_extra_lines": avg(extra),
        "model_size_mb": size,
    }


def main() -> None:
    rows = []
    for spec in SPECS:
        if spec["status"] != "ok":
            rows.append(build_row(method=spec["method"], status="failed",
                                  t4_deployable=spec["t4_deployable"], error=spec["error"]))
            continue
        cand_dir = REF if spec.get("is_reference") else RESULTS / spec["method"]
        if not any(cand_dir.glob("*.json")):
            rows.append(build_row(method=spec["method"], status="failed",
                                  t4_deployable=spec["t4_deployable"],
                                  error=spec.get("fallback_error", "no captures on disk")))
            continue
        s = _score(spec["method"], spec.get("is_reference", False))
        rows.append(build_row(
            method=spec["method"], status="ok", t4_deployable=spec["t4_deployable"],
            model_size_mb=s["model_size_mb"], mean_latency_s=s["mean_latency_s"],
            p50_latency_s=s["p50_latency_s"], p95_latency_s=s["p95_latency_s"],
            mean_cer=s["mean_cer"], max_cer=s["max_cer"], mean_bbox_iou=s["mean_bbox_iou"],
            mean_missed_lines=s["mean_missed_lines"], mean_extra_lines=s["mean_extra_lines"],
        ))
        print(f"{spec['method']}: {s['n_pages']} pages, "
              f"mean_latency={s['mean_latency_s']}s, mean_cer={s['mean_cer']}, iou={s['mean_bbox_iou']}")
    (WORK / "summary.csv").write_text(rows_to_csv(rows), encoding="utf-8")
    (WORK / "summary.json").write_text(json.dumps(rows, indent=2), encoding="utf-8")
    (WORK / "summary.md").write_text(rows_to_markdown(rows), encoding="utf-8")
    print(f"wrote summary for {len(rows)} methods to {WORK}")


if __name__ == "__main__":
    main()