feat: unified CLI (python -m ai_mouse {train,eval,balabit-adapter})

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

ai_mouse/__main__.py

@@ -0,0 +1,104 @@
+"""Unified CLI: `python -m ai_mouse {train,eval,balabit-adapter}`
+
+Subcommands dispatch to the underlying modules. This is the recommended
+top-level entry; you can also call `python -m ai_mouse.eval` etc. directly.
+"""
+from __future__ import annotations
+
+import argparse
+import logging
+import sys
+from pathlib import Path
+
+
+def _train_main(args: argparse.Namespace) -> int:
+    from ai_mouse.trainer import train
+
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+    train(
+        data_path=Path(args.data),
+        output_dir=Path(args.output),
+        epochs=args.epochs,
+        batch_size=args.batch_size,
+        lr=args.lr,
+        seq_len=args.seq_len,
+        resume_from=Path(args.resume_from) if args.resume_from else None,
+    )
+    return 0
+
+
+def _eval_main(args: argparse.Namespace) -> int:
+    from ai_mouse.eval.__main__ import main as eval_main
+    # Reconstruct argv for the sub-CLI
+    argv = [
+        "--model-dir", args.model_dir,
+        "--reference", str(args.reference),
+        "--n-samples", str(args.n_samples),
+        "--n-reference", str(args.n_reference),
+        "--output", str(args.output),
+        "--tag", args.tag,
+        "--seed", str(args.seed),
+    ]
+    return eval_main(argv)
+
+
+def _balabit_main(args: argparse.Namespace) -> int:
+    from ai_mouse.data_adapters.balabit import main as bal_main
+    argv = [
+        "--input", str(args.input),
+        "--output", str(args.output),
+        "--window-ms", str(args.window_ms),
+        "--min-dist", str(args.min_dist),
+        "--min-events", str(args.min_events),
+        "--max-span-ms", str(args.max_span_ms),
+        "--max-gap-ms", str(args.max_gap_ms),
+    ]
+    if args.overwrite:
+        argv.append("--overwrite")
+    return bal_main(argv)
+
+
+def main() -> int:
+    p = argparse.ArgumentParser(prog="ai_mouse", description="AI Mouse trajectory toolkit")
+    sub = p.add_subparsers(dest="cmd", required=True)
+
+    # train
+    pt = sub.add_parser("train", help="Train (or fine-tune) the Flow Matching model")
+    pt.add_argument("--data", required=True, help="Path to traces.jsonl")
+    pt.add_argument("--output", required=True, help="Output checkpoint dir")
+    pt.add_argument("--epochs", type=int, default=200)
+    pt.add_argument("--batch-size", type=int, default=64)
+    pt.add_argument("--lr", type=float, default=3e-4)
+    pt.add_argument("--seq-len", type=int, default=64)
+    pt.add_argument("--resume-from", default=None, help="Checkpoint dir to resume from (for fine-tune)")
+    pt.set_defaults(func=_train_main)
+
+    # eval
+    pe = sub.add_parser("eval", help="Generate evaluation report")
+    pe.add_argument("--model-dir", required=True)
+    pe.add_argument("--reference", type=Path, required=True)
+    pe.add_argument("--n-samples", type=int, default=200)
+    pe.add_argument("--n-reference", type=int, default=1000)
+    pe.add_argument("--output", type=Path, required=True)
+    pe.add_argument("--tag", default="eval")
+    pe.add_argument("--seed", type=int, default=0)
+    pe.set_defaults(func=_eval_main)
+
+    # balabit-adapter
+    pb = sub.add_parser("balabit-adapter", help="Convert Balabit dataset to traces.jsonl")
+    pb.add_argument("--input", type=Path, required=True)
+    pb.add_argument("--output", type=Path, default=Path("data/pretrain_traces.jsonl"))
+    pb.add_argument("--window-ms", type=int, default=1200)
+    pb.add_argument("--min-dist", type=int, default=50)
+    pb.add_argument("--min-events", type=int, default=5)
+    pb.add_argument("--max-span-ms", type=int, default=5000)
+    pb.add_argument("--max-gap-ms", type=int, default=200)
+    pb.add_argument("--overwrite", action="store_true")
+    pb.set_defaults(func=_balabit_main)
+
+    args = p.parse_args()
+    return args.func(args)
+
+
+if __name__ == "__main__":
+    sys.exit(main())