CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

Project

ai_mouse is an ONNX-Runtime SDK that generates human-like mouse trajectories and scroll wheel events. Runtime dependencies are numpy + onnxruntime only; training and the FastAPI web UI live under tools/ and are not packaged.

Package manager: uv, Python 3.12-3.13.

Library vs tools — hard boundary

src/ai_mouse/ — wheel content. NEVER add import torch / import fastapi / import scipy / import matplotlib here. CI's library job installs only runtime deps and would break.
tools/ — repo-only dev code (training, server, collector, eval, data adapters, ONNX export). May import library private modules (ai_mouse._coord, ai_mouse._postprocess) freely — they co-evolve.
Bundled assets: src/ai_mouse/assets/{flow_model,scroll_decoder}.onnx plus four JSON metadata files. Re-generated by tools/export_onnx.py after retraining.

Commands

# Web UI (collect + train + verify in browser)
uv run python tools/serve.py

# Tools CLI dispatch
uv run python -m tools train --data data/traces.jsonl --output data/models_v2
uv run python -m tools eval --model-dir data/models_v2 \
    --reference data/pretrain_traces.jsonl --output data/eval_reports/r.md
uv run python -m tools balabit-adapter --input data/balabit_raw \
    --output data/pretrain_traces.jsonl

# Re-export ONNX (after retraining)
uv run python -m tools.export_onnx --flow-ckpt data/models_v2 \
    --scroll-ckpt data/scroll_models --output src/ai_mouse/assets/

# Tests
uv run pytest tests/unit            # library-only (no torch)
uv run pytest tests/tools           # full dev suite (needs [dev] group)
uv run pytest tests/unit/test_mouse.py::test_mouse_model_seed_reproducibility

# Dependency sync
uv sync                              # runtime only
uv sync --group dev                  # dev-everything

Architecture

Two parallel ML subsystems share a collect → train → export → serve flow.

Mouse trajectories (`src/ai_mouse/mouse.py` library; `tools/trainer.py` training)

Model: TrajectoryFlowModel (Conditional Flow Matching with 4-layer pre-norm Transformer, d_model=128, defined in tools/models.py)
Inference: 10-step Euler ODE in Python; each step runs session.run(...) on src/ai_mouse/assets/flow_model.onnx. Followed by numpy post-processing in _postprocess.py (endpoint snapping, forward monotonicity, gaussian smoothing, log_dt → cumulative timestamps, truncated-normal click duration).
Rotated coordinate frame (_coord.py): trajectories normalised so start → (0, 0), end → (1, 0). Makes the model angle/distance invariant.

Scroll wheel (`src/ai_mouse/scroll.py`; `tools/scroll/trainer.py`)

Model: ScrollCVAE (bidirectional-GRU encoder + GRU decoder VAE, tools/scroll/models.py). Only the decoder is exported to ONNX (scroll_decoder.onnx); encoder is training-only.
Inference: sample z ~ N(0, 1) in numpy → one session.run(...) → softmax-normalise deltas → quantise (40 px precise / 120 px otherwise) → scale to target distance → cumulative timestamps.

Server (`tools/server/`) and frontend (`static/`)

App factory create_app() mounts four routers under /api. Frontend is vanilla Vue 3 + axios + ECharts via CDN. Note: the /api/verify and /api/scroll/verify endpoints always use the bundled ONNX weights (via from ai_mouse import generate / generate_scroll). If you retrain and want the Web UI to reflect new weights, re-run tools.export_onnx and restart the server.

Config

tools/config.py holds the training-side dataclasses (TrainConfig, ScrollTrainConfig, etc.). The library does NOT use these — its only "configuration" is what's embedded in src/ai_mouse/assets/train_config.json (architecture params needed to know seq_len etc. at inference time).

Tests

tests/unit/conftest.py — fixtures for library-only tests, no torch.
tests/tools/conftest.py — model_dir and scroll_model_dir fixtures that produce untrained torch weights in a temp dir. Used by training- /server-side tests.
tests/unit/test_golden.py — regression suite that pins library output against tests/unit/data/golden_{mouse,scroll}.npz captured before the ONNX migration. Tolerance is distance-scaled: mouse allows ±max(30 px, 20% of move distance) and ±700 ms; scroll requires exact total deltaY match and ±2 quanta per event.

Server tests use httpx.ASGITransport(app=create_app()) with pytest-asyncio — no live socket.

4.6 KiB Raw Blame History