feat: export ONNX weights and metadata into src/ai_mouse/assets/

src/ai_mouse/assets/click_dist.json

@@ -0,0 +1,6 @@
+{
+  "mu": 88.84602649006622,
+  "sigma": 22.60819023425299,
+  "low": 20.0,
+  "high": 500.0
+}

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
+    }
+  ]
+}

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
+}

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
+}

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())