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

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