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