128 lines
4.6 KiB
Python
128 lines
4.6 KiB
Python
"""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 tools.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())
|