Fu DaiandClaude Opus 4.8
1a585693be
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>
112 lines
4.7 KiB
Python
112 lines
4.7 KiB
Python
"""Orchestrate the quantization sweep: baseline -> each method -> aggregate -> plot.
|
|
|
|
Each method is isolated: any exception becomes a status="failed" row so one bad
|
|
method never aborts the sweep.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List
|
|
|
|
from scripts.quant.aggregate import build_row, rows_to_csv, rows_to_markdown
|
|
from scripts.quant.bbox_iou import bbox_iou_over_dirs
|
|
from scripts.quant.build_model import build_model
|
|
from scripts.quant.capture import capture_page
|
|
from scripts.quant.manifest import load_manifest
|
|
from scripts.quant.recipes import METHOD_SPECS, method_names, vllm_serve_args
|
|
from scripts.quant.serve import start_server, stop_server, wait_healthy
|
|
|
|
# Reuse the existing CER metric.
|
|
from scripts.cer_divergence import cer_over_dirs
|
|
|
|
OCR_URL = "http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/"
|
|
|
|
|
|
def _dir_size_mb(path: Path) -> float:
|
|
total = sum(f.stat().st_size for f in path.rglob("*") if f.is_file())
|
|
return round(total / (1024 * 1024), 1)
|
|
|
|
|
|
def _measure_method(method: str, base_model: str, work_dir: Path,
|
|
eval_images: List[Path], reference_dir: Path) -> Dict[str, Any]:
|
|
"""Build -> serve -> capture -> metrics for one method. Raises on any failure.
|
|
|
|
Assumes the API process (api.py) is already running and reads SURYA_INFERENCE_URL
|
|
from the env to point at the per-method vLLM server (port set by the runbook).
|
|
Returns a dict of measured summary fields (no method/status/t4 keys).
|
|
"""
|
|
model_dir = work_dir / "models" / method
|
|
out_capture = work_dir / "results" / method
|
|
port = 8000
|
|
|
|
built = build_model(method, base_model, model_dir, eval_images[: min(8, len(eval_images))])
|
|
serve_args = vllm_serve_args(method, str(built), base_model, port)
|
|
proc = start_server(serve_args, str(work_dir / f"vllm_{method}.log"))
|
|
try:
|
|
if not wait_healthy(port):
|
|
raise RuntimeError(f"vLLM did not become healthy for {method}")
|
|
latencies = []
|
|
for image in eval_images:
|
|
cap = capture_page(OCR_URL, image, out_capture, timeout=900)
|
|
if cap["elapsed_seconds"] is not None:
|
|
latencies.append(cap["elapsed_seconds"])
|
|
cer = cer_over_dirs(reference_dir, out_capture)
|
|
iou = bbox_iou_over_dirs(reference_dir, out_capture)
|
|
finally:
|
|
stop_server(port, proc)
|
|
|
|
latencies.sort()
|
|
n = len(latencies)
|
|
size = _dir_size_mb(built) if METHOD_SPECS[method]["kind"] in ("compressor", "bnb") else None
|
|
return {
|
|
"mean_latency_s": round(sum(latencies) / n, 3) if n else None,
|
|
"p50_latency_s": latencies[n // 2] if n else None,
|
|
"p95_latency_s": latencies[min(n - 1, int(n * 0.95))] if n else None,
|
|
"mean_cer": cer["mean_cer"],
|
|
"max_cer": cer["max_cer"],
|
|
"mean_bbox_iou": iou["mean_bbox_iou"],
|
|
"mean_missed_lines": iou["mean_missed_lines"],
|
|
"mean_extra_lines": iou["mean_extra_lines"],
|
|
"model_size_mb": size,
|
|
}
|
|
|
|
|
|
def run_method(method: str, base_model: str, work_dir: Path,
|
|
eval_images: List[Path], reference_dir: Path) -> Dict[str, Any]:
|
|
t4 = METHOD_SPECS[method]["t4_deployable"]
|
|
try:
|
|
metrics = _measure_method(method, base_model, work_dir, eval_images, reference_dir)
|
|
return build_row(method=method, status="ok", t4_deployable=t4, **metrics)
|
|
except Exception as exc: # feasibility probe: record and continue
|
|
return build_row(method=method, status="failed", t4_deployable=t4, error=str(exc))
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Run the quantization benchmark sweep.")
|
|
parser.add_argument("--base-model", default="datalab-to/surya-ocr-2")
|
|
parser.add_argument("--manifest", type=Path, default=Path("eval_set/manifest.txt"))
|
|
parser.add_argument("--image-root", type=Path, default=Path("."))
|
|
parser.add_argument("--work-dir", type=Path, default=Path("results/quant"))
|
|
parser.add_argument("--methods", nargs="*", default=method_names())
|
|
args = parser.parse_args()
|
|
|
|
eval_images = load_manifest(args.manifest, args.image_root)
|
|
reference_dir = args.work_dir / "reference"
|
|
args.work_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
rows = []
|
|
for method in args.methods:
|
|
print(f"=== {method} ===", flush=True)
|
|
rows.append(run_method(method, args.base_model, args.work_dir, eval_images, reference_dir))
|
|
|
|
(args.work_dir / "summary.csv").write_text(rows_to_csv(rows), encoding="utf-8")
|
|
(args.work_dir / "summary.json").write_text(json.dumps(rows, indent=2), encoding="utf-8")
|
|
(args.work_dir / "summary.md").write_text(rows_to_markdown(rows), encoding="utf-8")
|
|
print(f"wrote summary for {len(rows)} methods to {args.work_dir}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|