107 lines
4.6 KiB
Markdown
107 lines
4.6 KiB
Markdown
# 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
|
|
|
|
```bash
|
|
# 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.
|