feat(eval): CLI for generating evaluation reports

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

ai_mouse/eval/__main__.py

@@ -0,0 +1,127 @@
+"""CLI: `python -m ai_mouse.eval --model-dir ... --reference ... --output ...`
+
+Loads N synthetic start/end pairs, calls the generator, loads M reference
+traces from a Balabit-format jsonl, and writes a Markdown report.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import math
+import random
+import sys
+from pathlib import Path
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+def _load_reference_jsonl(path: Path, n_samples: int) -> list[dict]:
+    """Load up to n_samples reference traces from a JSONL file.
+
+    Returns list of {"xs","ys","ts"} 1-D ndarrays.
+    """
+    out: list[dict] = []
+    with path.open("r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                rec = json.loads(line)
+            except json.JSONDecodeError:
+                continue
+            moves = [e for e in rec.get("events", []) if e.get("type") == "move"]
+            if len(moves) < 4:
+                continue
+            xs = np.array([e["x"] for e in moves], dtype=float)
+            ys = np.array([e["y"] for e in moves], dtype=float)
+            ts = np.array([e["t"] for e in moves], dtype=float)
+            out.append({"xs": xs, "ys": ys, "ts": ts})
+            if len(out) >= n_samples:
+                break
+    return out
+
+
+def _generate_n_samples(
+    model_dir: str, n_samples: int, seed: int = 0
+) -> list[dict]:
+    """Call the project's generator N times with random start/end pairs."""
+    from ai_mouse.generator import generate
+
+    rng = random.Random(seed)
+    out: list[dict] = []
+    for i in range(n_samples):
+        sx = rng.randint(50, 750)
+        sy = rng.randint(50, 550)
+        angle = rng.uniform(0, 2 * math.pi)
+        dist = rng.randint(100, 600)
+        ex = int(sx + dist * math.cos(angle))
+        ey = int(sy + dist * math.sin(angle))
+        ex = max(0, min(800, ex))
+        ey = max(0, min(600, ey))
+        try:
+            pts = generate(start=(sx, sy), end=(ex, ey), model_dir=model_dir)
+        except Exception as exc:  # noqa: BLE001
+            logger.warning("generate() failed at i=%d: %s", i, exc)
+            continue
+        # Drop click events (last 2)
+        moves = pts[:-2]
+        if len(moves) < 4:
+            continue
+        xs = np.array([p[0] for p in moves], dtype=float)
+        ys = np.array([p[1] for p in moves], dtype=float)
+        ts = np.array([p[2] for p in moves], dtype=float)
+        out.append({"xs": xs, "ys": ys, "ts": ts})
+    return out
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Generate eval report comparing model output to reference traces.")
+    parser.add_argument("--model-dir", required=True, help="Path to trained model dir (with flow_model.pt)")
+    parser.add_argument("--reference", type=Path, required=True, help="JSONL reference traces (Balabit holdout)")
+    parser.add_argument("--n-samples", type=int, default=200, help="Number of generated samples")
+    parser.add_argument("--n-reference", type=int, default=1000, help="Number of reference samples to load")
+    parser.add_argument("--output", type=Path, required=True, help="Output Markdown file")
+    parser.add_argument("--tag", default="eval", help="Tag string used in plot filenames")
+    parser.add_argument("--seed", type=int, default=0)
+    args = parser.parse_args(argv)
+
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+
+    if not Path(args.model_dir).exists():
+        logger.error("Model dir not found: %s", args.model_dir)
+        return 2
+    if not args.reference.exists():
+        logger.error("Reference jsonl not found: %s", args.reference)
+        return 2
+
+    logger.info("Loading reference from %s ...", args.reference)
+    ref_traces = _load_reference_jsonl(args.reference, args.n_reference)
+    logger.info("Loaded %d reference traces", len(ref_traces))
+
+    logger.info("Generating %d samples from %s ...", args.n_samples, args.model_dir)
+    gen_traces = _generate_n_samples(args.model_dir, args.n_samples, seed=args.seed)
+    logger.info("Generated %d valid traces", len(gen_traces))
+
+    if not gen_traces or not ref_traces:
+        logger.error("Empty trace sets — aborting")
+        return 1
+
+    from ai_mouse.eval.report import build_report
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    build_report(
+        generated_traces=gen_traces,
+        reference_traces=ref_traces,
+        output_md=args.output,
+        tag=args.tag,
+        model_dir=args.model_dir,
+    )
+    logger.info("Done. Report at %s", args.output)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())