scripts/eval_lattice_checkpoint.py

#!/usr/bin/env python
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from dovla_cil.eval.lattice_eval import evaluate_lattice_checkpoint  # noqa: E402


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Evaluate a DoVLA checkpoint on all held-out same-state lattice edges."
    )
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--dataset", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--device", default="cpu", help="Device to use: 'cpu' or 'cuda'")
    parser.add_argument(
        "--training-k",
        type=int,
        default=None,
        help="Training intervention multiplicity when evaluating on a common fixed-K benchmark.",
    )
    parser.add_argument(
        "--all-groups",
        action="store_true",
        help="Evaluate every group, for a separate state-disjoint benchmark dataset.",
    )
    args = parser.parse_args(argv)
    result = evaluate_lattice_checkpoint(
        args.checkpoint,
        args.dataset,
        output_path=args.out,
        device=args.device,
        training_k=args.training_k,
        all_groups=args.all_groups,
    )
    print(json.dumps(result, indent=2))
    return 0


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