yellowdog/ai_mouse
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

docs: update CLAUDE.md for new src/tools layout

Browse Source
This commit is contained in:
Huang Qi
2026-05-12 01:33:36 +08:00
parent 0413769dd2
commit 1105026831
1 changed files with 77 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

114
CLAUDE.md
View File

@@ -4,63 +4,103 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
## Project ## Project
Local FastAPI tool that trains and serves ML models for generating **human-like mouse trajectories and scroll wheel events**. Frontend is Vue 3 + ECharts loaded from CDN. Package manager is **uv** (Python 3.12–3.13). `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 ## Commands
```bash ```bash
# Run the web app (opens http://127.0.0.1:8765 in browser) # Web UI (collect + train + verify in browser)
uv run python main.py uv run python tools/serve.py
# Tests (httpx + pytest-asyncio for ASGI integration tests) # Tools CLI dispatch
uv run pytest uv run python -m tools train --data data/traces.jsonl --output data/models_v2
uv run pytest tests/test_generator.py uv run python -m tools eval --model-dir data/models_v2 \
uv run pytest tests/test_server.py::TestStatus::test_status_returns_trace_count --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
# Sync dependencies # Re-export ONNX (after retraining)
uv sync 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
``` ```
There is no separate lint/format config; do not invent one.
## Architecture ## Architecture
Two parallel ML subsystems share an identical **collect → train → verify** workflow. Both pipelines persist JSONL traces, train via SSE-streamed progress, and load bundled weights for inference. Two parallel ML subsystems share a **collect → train → export → serve** flow.
### Mouse trajectories (`ai_mouse/`) ### Mouse trajectories (`src/ai_mouse/mouse.py` library; `tools/trainer.py` training)
- **Model**: `TrajectoryFlowModel` — Conditional Flow Matching (OT) with a 4-layer pre-norm Transformer backbone. Trained by sampling `t ~ U[0,1]`, interpolating `x_t = (1-t)·noise + t·data`, and regressing the velocity field `v = x1 - x0` via MSE. - **Model**: `TrajectoryFlowModel` (Conditional Flow Matching with 4-layer
- **Inference** ([generator.py](ai_mouse/generator.py)): 10-step Euler ODE from noise → trajectory in rotated frame → heavy spatial/temporal post-processing (endpoint snapping, forward monotonicity, log_dt clipping, asymmetric speed profile) → pixel decoding → cumulative timestamps → appended click-down/click-up events sampled from a truncated normal. pre-norm Transformer, d_model=128, defined in `tools/models.py`)
- **Rotated coordinate frame** ([coord.py](ai_mouse/coord.py)): All trajectories are normalised so `start → (0, 0)` and `end → (1, 0)`. Lateral is perpendicular. This makes the model **angle- and distance-invariant** — a 1000px diagonal and a 200px horizontal look identical to the network. `encode_trajectory` / `decode_trajectory` are the only bridge between pixel space and model space. - **Inference**: 10-step Euler ODE in Python; each step runs
- **Condition vector** (3 dims): `[dist/2000, log(dist/100), log(total_dur/500)]`. `total_duration` is sampled at inference from a per-distance-bin log-normal stored in `duration_dist.json`. `session.run(...)` on `src/ai_mouse/assets/flow_model.onnx`. Followed by
- **Training artefacts** in `data/models_v2/`: `flow_model.pt`, `click_dist.json` (truncated-normal click duration), `duration_dist.json` (per-bin log-normal), `train_config.json` (architecture params — required for inference to instantiate the model with matching hyperparameters). numpy post-processing in `_postprocess.py` (endpoint snapping, forward
- **6× data augmentation** in [trainer.py](ai_mouse/trainer.py): original, lateral flip, ±20% speed, temporal noise, flip+speed. monotonicity, gaussian smoothing, log_dt → cumulative timestamps,
- **Legacy** `JointCVAE` in `models.py` is kept only for backward compatibility; the active model is `TrajectoryFlowModel`. 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 (`ai_mouse/scroll/`) ### Scroll wheel (`src/ai_mouse/scroll.py`; `tools/scroll/trainer.py`)
- **Model**: `ScrollCVAE` — smaller bidirectional-GRU encoder + GRU decoder VAE. Sequences are 32 wheel events of `(delta_norm, log_Δt)`. - **Model**: `ScrollCVAE` (bidirectional-GRU encoder + GRU decoder VAE,
- **Condition vector** (7 dims): `[dist/5000, log(dist/500), direction, viewport_norm, mode_onehot×3]` where mode is `"target" | "fast" | "precise"` (see `SCROLL_MODES` in [config.py](ai_mouse/config.py)). `tools/scroll/models.py`). Only the **decoder** is exported to ONNX
- **Inference**: VAE prior sample → softmax-normalise deltas to sum to ~1 → scale to target distance → quantise to wheel increments (40px precise, 120px otherwise) → adjust last event so total matches exactly. (`scroll_decoder.onnx`); encoder is training-only.
- **Note**: scroll *collection* state is JS-side (the browser fires real `wheel` events); the Python `ScrollCollector` only generates targets and persists traces. - **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 (`ai_mouse/server/`) ### Server (`tools/server/`) and frontend (`static/`)
- App factory `create_app()` in [server/__init__.py](ai_mouse/server/__init__.py) mounts four routers under `/api`: `routes_collect`, `routes_train`, `routes_verify`, and `routes_scroll` (the scroll router uses prefix `/api/scroll`). App factory `create_app()` mounts four routers under `/api`. Frontend is
- **Session state** ([server/deps.py](ai_mouse/server/deps.py)): A module-level `SessionState` singleton holds the active `Collector` / `ScrollCollector`. Tests that need to override `_DATA_DIR` monkeypatch `ai_mouse.server.deps._DATA_DIR`. vanilla Vue 3 + axios + ECharts via CDN. Note: the `/api/verify` and
- **Training progress** is delivered via Server-Sent Events. The route launches `train()` on a background thread (`asyncio.to_thread`) and pushes `{epoch, total, loss}` dicts through an `asyncio.Queue` until `{done: True}` or `{error: ...}`. `/api/scroll/verify` endpoints always use the **bundled** ONNX weights
- **Static** is mounted from `static/` at the project root (not under the package); `index.html` is served at `/`. (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
### Frontend (`static/`) restart the server.
Vanilla Vue 3 + axios + ECharts pulled from CDN — **no bundler, no node_modules**. Single page with three tab views (`collect.js`, `train.js`, `verify.js`) registered as components in [app.js](static/js/app.js). `api.js` exports an axios instance and a `fetchSSE()` helper that reads `text/event-stream` from `fetch().body.getReader()` and parses `data: ...` frames. UI strings are in Chinese.
## Config ## Config
[ai_mouse/config.py](ai_mouse/config.py) is the **single source of truth** for hyperparameters. `TrainConfig`, `GenerateConfig`, `ScrollTrainConfig`, `ScrollModeConfig`, and `ServerConfig` are dataclasses. When changing model architecture (`d_model`, `nhead`, etc.) keep training and inference consistent — the `train_config.json` saved at training time is what `generate()` uses to reconstruct the model. `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
`tests/conftest.py` provides `model_dir` and `scroll_model_dir` fixtures that write **freshly initialised (untrained)** weights to a temp dir along with all required JSON metadata. Use these whenever a test calls `generate()` or `generate_scroll()`. The trajectories will be garbage but the inference path runs end-to-end. - `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. Server tests use `httpx.ASGITransport(app=create_app())` with
`pytest-asyncio` — no live socket.