From 9c4c7f6f0c1ae45756a5cdebf913b024cf448ebb Mon Sep 17 00:00:00 2001 From: Huang Qi Date: Sun, 10 May 2026 13:39:16 +0800 Subject: [PATCH] feat(eval): CLI for generating evaluation reports Co-Authored-By: Claude Sonnet 4.6 --- ai_mouse/eval/__main__.py | 127 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 127 insertions(+) create mode 100644 ai_mouse/eval/__main__.py diff --git a/ai_mouse/eval/__main__.py b/ai_mouse/eval/__main__.py new file mode 100644 index 0000000..27ff9cd --- /dev/null +++ b/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())