67 lines
5.5 KiB
Markdown
67 lines
5.5 KiB
Markdown
# CLAUDE.md
|
||
|
||
This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
|
||
|
||
## 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).
|
||
|
||
## Commands
|
||
|
||
```bash
|
||
# Run the web app (opens http://127.0.0.1:8765 in browser)
|
||
uv run python main.py
|
||
|
||
# Tests (httpx + pytest-asyncio for ASGI integration tests)
|
||
uv run pytest
|
||
uv run pytest tests/test_generator.py
|
||
uv run pytest tests/test_server.py::TestStatus::test_status_returns_trace_count
|
||
|
||
# Sync dependencies
|
||
uv sync
|
||
```
|
||
|
||
There is no separate lint/format config; do not invent one.
|
||
|
||
## 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.
|
||
|
||
### Mouse trajectories (`ai_mouse/`)
|
||
|
||
- **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.
|
||
- **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.
|
||
- **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.
|
||
- **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`.
|
||
- **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).
|
||
- **6× data augmentation** in [trainer.py](ai_mouse/trainer.py): original, lateral flip, ±20% speed, temporal noise, flip+speed.
|
||
- **Legacy** `JointCVAE` in `models.py` is kept only for backward compatibility; the active model is `TrajectoryFlowModel`.
|
||
|
||
### Scroll wheel (`ai_mouse/scroll/`)
|
||
|
||
- **Model**: `ScrollCVAE` — smaller bidirectional-GRU encoder + GRU decoder VAE. Sequences are 32 wheel events of `(delta_norm, log_Δt)`.
|
||
- **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)).
|
||
- **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.
|
||
- **Note**: scroll *collection* state is JS-side (the browser fires real `wheel` events); the Python `ScrollCollector` only generates targets and persists traces.
|
||
|
||
### Server (`ai_mouse/server/`)
|
||
|
||
- 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`).
|
||
- **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`.
|
||
- **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: ...}`.
|
||
- **Static** is mounted from `static/` at the project root (not under the package); `index.html` is served at `/`.
|
||
|
||
### Frontend (`static/`)
|
||
|
||
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
|
||
|
||
[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.
|
||
|
||
## 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.
|
||
|
||
Server tests use `httpx.ASGITransport(app=create_app())` with `pytest-asyncio` — no live socket.
|