scripts/verify_predict_subset.py

#!/usr/bin/env python3
"""Smoke-test ``predict_sheet`` on 1–2 holdout RAG eval sheets.

Usage:
  python scripts/verify_predict_subset.py
  python scripts/verify_predict_subset.py --checkpoint data/experiments/gat_seed42_fold3/best_model.pt --max-sheets 2
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

# Repo root on sys.path
_ROOT = Path(__file__).resolve().parents[1]
if str(_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROOT))


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--checkpoint",
        type=Path,
        default=None,
        help="Path to best_model.pt (default: first data/experiments/gat_seed*/best_model.pt)",
    )
    parser.add_argument("--max-sheets", type=int, default=2, help="Number of holdout sheets to run")
    args = parser.parse_args()

    import config
    import rag_eval_storage
    from predict import load_model_from_path, predict_sheet

    ckpt = args.checkpoint
    if ckpt is None:
        candidates = sorted(config.EXPERIMENTS_DIR.glob("gat_seed*/best_model.pt"))
        if not candidates:
            print("ERROR: No data/experiments/gat_seed*/best_model.pt found.", file=sys.stderr)
            return 1
        ckpt = candidates[0]
        print(f"Using checkpoint: {ckpt}")

    model = load_model_from_path(ckpt)
    if model is None:
        print(f"ERROR: Failed to load {ckpt}", file=sys.stderr)
        return 1

    ds = rag_eval_storage.load_rag_eval_dataset()
    held = list(dict.fromkeys(ds.get("held_out_sheets") or []))
    if not held:
        print("ERROR: No held_out_sheets in rag eval dataset.", file=sys.stderr)
        return 1

    ok = 0
    for sid in held[: max(1, args.max_sheets)]:
        grid = rag_eval_storage.load_rag_eval_grid_for_sheet(sid)
        n_cells = len(grid.get("cells") or []) if grid else 0
        if not grid or n_cells < 3:
            print(f"SKIP {sid}: grid missing or <3 cells (n={n_cells})")
            continue
        preds = predict_sheet(grid, model_override=model)
        if preds is None or len(preds) == 0:
            print(f"FAIL {sid}: predict_sheet returned empty (preds={preds!r})", file=sys.stderr)
            return 1
        print(f"OK   {sid}: {len(preds)} cell predictions")

        from eval_rag import _predict_and_label
        from chunk_builder import extract_chunks

        labeled = _predict_and_label(grid, model=model)
        chunks = extract_chunks(labeled)
        if len(chunks) == 0:
            print(f"WARN {sid}: extract_chunks returned 0 (graph chunking may fallback elsewhere)")
        else:
            print(f"     {len(chunks)} graph chunks from extract_chunks")
        ok += 1

    if ok == 0:
        print("ERROR: No sheet produced predictions.", file=sys.stderr)
        return 1
    print(f"PASS: {ok} sheet(s).")
    return 0


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