benchmark/parser/run_parser/run_parser.py

"""Benchmark ``StageAParser.parse_pdf`` over the batched OmniDocBench PDFs.

Workflow (see build_pdfs.py first):
  images -> batched PDFs (32 pages each) -> THIS script -> per-PDF JSON + timing.

To measure parse time accurately, all three models (layout / OCR / table) are
loaded into VRAM *before* the timed loop, and an optional warmup parse is run
on the first PDF (discarded) so CUDA kernel autotuning does not pollute the
first real measurement.

For each PDF we call ``parse_pdf`` and save the ``ParsedDocument`` as
``<pdf_stem>.json``; split back to per-image results later using the
``mapping.json`` produced by build_pdfs.py.

Example
-------
    # chạy từ benchmark/parser/ (cần GPU)
    python run_parser/run_parser.py \
        --pdfs   data/pdfs \
        --out    parser_results \
        --timing eval_results/parser_timing.json \
        --device cuda
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--pdfs", required=True, type=Path,
                        help="Folder of batched PDFs (from build_pdfs.py).")
    parser.add_argument("--out", required=True, type=Path,
                        help="Output folder for per-PDF ParsedDocument JSON.")
    parser.add_argument("--timing", type=Path, default=None,
                        help="JSON timing report path (default: <out>/_timing.json).")
    parser.add_argument("--pdftranslator", type=Path, default=None,
                        help="PDFTranslator repo root (default: sibling ../PDFTranslator).")
    parser.add_argument("--device", default="auto",
                        help="Torch device for StageAParser (auto|cuda|cpu).")
    # Batch size + ngưỡng detector (điều chỉnh theo VRAM; mặc định hợp cho A100).
    parser.add_argument("--page-batch-size", type=int, default=32)
    parser.add_argument("--layout-batch-size", type=int, default=32)
    parser.add_argument("--detection-batch-size", type=int, default=32)
    parser.add_argument("--ocr-batch-size", type=int, default=512)
    parser.add_argument("--table-batch-size", type=int, default=512)
    parser.add_argument("--blank-threshold", type=float, default=0.5,
                        help="detector_blank_threshold cho Surya OCR.")
    parser.add_argument("--text-threshold", type=float, default=0.6,
                        help="detector_text_threshold cho Surya OCR.")
    parser.add_argument("--limit", type=int, default=None,
                        help="Only process the first N PDFs (quick test).")
    parser.add_argument("--overwrite", action="store_true",
                        help="Re-run PDFs whose JSON already exists (default: skip/resume).")
    parser.add_argument("--no-warmup", dest="warmup", action="store_false",
                        help="Skip the warmup parse of the first PDF.")
    return parser.parse_args()


def resolve_pdftranslator(arg: Path | None) -> Path:
    if arg is not None:
        if (arg / "pdf2zh").is_dir():
            return arg.resolve()
        raise FileNotFoundError(f"pdf2zh not found under --pdftranslator {arg}")
    # Walk up from this file until we find the PDFTranslator repo root (has pdf2zh).
    for parent in Path(__file__).resolve().parents:
        if (parent / "pdf2zh").is_dir():
            return parent
    raise FileNotFoundError(
        "Cannot locate the PDFTranslator repo root (no pdf2zh/ found above "
        f"{__file__}). Pass --pdftranslator explicitly."
    )


def preload_models(parser) -> float:
    """Load all model weights into VRAM. Returns seconds spent."""
    start = time.perf_counter()
    for name in ("layout_model", "ocr_model", "table_model"):
        model = getattr(parser, name, None)
        if model is not None and getattr(model, "model", None) is None:
            print(f"[run_parser] loading {name} ...", flush=True)
            model.load_model()
    return time.perf_counter() - start


def main() -> int:
    args = parse_args()

    pdftranslator = resolve_pdftranslator(args.pdftranslator)
    sys.path.insert(0, str(pdftranslator))

    from pdf2zh.parser.main import StageAParser

    pdfs = sorted(p for p in args.pdfs.glob("*.pdf"))
    if not pdfs:
        raise FileNotFoundError(f"No .pdf files found in {args.pdfs}")
    if args.limit is not None:
        pdfs = pdfs[: args.limit]

    args.out.mkdir(parents=True, exist_ok=True)
    timing_path = args.timing or (args.out / "_timing.json")

    print(f"[run_parser] {len(pdfs)} PDFs from {args.pdfs}", flush=True)
    print(f"[run_parser] device={args.device}  out={args.out}", flush=True)

    # 1) Construct parser (declares models, no weights yet — lazy loading).
    init_start = time.perf_counter()
    parser = StageAParser(
        device=args.device,
        page_batch_size=args.page_batch_size,
        layout_batch_size=args.layout_batch_size,
        detection_batch_size=args.detection_batch_size,
        ocr_batch_size=args.ocr_batch_size,
        table_batch_size=args.table_batch_size,
        detector_blank_threshold=args.blank_threshold,
        detector_text_threshold=args.text_threshold,
    )
    init_seconds = time.perf_counter() - init_start

    # 2) Preload ALL model weights before timing anything.
    model_load_seconds = preload_models(parser)
    print(f"[run_parser] parser init: {init_seconds:.1f}s  "
          f"model load: {model_load_seconds:.1f}s", flush=True)

    # 3) Warmup parse (discarded) so CUDA autotune doesn't skew the first PDF.
    warmup_seconds = None
    if args.warmup and pdfs:
        w_start = time.perf_counter()
        try:
            parser.parse_pdf(str(pdfs[0]))
            warmup_seconds = time.perf_counter() - w_start
            print(f"[run_parser] warmup parse: {warmup_seconds:.2f}s (discarded)",
                  flush=True)
        except Exception as exc:
            print(f"[run_parser] warmup failed (ignored): {exc!r}", flush=True)

    # 4) Timed loop over all PDFs.
    per_pdf: list[dict] = []
    processed = skipped = failed = 0
    total_parse_seconds = 0.0
    total_pages = 0

    for idx, pdf_path in enumerate(pdfs, start=1):
        out_path = args.out / f"{pdf_path.stem}.json"
        if out_path.exists() and not args.overwrite:
            skipped += 1
            continue

        record: dict = {"pdf": pdf_path.name, "index": idx}
        start = time.perf_counter()
        try:
            doc = parser.parse_pdf(str(pdf_path))
            elapsed = time.perf_counter() - start
            doc.save(out_path)

            num_pages = len(doc.pages)
            num_elements = sum(len(p.elements) for p in doc.pages)
            record.update(
                seconds=round(elapsed, 4),
                num_pages=num_pages,
                num_elements=num_elements,
                seconds_per_page=round(elapsed / num_pages, 4) if num_pages else None,
                ok=True,
            )
            total_parse_seconds += elapsed
            total_pages += num_pages
            processed += 1
            print(f"[{idx}/{len(pdfs)}] {pdf_path.name}  {elapsed:.2f}s  "
                  f"pages={num_pages}  elements={num_elements}", flush=True)
        except Exception as exc:
            elapsed = time.perf_counter() - start
            record.update(seconds=round(elapsed, 4), ok=False, error=repr(exc))
            failed += 1
            print(f"[{idx}/{len(pdfs)}] {pdf_path.name}  FAILED: {exc!r}", flush=True)

        per_pdf.append(record)

    report = {
        "device": args.device,
        "pdfs_dir": str(args.pdfs),
        "out_dir": str(args.out),
        "num_pdfs_total": len(pdfs),
        "num_processed": processed,
        "num_skipped_existing": skipped,
        "num_failed": failed,
        "parser_init_seconds": round(init_seconds, 4),
        "model_load_seconds": round(model_load_seconds, 4),
        "warmup_seconds": round(warmup_seconds, 4) if warmup_seconds is not None else None,
        "total_parse_seconds": round(total_parse_seconds, 4),
        "total_pages": total_pages,
        "avg_seconds_per_pdf": round(total_parse_seconds / processed, 4) if processed else None,
        "avg_seconds_per_page": round(total_parse_seconds / total_pages, 4) if total_pages else None,
        "per_pdf": per_pdf,
    }
    timing_path.write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")

    print("\n[run_parser] done.", flush=True)
    print(f"  processed={processed}  skipped={skipped}  failed={failed}", flush=True)
    print(f"  model load : {model_load_seconds:.1f}s", flush=True)
    print(f"  total parse: {total_parse_seconds:.1f}s over {total_pages} pages", flush=True)
    if report["avg_seconds_per_page"] is not None:
        print(f"  avg/page   : {report['avg_seconds_per_page']:.3f}s", flush=True)
    print(f"  timing     : {timing_path}", flush=True)

    return 1 if (processed == 0 and skipped == 0) else 0


if __name__ == "__main__":
    raise SystemExit(main())