diff --git a/src/ai_mouse/assets/.gitkeep b/src/ai_mouse/assets/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/src/ai_mouse/assets/click_dist.json b/src/ai_mouse/assets/click_dist.json new file mode 100644 index 0000000..f0b15e6 --- /dev/null +++ b/src/ai_mouse/assets/click_dist.json @@ -0,0 +1,6 @@ +{ + "mu": 88.84602649006622, + "sigma": 22.60819023425299, + "low": 20.0, + "high": 500.0 +} \ No newline at end of file diff --git a/src/ai_mouse/assets/duration_dist.json b/src/ai_mouse/assets/duration_dist.json new file mode 100644 index 0000000..466d2c5 --- /dev/null +++ b/src/ai_mouse/assets/duration_dist.json @@ -0,0 +1,47 @@ +{ + "bins": [ + 0, + 50, + 100, + 200, + 400, + 600, + 800, + 1200, + Infinity + ], + "params": [ + { + "mu_log": 5.881647501269015, + "sigma_log": 0.05 + }, + { + "mu_log": 6.141460896343142, + "sigma_log": 0.43685929770118687 + }, + { + "mu_log": 6.409513572943365, + "sigma_log": 0.33676219820945824 + }, + { + "mu_log": 6.590028978462768, + "sigma_log": 0.2675760001657101 + }, + { + "mu_log": 6.70168704031046, + "sigma_log": 0.2809202882427188 + }, + { + "mu_log": 6.758762243710234, + "sigma_log": 0.21524718926818573 + }, + { + "mu_log": 6.214608098422191, + "sigma_log": 0.5 + }, + { + "mu_log": 6.214608098422191, + "sigma_log": 0.5 + } + ] +} \ No newline at end of file diff --git a/src/ai_mouse/assets/flow_model.onnx b/src/ai_mouse/assets/flow_model.onnx new file mode 100644 index 0000000..31652c7 Binary files /dev/null and b/src/ai_mouse/assets/flow_model.onnx differ diff --git a/src/ai_mouse/assets/scroll_config.json b/src/ai_mouse/assets/scroll_config.json new file mode 100644 index 0000000..9275938 --- /dev/null +++ b/src/ai_mouse/assets/scroll_config.json @@ -0,0 +1,13 @@ +{ + "seq_len": 32, + "latent_dim": 16, + "hidden": 64, + "cond_dim": 7, + "epochs": 100, + "batch_size": 32, + "lr": 0.0005, + "beta_max": 0.3, + "beta_anneal_epochs": 15, + "w_delta": 1.0, + "w_logdt": 1.5 +} \ No newline at end of file diff --git a/src/ai_mouse/assets/scroll_decoder.onnx b/src/ai_mouse/assets/scroll_decoder.onnx new file mode 100644 index 0000000..96829a2 Binary files /dev/null and b/src/ai_mouse/assets/scroll_decoder.onnx differ diff --git a/src/ai_mouse/assets/train_config.json b/src/ai_mouse/assets/train_config.json new file mode 100644 index 0000000..79d2ca9 --- /dev/null +++ b/src/ai_mouse/assets/train_config.json @@ -0,0 +1,12 @@ +{ + "seq_len": 64, + "epochs": 50, + "batch_size": 64, + "lr": 5e-06, + "d_model": 128, + "nhead": 4, + "num_layers": 4, + "dim_feedforward": 256, + "dropout": 0.1, + "cond_dim": 3 +} \ No newline at end of file diff --git a/tools/export_onnx.py b/tools/export_onnx.py index c665381..c013f59 100644 --- a/tools/export_onnx.py +++ b/tools/export_onnx.py @@ -1,7 +1,7 @@ """Export trained PyTorch checkpoints to ONNX for the inference SDK. Usage: - uv run python tools/export_onnx.py \ + uv run python -m tools.export_onnx \ --flow-ckpt data/models_v2 \ --scroll-ckpt data/scroll_models \ --output src/ai_mouse/assets/ @@ -258,3 +258,48 @@ def _check_scroll_parity(ckpt_dir: Path, onnx_path: Path) -> None: f"Scroll decoder parity FAILED: max abs diff = {max_diff:.2e} > {_ATOL:.2e}" ) logger.info("Scroll decoder parity OK (atol=%.0e)", _ATOL) + + +def _copy_metadata(flow_dir: Path, scroll_dir: Path, out_dir: Path) -> None: + """Copy JSON metadata files alongside the ONNX models.""" + for name in ("click_dist.json", "duration_dist.json", "train_config.json"): + src = flow_dir / name + if not src.exists(): + raise FileNotFoundError(f"Required metadata missing: {src}") + shutil.copy2(src, out_dir / name) + src = scroll_dir / "scroll_config.json" + if not src.exists(): + raise FileNotFoundError(f"Required metadata missing: {src}") + shutil.copy2(src, out_dir / "scroll_config.json") + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(prog="export_onnx", description=__doc__.splitlines()[0]) + p.add_argument("--flow-ckpt", type=Path, required=True) + p.add_argument("--scroll-ckpt", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + args = p.parse_args(argv) + + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + + args.output.mkdir(parents=True, exist_ok=True) + + flow_onnx = export_flow_model(args.flow_ckpt, args.output) + scroll_onnx = export_scroll_decoder(args.scroll_ckpt, args.output) + + try: + _check_flow_parity(args.flow_ckpt, flow_onnx) + _check_scroll_parity(args.scroll_ckpt, scroll_onnx) + except RuntimeError as exc: + logger.error("Parity check failed: %s", exc) + flow_onnx.unlink(missing_ok=True) + scroll_onnx.unlink(missing_ok=True) + return 1 + + _copy_metadata(args.flow_ckpt, args.scroll_ckpt, args.output) + logger.info("Export complete: %s", args.output) + return 0 + + +if __name__ == "__main__": + sys.exit(main())