diff --git a/ai_mouse/__main__.py b/ai_mouse/__main__.py new file mode 100644 index 0000000..a60eb5d --- /dev/null +++ b/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())