ai_mouse/CLAUDE.md

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

# 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): 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): 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: 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).
• 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 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): 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. 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 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.