diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..c45d3db --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,37 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + branches: [main] + +jobs: + library: + name: Library tests (no torch) + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, windows-latest] + python: ["3.12", "3.13"] + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v3 + - run: uv venv --python ${{ matrix.python }} + - run: uv pip install -e . pytest + - run: uv run pytest tests/unit -v + + dev: + name: Full dev suite (with torch) + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, windows-latest] + python: ["3.12", "3.13"] + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v3 + - run: uv sync --group dev --python ${{ matrix.python }} + - run: uv run pytest tests/ -v diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..b8956f6 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,39 @@ +# Changelog + +All notable changes to this project will be documented here. Format follows +[Keep a Changelog](https://keepachangelog.com/en/1.1.0/); versioning follows +[Semantic Versioning](https://semver.org/). + +## [0.2.0] - 2026-05-12 + +### Changed (breaking) + +- Inference no longer requires PyTorch. Runtime dependencies are now + `numpy + onnxruntime` only. +- Public API additions: `MouseModel` and `ScrollModel` classes wrapping a + persistent ORT `InferenceSession`. +- Function signatures `generate()` and `generate_scroll()` are now keyword-only + past the positional `start`/`end` (or `start_scroll_y`/`target_scroll_y`). +- New parameters: `click=True` (mouse), `seed=` (both), `viewport_height=` (scroll). +- Removed `config=` parameter; use kwargs directly. +- `model_dir=` renamed to `model_path=`; accepts `str` or `pathlib.Path`. +- `start_scrollY` / `target_scrollY` renamed to `start_scroll_y` / `target_scroll_y`. +- Training, web UI, collector, eval, and data adapter code moved to repo-level + `tools/`; no longer packaged in the wheel. + +### Added + +- ONNX-format pre-trained weights bundled inside the wheel via + `importlib.resources` (~3 MB). +- `tools/export_onnx.py` script with PyTorch/ORT parity check. +- Errors namespace `ai_mouse.errors` with `AiMouseError`, `ModelLoadError`, + `GenerationError`. +- Custom ORT providers parameter for GPU / DirectML. +- Per-process `lru_cache` so `generate()` / `generate_scroll()` reuse the + default model across calls. + +### Removed + +- Legacy `JointCVAE` class. +- `ai_mouse.config.GenerateConfig` top-level export (parameters moved to kwargs). +- Source dependency on `scipy.stats.truncnorm` (replaced by numpy rejection sampling). diff --git a/CLAUDE.md b/CLAUDE.md index ea5654a..3f7d82d 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -4,63 +4,103 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co ## 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 ```bash -# Run the web app (opens http://127.0.0.1:8765 in browser) -uv run python main.py +# Web UI (collect + train + verify in browser) +uv run python tools/serve.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 +# 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 -# Sync dependencies -uv sync +# 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 ``` -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. +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. -- **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`. +- **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 (`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)`. -- **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. +- **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 (`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`). -- **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. +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 -[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/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. diff --git a/README.md b/README.md new file mode 100644 index 0000000..7e6de61 --- /dev/null +++ b/README.md @@ -0,0 +1,120 @@ +# ai_mouse + +Human-like mouse trajectory and scroll wheel event generator. Inference runs on +ONNX Runtime; the only runtime dependencies are `numpy` and `onnxruntime`. + +## Install + +```bash +pip install git+https://github.com//ai_mouse.git +``` + +For GPU inference (optional), replace `onnxruntime` with the GPU variant: + +```bash +pip install onnxruntime-gpu # CUDA / TensorRT +# or +pip install onnxruntime-directml # Windows DirectML +``` + +## Quick start + +### Mouse trajectory + +```python +from ai_mouse import generate + +points = generate(start=(100, 200), end=(900, 400)) +# [(x, y, t_ms), ..., (cx, cy, t_down), (cx, cy, t_up)] +``` + +### Scroll wheel + +```python +from ai_mouse import generate_scroll + +events = generate_scroll(start_scroll_y=0, target_scroll_y=2000) +# [{"deltaY": 120, "deltaMode": 0, "t": 32}, ...] +``` + +### Class API (recommended for repeated calls) + +```python +from ai_mouse import MouseModel + +m = MouseModel() # session created once +for target in target_list: + pts = m.generate((cx, cy), target) +``` + +### Custom providers / GPU + +```python +from ai_mouse import MouseModel + +m = MouseModel(providers=["CUDAExecutionProvider", "CPUExecutionProvider"]) +# or +m = MouseModel(providers=["DmlExecutionProvider"]) +``` + +### Reproducibility + +```python +m.generate(start, end, seed=42) +``` + +## API summary + +| Name | Purpose | +|---|---| +| `generate(start, end, *, n_points=64, speed=None, click=True, seed=None)` | One-shot call; internal lru_cache singleton | +| `MouseModel(model_path=None, providers=None, seed=None)` | Persistent session | +| `generate_scroll(...)` / `ScrollModel(...)` | Same shape for scroll | +| `ai_mouse.errors.{ModelLoadError, GenerationError}` | Exception hierarchy | + +## Thread safety + +`MouseModel.generate` and `ScrollModel.generate` are safe to call concurrently +from multiple threads — ORT `InferenceSession` is itself thread-safe. + +## Development + +The repo contains optional dev-only tooling under `tools/` for training your +own models, running the FastAPI web UI, and evaluating output quality. Install +with the `dev` group: + +```bash +uv sync --group dev +``` + +Common commands: + +```bash +# Web UI (collect + train + verify in browser) +uv run python tools/serve.py + +# Training (after collecting your own data) +uv run python -m tools train --data data/traces.jsonl --output data/models_v2 + +# Convert Balabit corpus to trace format +uv run python -m tools balabit-adapter --input data/balabit_raw \ + --output data/pretrain_traces.jsonl + +# Eval report +uv run python -m tools eval --model-dir data/models_v2 \ + --reference data/pretrain_traces.jsonl --output data/eval_reports/report.md + +# 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: + +```bash +uv run pytest tests/unit # library-only (no torch) +uv run pytest tests/tools # full dev suite +``` + +After retraining you need to re-export and rebuild the wheel for the new +weights to ship; the in-app Verify endpoint always uses bundled weights. diff --git a/ai_mouse/__init__.py b/ai_mouse/__init__.py deleted file mode 100644 index f0e1957..0000000 --- a/ai_mouse/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# sites/ai_mouse/ai_mouse/__init__.py -from ai_mouse.generator import generate -from ai_mouse.scroll.generator import generate_scroll - -__all__ = ["generate", "generate_scroll"] diff --git a/ai_mouse/generator.py b/ai_mouse/generator.py deleted file mode 100644 index 6b05055..0000000 --- a/ai_mouse/generator.py +++ /dev/null @@ -1,353 +0,0 @@ -"""Inference layer: Flow Matching trajectory generation. - -Pipeline: -1. Load model from model_dir (flow_model.pt, click_dist.json, - duration_dist.json, train_config.json). -2. Compute condition vector: [dist/2000, log(dist/100), log(total_dur/500)]. -3. Sample total_duration from duration_dist.json by distance bin (log-normal). -4. 10-step Euler ODE: start from noise, integrate velocity field to get trajectory. -5. Spatial post-processing: - a. Endpoint snapping: force first=(0,0), last=(1,0), lerp last 6 points. - b. Smooth start: dampen lateral near start (first 4 points). - c. Enforce forward monotonicity (prevent x-axis jitter). - d. 5-point gaussian smooth on lateral (preserve endpoints). -6. Temporal post-processing: - a. Clip log_dt to [0, 5] to prevent exponential explosion. - (speed profile and median±1.1 hard clip removed in 2026-05 refactor — - let the model's learned timing distribution come through naturally.) -7. Decode to pixels via decode_trajectory. -8. Resample to n_points if n_points != model seq_len. -9. Convert log_dt → ms timestamps, scale to total_duration, clip [2, 150]. -10. Ensure timestamps monotonically increasing. -11. Append click events sampled from truncated normal. -""" -from __future__ import annotations - -import json -import logging -import math -from pathlib import Path - -import numpy as np -import torch -from scipy.stats import truncnorm - -from ai_mouse.config import GenerateConfig -from ai_mouse.coord import decode_trajectory -from ai_mouse.models import TrajectoryFlowModel -from ai_mouse.utils import resample_arc - -logger = logging.getLogger(__name__) - -_BUNDLED_MODELS_DIR = Path(__file__).parent.parent / "data" / "models_v2" - - -# --------------------------------------------------------------------------- -# Duration sampling helper -# --------------------------------------------------------------------------- - - -def _sample_duration(duration_dist: dict, dist: float) -> float: - """Sample a total movement duration (ms) for the given pixel distance. - - Uses per-distance-bin log-normal parameters from duration_dist. - - Args: - duration_dist: dict with "bins" and "params" keys. - dist: pixel distance between start and end. - - Returns: - Sampled duration in milliseconds. - """ - bins = duration_dist["bins"] - params = duration_dist["params"] - # Find bin for this distance - bin_idx = len(bins) - 1 - for i in range(len(bins) - 1): - if dist < bins[i + 1]: - bin_idx = i - break - # Clamp to valid params index - bin_idx = min(bin_idx, len(params) - 1) - mu_log = params[bin_idx]["mu_log"] - sigma_log = params[bin_idx]["sigma_log"] - return float(np.exp(np.random.normal(mu_log, sigma_log))) - - -# --------------------------------------------------------------------------- -# Smoothing helper -# --------------------------------------------------------------------------- - - -def _gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray: - """5-point gaussian smoothing along a 1-D array, preserving endpoints. - - Args: - x: 1-D input array. - sigma: Gaussian std (px); larger = more smoothing. Default 1.0 gives - weights ≈ [0.054, 0.244, 0.403, 0.244, 0.054]. - - Returns: - Smoothed array of the same shape. x[0] and x[-1] are unchanged. - If len(x) < 5, returns x unchanged (kernel won't fit). - """ - if len(x) < 5: - return x.copy() - kernel = np.exp(-0.5 * (np.arange(-2, 3) / sigma) ** 2) - kernel /= kernel.sum() - # Pad with edge values to avoid boundary artifacts, then slice back - padded = np.pad(x, pad_width=2, mode="edge") - smoothed = np.convolve(padded, kernel, mode="valid") - smoothed[0] = x[0] - smoothed[-1] = x[-1] - return smoothed - - -# --------------------------------------------------------------------------- -# Main generate function -# --------------------------------------------------------------------------- - - -def generate( - start: tuple[int, int], - end: tuple[int, int], - n_points: int = 64, - speed: float | None = None, - model_dir: str | None = None, - config: GenerateConfig | None = None, -) -> list[tuple[int, int, int]]: - """Generate a human-like mouse trajectory from start to end. - - Uses a Flow Matching model with 4-step Euler ODE integration. - - Args: - start: (x, y) starting pixel coordinate. - end: (x, y) target pixel coordinate. - n_points: number of movement points in the path (default 64). - speed: optional speed multiplier; speed=2 halves the duration. - model_dir: directory containing flow_model.pt, click_dist.json, - duration_dist.json, train_config.json. - None → use bundled pre-trained weights. - config: GenerateConfig instance; None → use defaults. - - Returns: - List of (x, y, t_ms) tuples. All values are ints. - Last two entries are the mouse-down and mouse-up click events. - """ - if config is None: - config = GenerateConfig() - - model_dir_path = Path(model_dir) if model_dir else _BUNDLED_MODELS_DIR - - flow_pt = model_dir_path / "flow_model.pt" - click_json = model_dir_path / "click_dist.json" - duration_json = model_dir_path / "duration_dist.json" - config_json = model_dir_path / "train_config.json" - - if not flow_pt.exists(): - if model_dir is not None: - raise FileNotFoundError( - f"Model weights not found in {model_dir_path}. " - "Run training first or omit model_dir to use bundled weights." - ) - raise FileNotFoundError( - f"Bundled model weights missing at {_BUNDLED_MODELS_DIR}. " - "Run training first." - ) - - # Load train config for model architecture params - seq_len = config.seq_len - d_model = 128 - nhead = 4 - num_layers = 4 - dim_feedforward = 256 - cond_dim = 3 - if config_json.exists(): - cfg = json.loads(config_json.read_text()) - seq_len = int(cfg.get("seq_len", seq_len)) - d_model = int(cfg.get("d_model", d_model)) - nhead = int(cfg.get("nhead", nhead)) - num_layers = int(cfg.get("num_layers", num_layers)) - dim_feedforward = int(cfg.get("dim_feedforward", dim_feedforward)) - cond_dim = int(cfg.get("cond_dim", cond_dim)) - - # Load model - model = TrajectoryFlowModel( - seq_len=seq_len, - d_model=d_model, - nhead=nhead, - num_layers=num_layers, - dim_feedforward=dim_feedforward, - cond_dim=cond_dim, - ) - model.load_state_dict( - torch.load(flow_pt, map_location="cpu", weights_only=True) - ) - model.eval() - - # Load auxiliary distributions - click_params: dict = {"mu": 80.0, "sigma": 30.0, "low": 20.0, "high": 300.0} - if click_json.exists(): - click_params = json.loads(click_json.read_text()) - - duration_dist: dict | None = None - if duration_json.exists(): - duration_dist = json.loads(duration_json.read_text()) - - # Compute pixel distance - sx, sy = float(start[0]), float(start[1]) - ex, ey = float(end[0]), float(end[1]) - dist = math.hypot(ex - sx, ey - sy) - dist = max(dist, 1.0) - - # Sample total duration - if duration_dist is not None: - total_duration = _sample_duration(duration_dist, dist) - else: - # Fallback: simple heuristic ~2px/ms - total_duration = dist / 2.0 - if speed is not None and speed > 0: - total_duration /= speed - total_duration = max(total_duration, 10.0) - - # Build condition vector: [dist_norm, log_dist, log_total_dur] - cond_arr = np.array( - [ - dist / 2000.0, - math.log(dist / 100.0), - math.log(total_duration / 500.0), - ], - dtype=np.float32, - ) - cond_t = torch.from_numpy(cond_arr).unsqueeze(0) # (1, 3) - - # ----------------------------------------------------------------------- - # 4-step Euler ODE integration - # ----------------------------------------------------------------------- - n_steps = config.n_steps - dt = 1.0 / n_steps - - with torch.no_grad(): - x = torch.randn(1, seq_len, 3) # start from noise - for step in range(n_steps): - t_val = step * dt - t_tensor = torch.tensor([t_val]) - v = model(x, t_tensor, cond_t) - x = x + v * dt - # x is now the generated trajectory in (forward, lateral, log_dt) space - - decoded = x.squeeze(0).numpy() # (seq_len, 3) - - forward = decoded[:, 0].copy() # (seq_len,) - lateral = decoded[:, 1].copy() # (seq_len,) - log_dt = decoded[:, 2].copy() # (seq_len,) - - # ------------------------------------------------------------------ - # Spatial post-processing - # ------------------------------------------------------------------ - - # Endpoint snapping: lerp last 6 points towards (1.0, 0.0) - n_snap = min(6, seq_len // 4) - for i in range(n_snap): - alpha = ((i + 1) / n_snap) ** 2 # quadratic ease-in - k = seq_len - n_snap + i - forward[k] = forward[k] * (1.0 - alpha) + 1.0 * alpha - lateral[k] = lateral[k] * (1.0 - alpha) + 0.0 * alpha - - # Force first and last points to canonical values - forward[0], lateral[0] = 0.0, 0.0 - forward[-1], lateral[-1] = 1.0, 0.0 - - # Smooth start: dampen lateral near start (first 4 points) - n_start_fix = min(4, seq_len // 4) - for i in range(1, n_start_fix + 1): - blend = i / (n_start_fix + 1) # 0.2, 0.4, 0.6, 0.8 - forward[i] = max(forward[i], forward[i - 1]) # ensure monotonic start - lateral[i] = lateral[i] * blend # dampen lateral near start - - # Enforce forward monotonicity with soft correction (prevent x-jitter) - for i in range(1, seq_len - 1): # skip last point (already snapped to 1.0) - if forward[i] < forward[i - 1]: - forward[i] = forward[i - 1] + 0.001 - - # Clamp forward to [0, 1] and re-force endpoints after monotonicity fix - forward = np.clip(forward, 0.0, 1.0) - forward[0] = 0.0 - forward[-1] = 1.0 - - # Lateral 5-point gaussian smoothing (endpoints preserved) - lateral = _gaussian_smooth(lateral, sigma=1.0) - - # ------------------------------------------------------------------ - # Temporal post-processing (log_dt) - # ------------------------------------------------------------------ - - # Clip log_dt to prevent extreme values after exp() - # Training data log_dt = log(Δt_ms + 1), typical range [0, 4.5] - # (e.g., Δt=1ms → 0.69, Δt=10ms → 2.40, Δt=80ms → 4.39) - log_dt = np.clip(log_dt, 0.0, 5.0) - - # First point has no interval (padding from training) - log_dt[0] = 0.0 - - # Decode spatial coordinates to pixels - normalised = np.stack([forward, lateral], axis=1) # (seq_len, 2) - pixels = decode_trajectory(normalised, start, end) # (seq_len, 2) - - # Resample to n_points if needed - if n_points != seq_len: - pixels = resample_arc(pixels, n_points) - # Also resample log_dt via linear interpolation in uniform arc - log_dt = np.interp( - np.linspace(0, 1, n_points), - np.linspace(0, 1, seq_len), - log_dt, - ) - - xs = pixels[:, 0] - ys = pixels[:, 1] - - # Convert log_dt → dt (ms), scale to total_duration, clip - dt_raw = np.exp(log_dt) - dt_raw = np.clip(dt_raw, 0.0, None) - dt_sum = dt_raw.sum() - if dt_sum > 1e-6: - scale = total_duration / dt_sum - else: - scale = total_duration / max(n_points, 1) - dt_ms = np.clip( - dt_raw * scale, - config.dt_clip_min_ms, - config.dt_clip_max_ms, - ) - - # Cumulative timestamps (start at 0) - t_abs = np.cumsum(dt_ms) - t_abs = np.concatenate([[0.0], t_abs[:-1]]) # shift so first point = 0 - - # Ensure monotonically increasing - for i in range(1, len(t_abs)): - if t_abs[i] <= t_abs[i - 1]: - t_abs[i] = t_abs[i - 1] + 1.0 - - move_points: list[tuple[int, int, int]] = [ - (int(round(xs[i])), int(round(ys[i])), int(round(t_abs[i]))) - for i in range(n_points) - ] - - # Sample click duration from truncated normal - mu = float(click_params["mu"]) - sigma = float(click_params["sigma"]) - low = float(click_params["low"]) - high = float(click_params["high"]) - a, b = (low - mu) / sigma, (high - mu) / sigma - click_duration = int(truncnorm.rvs(a, b, loc=mu, scale=sigma)) - click_duration = max(click_duration, int(low)) - - last_t = move_points[-1][2] - click_x = int(round(xs[-1])) - click_y = int(round(ys[-1])) - return move_points + [ - (click_x, click_y, last_t), - (click_x, click_y, last_t + click_duration), - ] diff --git a/ai_mouse/scroll/__init__.py b/ai_mouse/scroll/__init__.py deleted file mode 100644 index 3db526b..0000000 --- a/ai_mouse/scroll/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Scroll wheel event generation subsystem.""" -from ai_mouse.scroll.generator import generate_scroll -from ai_mouse.scroll.trainer import train_scroll -from ai_mouse.scroll.collector import ScrollCollector - -__all__ = ["generate_scroll", "train_scroll", "ScrollCollector"] diff --git a/ai_mouse/scroll/generator.py b/ai_mouse/scroll/generator.py deleted file mode 100644 index ec23ff6..0000000 --- a/ai_mouse/scroll/generator.py +++ /dev/null @@ -1,148 +0,0 @@ -"""Scroll wheel event sequence generator.""" -from __future__ import annotations - -import json -import logging -import math -from pathlib import Path - -import numpy as np -import torch - -from ai_mouse.scroll.models import ScrollCVAE - -logger = logging.getLogger(__name__) - -_BUNDLED_SCROLL_MODELS = Path(__file__).resolve().parent.parent.parent / "data" / "scroll_models" - - -def _build_condition( - distance: float, - direction: int, - mode: str, - viewport_height: float = 900.0, -) -> np.ndarray: - """Build 7-dim condition vector matching the trainer layout. - - Dims: [dist/5000, log(dist/500), direction, viewport_norm, mode_onehot*3] - """ - mode_onehot = [0.0, 0.0, 0.0] - if mode == "target": - mode_onehot[0] = 1.0 - elif mode == "fast": - mode_onehot[1] = 1.0 - elif mode == "precise": - mode_onehot[2] = 1.0 - - viewport_norm = viewport_height / 1000.0 - - return np.array([ - distance / 5000.0, - math.log(max(distance, 1.0) / 500.0), - float(direction), - viewport_norm, - *mode_onehot, - ], dtype=np.float32) - - -def generate_scroll( - start_scrollY: int, - target_scrollY: int, - mode: str = "target", - model_dir: str | None = None, -) -> list[dict]: - """Generate a realistic scroll event sequence. - - Args: - start_scrollY: Current scroll position (px from top). - target_scrollY: Target scroll position. - mode: "target" | "fast" | "precise" - model_dir: Path to scroll model files. None = bundled. - - Returns: - List of {"deltaY": int, "deltaMode": 0, "t": int}. - Positive deltaY = scroll down, negative = scroll up. - """ - model_dir_path = Path(model_dir) if model_dir else _BUNDLED_SCROLL_MODELS - model_pt = model_dir_path / "scroll_model.pt" - config_json = model_dir_path / "scroll_config.json" - - if not model_pt.exists(): - raise FileNotFoundError(f"Scroll model not found at {model_pt}") - - seq_len = 32 - if config_json.exists(): - cfg = json.loads(config_json.read_text()) - seq_len = cfg.get("seq_len", 32) - - model = ScrollCVAE(seq_len=seq_len) - model.load_state_dict(torch.load(model_pt, map_location="cpu", weights_only=True)) - model.eval() - - distance = abs(target_scrollY - start_scrollY) - direction = 1 if target_scrollY > start_scrollY else -1 - distance = max(distance, 10) - - cond = _build_condition(float(distance), direction, mode) - cond_t = torch.from_numpy(cond).unsqueeze(0) - - with torch.no_grad(): - z = torch.randn(1, model.latent_dim) - decoded = model.decode(z, cond_t).squeeze(0).numpy() - - delta_norm = decoded[:, 0] - log_dt = decoded[:, 1] - - # De-normalise delta: use softmax-like normalisation so they sum to ~1 - delta_weights = np.exp(delta_norm) - delta_weights = delta_weights / delta_weights.sum() - delta_px = delta_weights * distance * direction - - # Quantise to realistic wheel increments - quantum = 40 if mode == "precise" else 120 - - delta_quantised = np.round(delta_px / quantum) * quantum - for i in range(len(delta_quantised)): - if delta_quantised[i] == 0: - delta_quantised[i] = quantum * direction - - # Adjust last event so total matches target distance - current_total = delta_quantised.sum() - diff = (distance * direction) - current_total - delta_quantised[-1] += diff - - # Timestamps from log_dt - if len(log_dt) > 3: - median_log = float(np.median(log_dt)) - log_dt[:2] = np.clip(log_dt[:2], None, median_log + 0.5) - log_dt[-2:] = np.clip(log_dt[-2:], None, median_log + 0.5) - - dt_ms = np.exp(log_dt).clip(5, 80) - - # Scale to realistic total duration - if mode == "fast": - expected_duration = distance * 0.2 + 100 - elif mode == "precise": - expected_duration = distance * 1.5 + 300 - else: - expected_duration = distance * 0.4 + 200 - - dt_ms = dt_ms * (expected_duration / max(dt_ms.sum(), 1.0)) - dt_ms = dt_ms.clip(5, 80) - - t_abs = np.cumsum(dt_ms).astype(int) - t_abs = np.concatenate([[0], t_abs[:-1]]) - - # Ensure monotonic - for i in range(1, len(t_abs)): - if t_abs[i] <= t_abs[i - 1]: - t_abs[i] = t_abs[i - 1] + 5 - - # Build events, removing zero-delta (keep at least 5) - events = [] - for i in range(seq_len): - dy = int(delta_quantised[i]) - if dy != 0 or len(events) < 5: - events.append({"deltaY": dy, "deltaMode": 0, "t": int(t_abs[i])}) - - return events diff --git a/data/scroll_models/scroll_config.json b/data/scroll_models/scroll_config.json new file mode 100644 index 0000000..9275938 --- /dev/null +++ b/data/scroll_models/scroll_config.json @@ -0,0 +1,13 @@ +{ + "seq_len": 32, + "latent_dim": 16, + "hidden": 64, + "cond_dim": 7, + "epochs": 100, + "batch_size": 32, + "lr": 0.0005, + "beta_max": 0.3, + "beta_anneal_epochs": 15, + "w_delta": 1.0, + "w_logdt": 1.5 +} \ No newline at end of file diff --git a/data/scroll_models/scroll_model.pt b/data/scroll_models/scroll_model.pt new file mode 100644 index 0000000..90788c3 Binary files /dev/null and b/data/scroll_models/scroll_model.pt differ diff --git a/examples/quickstart.py b/examples/quickstart.py new file mode 100644 index 0000000..149a644 --- /dev/null +++ b/examples/quickstart.py @@ -0,0 +1,30 @@ +"""Minimal example: generate one trajectory + click event. + +Run: uv run python examples/quickstart.py +""" +from __future__ import annotations + +from ai_mouse import generate + +start = (100, 200) +end = (900, 400) +points = generate(start, end, seed=0) + +print(f"Generated {len(points)} events:") +print(f" first move: {points[0]}") +print(f" middle move: {points[len(points) // 2]}") +print(f" last move: {points[-3]}") +print(f" click-down: {points[-2]}") +print(f" click-up: {points[-1]}") + +# Typical replay loop pattern. t_ms is cumulative from the start of the trace, +# so block your sender thread until time-since-start reaches each event's t_ms. +# +# import time +# t0 = time.monotonic() +# for x, y, t_ms in points: +# target_wallclock = t0 + t_ms / 1000.0 +# while time.monotonic() < target_wallclock: +# pass +# # replace this with pyautogui / pynput / win32 mouse_event: +# # send_mouse_move(x, y) diff --git a/pyproject.toml b/pyproject.toml index dd3197f..b97161a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,15 +1,37 @@ [project] name = "ai-mouse" -version = "0.1.0" +version = "0.2.0" +description = "Human-like mouse trajectory and scroll wheel event generator (ONNX Runtime SDK)." requires-python = ">=3.12,<3.14" dependencies = [ - "torch>=2.2.0", "numpy>=1.26.0", + "onnxruntime>=1.17.0", +] + +[project.urls] +Repository = "https://github.com//ai_mouse" + +[dependency-groups] +dev = [ + "torch>=2.2.0", "fastapi>=0.111.0", "uvicorn>=0.29.0", "scipy>=1.10.0", "matplotlib>=3.8.0", + "pytest>=8.0.0", + "pytest-asyncio>=0.23.0", + "httpx>=0.27.0", + "onnx>=1.15.0", + "onnxscript>=0.1", ] -[dependency-groups] -dev = ["pytest>=8.0.0", "pytest-asyncio>=0.23.0", "httpx>=0.27.0"] +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/ai_mouse"] + +[tool.pytest.ini_options] +asyncio_mode = "auto" +testpaths = ["tests"] diff --git a/src/ai_mouse/__init__.py b/src/ai_mouse/__init__.py new file mode 100644 index 0000000..7f2a4c5 --- /dev/null +++ b/src/ai_mouse/__init__.py @@ -0,0 +1,78 @@ +"""ai_mouse — ONNX Runtime SDK for human-like mouse trajectories and scroll events. + +Public API: + + from ai_mouse import generate, generate_scroll, MouseModel, ScrollModel + +See the project README for usage examples. +""" +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path +from typing import Literal + +from ai_mouse import errors +from ai_mouse._model_cache import get_mouse_model, get_scroll_model +from ai_mouse.mouse import MouseModel +from ai_mouse.scroll import ScrollModel + +__version__ = "0.2.0" + +__all__ = [ + "MouseModel", + "ScrollModel", + "errors", + "generate", + "generate_scroll", + "__version__", +] + + +def generate( + start: tuple[int, int], + end: tuple[int, int], + *, + n_points: int = 64, + speed: float | None = None, + click: bool = True, + seed: int | None = None, + model_path: str | Path | None = None, + providers: Sequence[str] | None = None, +) -> list[tuple[int, int, int]]: + """Generate a human-like mouse trajectory. + + See :meth:`MouseModel.generate` for argument semantics. + The underlying :class:`MouseModel` is cached process-wide; repeat + calls with the same ``(model_path, providers)`` reuse the session. + """ + model = get_mouse_model(model_path, providers) + return model.generate( + start=start, + end=end, + n_points=n_points, + speed=speed, + click=click, + seed=seed, + ) + + +def generate_scroll( + start_scroll_y: int, + target_scroll_y: int, + *, + mode: Literal["target", "fast", "precise"] = "target", + viewport_height: int = 800, + seed: int | None = None, + model_path: str | Path | None = None, + providers: Sequence[str] | None = None, +) -> list[dict]: + """Generate a sequence of mouse-wheel events. See :class:`ScrollModel.generate`.""" + model = get_scroll_model(model_path, providers) + return model.generate( + start_scroll_y=start_scroll_y, + target_scroll_y=target_scroll_y, + mode=mode, + viewport_height=viewport_height, + seed=seed, + ) diff --git a/src/ai_mouse/_assets.py b/src/ai_mouse/_assets.py new file mode 100644 index 0000000..cf5c6f5 --- /dev/null +++ b/src/ai_mouse/_assets.py @@ -0,0 +1,52 @@ +"""Asset path resolution for bundled ONNX weights and JSON metadata. + +Uses :mod:`importlib.resources` to locate files inside the installed +package, falling back to a user-supplied directory if provided. +""" +from __future__ import annotations + +from importlib.resources import as_file, files +from pathlib import Path + +from ai_mouse.errors import ModelLoadError + +_PACKAGE_ASSETS = "ai_mouse.assets" + + +def bundled_path(name: str) -> Path: + """Return a filesystem path to a bundled asset. + + Args: + name: filename inside the assets/ directory. + + Returns: + A concrete :class:`pathlib.Path`. For non-zip wheels (the common + case) the path points directly into the installed package; for + zipapp installations :func:`importlib.resources.as_file` + materialises a temp file. + """ + ref = files(_PACKAGE_ASSETS) / name + with as_file(ref) as p: + return Path(p) + + +def resolve(model_path: Path | None, filename: str) -> Path: + """Locate an asset given an optional user-supplied directory. + + Args: + model_path: user-supplied directory, or None to use bundled assets. + filename: file to locate inside the directory. + + Returns: + Absolute path to the asset. + + Raises: + ModelLoadError: if the file does not exist. + """ + if model_path is None: + p = bundled_path(filename) + else: + p = Path(model_path) / filename + if not p.exists(): + raise ModelLoadError(f"Required asset missing: {p}") + return p diff --git a/ai_mouse/coord.py b/src/ai_mouse/_coord.py similarity index 100% rename from ai_mouse/coord.py rename to src/ai_mouse/_coord.py diff --git a/src/ai_mouse/_model_cache.py b/src/ai_mouse/_model_cache.py new file mode 100644 index 0000000..dd88e7e --- /dev/null +++ b/src/ai_mouse/_model_cache.py @@ -0,0 +1,45 @@ +"""Process-level lru_cache for default MouseModel / ScrollModel instances.""" +from __future__ import annotations + +from collections.abc import Sequence +from functools import lru_cache +from pathlib import Path + +from ai_mouse.mouse import MouseModel +from ai_mouse.scroll import ScrollModel + + +@lru_cache(maxsize=4) +def _get_mouse_model( + model_key: str, + providers_key: tuple[str, ...], +) -> MouseModel: + path = None if model_key == "__bundled__" else Path(model_key) + providers = list(providers_key) if providers_key else None + return MouseModel(model_path=path, providers=providers) + + +@lru_cache(maxsize=4) +def _get_scroll_model( + model_key: str, + providers_key: tuple[str, ...], +) -> ScrollModel: + path = None if model_key == "__bundled__" else Path(model_key) + providers = list(providers_key) if providers_key else None + return ScrollModel(model_path=path, providers=providers) + + +def get_mouse_model( + model_path: str | Path | None, + providers: Sequence[str] | None, +) -> MouseModel: + key = "__bundled__" if model_path is None else str(model_path) + return _get_mouse_model(key, tuple(providers or ())) + + +def get_scroll_model( + model_path: str | Path | None, + providers: Sequence[str] | None, +) -> ScrollModel: + key = "__bundled__" if model_path is None else str(model_path) + return _get_scroll_model(key, tuple(providers or ())) diff --git a/src/ai_mouse/_postprocess.py b/src/ai_mouse/_postprocess.py new file mode 100644 index 0000000..3091703 --- /dev/null +++ b/src/ai_mouse/_postprocess.py @@ -0,0 +1,180 @@ +"""Pure-numpy post-processing primitives for trajectory generation. + +All functions are pure (no I/O, no global state) and accept an explicit +:class:`numpy.random.Generator` when randomness is involved. +""" +from __future__ import annotations + +import numpy as np + + +def gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray: + """5-tap gaussian smoothing along a 1-D array; endpoints preserved. + + Args: + x: 1-D input array. + sigma: gaussian std. Default 1.0 gives weights approximately + [0.054, 0.244, 0.403, 0.244, 0.054]. + + Returns: + Smoothed array of the same shape. ``x[0]`` and ``x[-1]`` unchanged. + If ``len(x) < 5`` returns a copy of ``x`` (kernel won't fit). + """ + if len(x) < 5: + return x.copy() + kernel = np.exp(-0.5 * (np.arange(-2, 3) / sigma) ** 2) + kernel /= kernel.sum() + padded = np.pad(x, pad_width=2, mode="edge") + smoothed = np.convolve(padded, kernel, mode="valid") + smoothed[0] = x[0] + smoothed[-1] = x[-1] + return smoothed + + +def snap_endpoints( + forward: np.ndarray, + lateral: np.ndarray, + seq_len: int, + n_snap: int = 6, +) -> tuple[np.ndarray, np.ndarray]: + """Force first point to (0,0) and last point to (1,0) with quadratic ease. + + The last ``n_snap`` points are linearly interpolated towards (1, 0) + with quadratic easing, then the first/last points are pinned exactly. + + Args: + forward: (T,) forward coordinates (modified in place). + lateral: (T,) lateral coordinates (modified in place). + seq_len: length of forward/lateral. + n_snap: number of trailing points to ease (capped at seq_len//4). + + Returns: + ``(forward, lateral)`` after modification. + """ + n_snap = min(n_snap, seq_len // 4) + for i in range(n_snap): + alpha = ((i + 1) / n_snap) ** 2 + k = seq_len - n_snap + i + forward[k] = forward[k] * (1.0 - alpha) + 1.0 * alpha + lateral[k] = lateral[k] * (1.0 - alpha) + 0.0 * alpha + forward[0], lateral[0] = 0.0, 0.0 + forward[-1], lateral[-1] = 1.0, 0.0 + return forward, lateral + + +def smooth_start( + forward: np.ndarray, + lateral: np.ndarray, + n: int = 4, +) -> tuple[np.ndarray, np.ndarray]: + """Dampen lateral oscillation in the first ``n`` points. + + Assumes :func:`snap_endpoints` has already pinned (0,0). Forward is + forced non-decreasing locally; lateral is linearly damped towards 0. + """ + n_start_fix = min(n, len(forward) // 4) + for i in range(1, n_start_fix + 1): + blend = i / (n_start_fix + 1) + forward[i] = max(forward[i], forward[i - 1]) + lateral[i] = lateral[i] * blend + return forward, lateral + + +def enforce_forward_monotonic(forward: np.ndarray) -> np.ndarray: + """Force ``forward`` non-decreasing, clip to [0,1], pin endpoints.""" + seq_len = len(forward) + for i in range(1, seq_len - 1): + if forward[i] < forward[i - 1]: + forward[i] = forward[i - 1] + 0.001 + forward = np.clip(forward, 0.0, 1.0) + forward[0] = 0.0 + forward[-1] = 1.0 + return forward + + +def resample_arc(xy: np.ndarray, n_points: int) -> np.ndarray: + """Resample a 2-D polyline to ``n_points`` along cumulative arc length.""" + arc = np.concatenate( + [[0], np.cumsum(np.linalg.norm(np.diff(xy, axis=0), axis=1))] + ) + s_new = np.linspace(0, arc[-1], n_points) + return np.stack( + [np.interp(s_new, arc, xy[:, 0]), np.interp(s_new, arc, xy[:, 1])], + axis=1, + ) + + +def build_timestamps( + log_dt: np.ndarray, + total_duration_ms: float, + dt_clip: tuple[float, float] = (2.0, 150.0), +) -> np.ndarray: + """Convert per-step log_dt + total duration to cumulative ms timestamps. + + Args: + log_dt: (N,) array of natural-log step intervals. + total_duration_ms: target total span. The output is scaled so the + sum approximately matches this (modulo dt_clip). + dt_clip: (min, max) per-step clamp in milliseconds. + + Returns: + (N,) integer-rounded cumulative timestamps starting at 0, + strictly increasing. + """ + n = len(log_dt) + dt_raw = np.clip(np.exp(log_dt), 0.0, None) + dt_sum = dt_raw.sum() + if dt_sum > 1e-6: + scale = total_duration_ms / dt_sum + else: + scale = total_duration_ms / max(n, 1) + dt_ms = np.clip(dt_raw * scale, dt_clip[0], dt_clip[1]) + + t_abs = np.cumsum(dt_ms) + t_abs = np.concatenate([[0.0], t_abs[:-1]]) + + for i in range(1, n): + if t_abs[i] <= t_abs[i - 1]: + t_abs[i] = t_abs[i - 1] + 1.0 + return t_abs + + +def truncnorm_sample( + mu: float, + sigma: float, + low: float, + high: float, + rng: np.random.Generator, + max_tries: int = 32, +) -> float: + """Sample from N(mu, sigma) truncated to [low, high] via rejection. + + Falls back to clipping if rejection fails ``max_tries`` times. + """ + for _ in range(max_tries): + v = rng.normal(mu, sigma) + if low <= v <= high: + return float(v) + return float(np.clip(rng.normal(mu, sigma), low, high)) + + +def sample_duration( + duration_dist: dict, + dist: float, + rng: np.random.Generator, +) -> float: + """Sample total trajectory duration (ms) for the given pixel distance. + + Uses per-bin log-normal parameters in ``duration_dist``. + """ + bins = duration_dist["bins"] + params = duration_dist["params"] + bin_idx = len(bins) - 1 + for i in range(len(bins) - 1): + if dist < bins[i + 1]: + bin_idx = i + break + bin_idx = min(bin_idx, len(params) - 1) + mu_log = params[bin_idx]["mu_log"] + sigma_log = params[bin_idx]["sigma_log"] + return float(np.exp(rng.normal(mu_log, sigma_log))) diff --git a/src/ai_mouse/assets/click_dist.json b/src/ai_mouse/assets/click_dist.json new file mode 100644 index 0000000..f0b15e6 --- /dev/null +++ b/src/ai_mouse/assets/click_dist.json @@ -0,0 +1,6 @@ +{ + "mu": 88.84602649006622, + "sigma": 22.60819023425299, + "low": 20.0, + "high": 500.0 +} \ No newline at end of file diff --git a/src/ai_mouse/assets/duration_dist.json b/src/ai_mouse/assets/duration_dist.json new file mode 100644 index 0000000..466d2c5 --- /dev/null +++ b/src/ai_mouse/assets/duration_dist.json @@ -0,0 +1,47 @@ +{ + "bins": [ + 0, + 50, + 100, + 200, + 400, + 600, + 800, + 1200, + Infinity + ], + "params": [ + { + "mu_log": 5.881647501269015, + "sigma_log": 0.05 + }, + { + "mu_log": 6.141460896343142, + "sigma_log": 0.43685929770118687 + }, + { + "mu_log": 6.409513572943365, + "sigma_log": 0.33676219820945824 + }, + { + "mu_log": 6.590028978462768, + "sigma_log": 0.2675760001657101 + }, + { + "mu_log": 6.70168704031046, + "sigma_log": 0.2809202882427188 + }, + { + "mu_log": 6.758762243710234, + "sigma_log": 0.21524718926818573 + }, + { + "mu_log": 6.214608098422191, + "sigma_log": 0.5 + }, + { + "mu_log": 6.214608098422191, + "sigma_log": 0.5 + } + ] +} \ No newline at end of file diff --git a/src/ai_mouse/assets/flow_model.onnx b/src/ai_mouse/assets/flow_model.onnx new file mode 100644 index 0000000..31652c7 Binary files /dev/null and b/src/ai_mouse/assets/flow_model.onnx differ diff --git a/src/ai_mouse/assets/scroll_config.json b/src/ai_mouse/assets/scroll_config.json new file mode 100644 index 0000000..9275938 --- /dev/null +++ b/src/ai_mouse/assets/scroll_config.json @@ -0,0 +1,13 @@ +{ + "seq_len": 32, + "latent_dim": 16, + "hidden": 64, + "cond_dim": 7, + "epochs": 100, + "batch_size": 32, + "lr": 0.0005, + "beta_max": 0.3, + "beta_anneal_epochs": 15, + "w_delta": 1.0, + "w_logdt": 1.5 +} \ No newline at end of file diff --git a/src/ai_mouse/assets/scroll_decoder.onnx b/src/ai_mouse/assets/scroll_decoder.onnx new file mode 100644 index 0000000..96829a2 Binary files /dev/null and b/src/ai_mouse/assets/scroll_decoder.onnx differ diff --git a/src/ai_mouse/assets/train_config.json b/src/ai_mouse/assets/train_config.json new file mode 100644 index 0000000..79d2ca9 --- /dev/null +++ b/src/ai_mouse/assets/train_config.json @@ -0,0 +1,12 @@ +{ + "seq_len": 64, + "epochs": 50, + "batch_size": 64, + "lr": 5e-06, + "d_model": 128, + "nhead": 4, + "num_layers": 4, + "dim_feedforward": 256, + "dropout": 0.1, + "cond_dim": 3 +} \ No newline at end of file diff --git a/src/ai_mouse/errors.py b/src/ai_mouse/errors.py new file mode 100644 index 0000000..d5cd1fd --- /dev/null +++ b/src/ai_mouse/errors.py @@ -0,0 +1,18 @@ +"""Exception hierarchy for the ai_mouse library. + +Downstream consumers can catch the umbrella :class:`AiMouseError` +or the specific subclasses for finer control. +""" +from __future__ import annotations + + +class AiMouseError(Exception): + """Base class for all ai_mouse errors.""" + + +class ModelLoadError(AiMouseError): + """Raised when ONNX weights / metadata cannot be loaded.""" + + +class GenerationError(AiMouseError): + """Raised when inference produces an invalid result (e.g. NaN).""" diff --git a/src/ai_mouse/mouse.py b/src/ai_mouse/mouse.py new file mode 100644 index 0000000..80629f8 --- /dev/null +++ b/src/ai_mouse/mouse.py @@ -0,0 +1,166 @@ +"""MouseModel — ONNX Runtime-backed mouse trajectory generation.""" +from __future__ import annotations + +import json +import math +from collections.abc import Sequence +from pathlib import Path +from typing import Optional + +import numpy as np +import onnxruntime as ort + +from ai_mouse._assets import resolve +from ai_mouse._coord import decode_trajectory +from ai_mouse._postprocess import ( + build_timestamps, + enforce_forward_monotonic, + gaussian_smooth, + resample_arc, + sample_duration, + smooth_start, + snap_endpoints, + truncnorm_sample, +) +from ai_mouse.errors import GenerationError, ModelLoadError + +_N_EULER_STEPS = 10 + + +class MouseModel: + """Persistent ONNX Runtime session for mouse trajectory generation.""" + + def __init__( + self, + model_path: str | Path | None = None, + providers: Sequence[str] | None = None, + seed: int | None = None, + ) -> None: + path_obj: Optional[Path] = Path(model_path) if model_path is not None else None + + onnx_path = resolve(path_obj, "flow_model.onnx") + cfg_path = resolve(path_obj, "train_config.json") + click_path = resolve(path_obj, "click_dist.json") + dur_path = resolve(path_obj, "duration_dist.json") + + cfg = json.loads(cfg_path.read_text()) + self._seq_len = int(cfg["seq_len"]) + self._cond_dim = int(cfg.get("cond_dim", 3)) + self._click_params = json.loads(click_path.read_text()) + self._duration_dist = json.loads(dur_path.read_text()) + + try: + self._session = ort.InferenceSession( + str(onnx_path), + providers=list(providers) if providers else ["CPUExecutionProvider"], + ) + except Exception as exc: + raise ModelLoadError(f"Failed to load ONNX session: {exc}") from exc + + self._default_seed = seed + self._rng = np.random.default_rng(seed) + + def generate( + self, + start: tuple[int, int], + end: tuple[int, int], + n_points: int = 64, + speed: float | None = None, + click: bool = True, + seed: int | None = None, + ) -> list[tuple[int, int, int]]: + rng = np.random.default_rng(seed) if seed is not None else self._rng + + sx, sy = float(start[0]), float(start[1]) + ex, ey = float(end[0]), float(end[1]) + dist = max(math.hypot(ex - sx, ey - sy), 1.0) + + total_duration = sample_duration(self._duration_dist, dist, rng) + if speed is not None and speed > 0: + total_duration /= speed + total_duration = max(total_duration, 10.0) + + cond = np.array( + [ + dist / 2000.0, + math.log(dist / 100.0), + math.log(total_duration / 500.0), + ], + dtype=np.float32, + )[None] + + x = rng.standard_normal((1, self._seq_len, 3)).astype(np.float32) + dt = 1.0 / _N_EULER_STEPS + for step in range(_N_EULER_STEPS): + t = np.full((1,), step * dt, dtype=np.float32) + v = self._session.run(["v"], {"x_t": x, "t": t, "cond": cond})[0] + x = x + v * dt + + if not np.all(np.isfinite(x)): + raise GenerationError("Trajectory contains NaN/Inf after Euler integration") + + forward = x[0, :, 0].copy() + lateral = x[0, :, 1].copy() + log_dt = x[0, :, 2].copy() + + forward, lateral = snap_endpoints(forward, lateral, self._seq_len) + forward, lateral = smooth_start(forward, lateral) + forward = enforce_forward_monotonic(forward) + lateral = gaussian_smooth(lateral, sigma=1.0) + + log_dt = np.clip(log_dt, 0.0, 5.0) + log_dt[0] = 0.0 + + normalised = np.stack([forward, lateral], axis=1) + pixels = decode_trajectory(normalised, start, end) + + if n_points != self._seq_len: + pixels = resample_arc(pixels, n_points) + log_dt = np.interp( + np.linspace(0, 1, n_points), + np.linspace(0, 1, self._seq_len), + log_dt, + ) + + ts = build_timestamps(log_dt, total_duration) + + moves: list[tuple[int, int, int]] = [ + (int(round(pixels[i, 0])), int(round(pixels[i, 1])), int(round(ts[i]))) + for i in range(n_points) + ] + if not click: + return moves + + click_dur = int( + truncnorm_sample( + float(self._click_params["mu"]), + float(self._click_params["sigma"]), + float(self._click_params["low"]), + float(self._click_params["high"]), + rng, + ) + ) + click_dur = max(click_dur, int(float(self._click_params["low"]))) + last_t = moves[-1][2] + cx, cy = moves[-1][0], moves[-1][1] + return moves + [(cx, cy, last_t), (cx, cy, last_t + click_dur)] + + def sample_click_duration_ms(self, seed: int | None = None) -> int: + rng = np.random.default_rng(seed) if seed is not None else self._rng + v = truncnorm_sample( + float(self._click_params["mu"]), + float(self._click_params["sigma"]), + float(self._click_params["low"]), + float(self._click_params["high"]), + rng, + ) + return max(int(v), int(float(self._click_params["low"]))) + + def close(self) -> None: + self._session = None # type: ignore[assignment] + + def __enter__(self) -> "MouseModel": + return self + + def __exit__(self, *exc) -> None: + self.close() diff --git a/src/ai_mouse/py.typed b/src/ai_mouse/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/src/ai_mouse/scroll.py b/src/ai_mouse/scroll.py new file mode 100644 index 0000000..52d8812 --- /dev/null +++ b/src/ai_mouse/scroll.py @@ -0,0 +1,139 @@ +"""ScrollModel — ONNX Runtime-backed scroll event generation.""" +from __future__ import annotations + +import json +import math +from collections.abc import Sequence +from pathlib import Path +from typing import Literal, Optional + +import numpy as np +import onnxruntime as ort + +from ai_mouse._assets import resolve +from ai_mouse.errors import ModelLoadError + +_DURATION_TABLE = { + "fast": lambda d: d * 0.2 + 100.0, + "precise": lambda d: d * 1.5 + 300.0, + "target": lambda d: d * 0.4 + 200.0, +} + +_QUANTUM = {"precise": 40, "fast": 120, "target": 120} + + +class ScrollModel: + """Persistent ONNX Runtime session for scroll event generation.""" + + def __init__( + self, + model_path: str | Path | None = None, + providers: Sequence[str] | None = None, + seed: int | None = None, + ) -> None: + path_obj: Optional[Path] = Path(model_path) if model_path is not None else None + + onnx_path = resolve(path_obj, "scroll_decoder.onnx") + cfg_path = resolve(path_obj, "scroll_config.json") + cfg = json.loads(cfg_path.read_text()) + + self._seq_len = int(cfg["seq_len"]) + self._latent_dim = int(cfg["latent_dim"]) + self._cond_dim = int(cfg["cond_dim"]) + + try: + self._session = ort.InferenceSession( + str(onnx_path), + providers=list(providers) if providers else ["CPUExecutionProvider"], + ) + except Exception as exc: + raise ModelLoadError(f"Failed to load scroll ONNX session: {exc}") from exc + + self._rng = np.random.default_rng(seed) + + def generate( + self, + start_scroll_y: int, + target_scroll_y: int, + mode: Literal["target", "fast", "precise"] = "target", + viewport_height: int = 800, + seed: int | None = None, + ) -> list[dict]: + rng = np.random.default_rng(seed) if seed is not None else self._rng + + distance = abs(target_scroll_y - start_scroll_y) + direction = 1 if target_scroll_y > start_scroll_y else -1 + distance = max(distance, 10) + + cond = self._build_condition(float(distance), direction, mode, viewport_height) + z = rng.standard_normal((1, self._latent_dim)).astype(np.float32) + decoded = self._session.run(["seq"], {"z": z, "cond": cond[None]})[0][0] + + delta_norm = decoded[:, 0] + log_dt = decoded[:, 1] + + delta_weights = np.exp(delta_norm) + delta_weights /= delta_weights.sum() + delta_px = delta_weights * distance * direction + + quantum = _QUANTUM[mode] + delta_q = np.round(delta_px / quantum) * quantum + for i in range(len(delta_q)): + if delta_q[i] == 0: + delta_q[i] = quantum * direction + delta_q[-1] += (distance * direction) - delta_q.sum() + + if len(log_dt) > 3: + median_log = float(np.median(log_dt)) + log_dt[:2] = np.clip(log_dt[:2], None, median_log + 0.5) + log_dt[-2:] = np.clip(log_dt[-2:], None, median_log + 0.5) + dt_ms = np.clip(np.exp(log_dt), 5, 80) + expected = _DURATION_TABLE[mode](distance) + dt_ms = np.clip(dt_ms * (expected / max(dt_ms.sum(), 1.0)), 5, 80) + + t_abs = np.cumsum(dt_ms).astype(int) + t_abs = np.concatenate([[0], t_abs[:-1]]) + for i in range(1, len(t_abs)): + if t_abs[i] <= t_abs[i - 1]: + t_abs[i] = t_abs[i - 1] + 5 + + events: list[dict] = [] + for i in range(self._seq_len): + dy = int(delta_q[i]) + if dy != 0 or len(events) < 5: + events.append({"deltaY": dy, "deltaMode": 0, "t": int(t_abs[i])}) + return events + + def _build_condition( + self, + distance: float, + direction: int, + mode: str, + viewport_height: int, + ) -> np.ndarray: + mode_onehot = [0.0, 0.0, 0.0] + if mode == "target": + mode_onehot[0] = 1.0 + elif mode == "fast": + mode_onehot[1] = 1.0 + elif mode == "precise": + mode_onehot[2] = 1.0 + return np.array( + [ + distance / 5000.0, + math.log(max(distance, 1.0) / 500.0), + float(direction), + viewport_height / 1000.0, + *mode_onehot, + ], + dtype=np.float32, + ) + + def close(self) -> None: + self._session = None # type: ignore[assignment] + + def __enter__(self) -> "ScrollModel": + return self + + def __exit__(self, *exc) -> None: + self.close() diff --git a/tests/test_coord.py b/tests/test_coord.py deleted file mode 100644 index 0fd376f..0000000 --- a/tests/test_coord.py +++ /dev/null @@ -1,113 +0,0 @@ -"""Tests for rotated coordinate system transforms.""" -from __future__ import annotations - -import math - -import numpy as np -import pytest - -from ai_mouse.coord import encode_trajectory, decode_trajectory - - -class TestEncodeTrajectory: - """Test pixel → rotated normalised frame.""" - - def test_start_maps_to_origin(self): - start = (100, 200) - end = (400, 500) - points = np.array([[100, 200]], dtype=float) - result = encode_trajectory(points, start, end) - np.testing.assert_allclose(result[0], [0.0, 0.0], atol=1e-10) - - def test_end_maps_to_one_zero(self): - start = (100, 200) - end = (400, 500) - points = np.array([[400, 500]], dtype=float) - result = encode_trajectory(points, start, end) - np.testing.assert_allclose(result[0], [1.0, 0.0], atol=1e-10) - - def test_midpoint_maps_to_half_zero(self): - start = (0, 0) - end = (200, 0) - points = np.array([[100, 0]], dtype=float) - result = encode_trajectory(points, start, end) - np.testing.assert_allclose(result[0], [0.5, 0.0], atol=1e-10) - - def test_lateral_offset_positive(self): - """Point at (100, 50) with horizontal start→end has lateral = 50/200 = 0.25.""" - start = (0, 0) - end = (200, 0) - # For horizontal u=(1,0), v=(-0, 1)=(0,1). - # Point (100, 50): forward = 100/200=0.5, lateral = 50/200=0.25 - points = np.array([[100, 50]], dtype=float) - result = encode_trajectory(points, start, end) - np.testing.assert_allclose(result[0], [0.5, 0.25], atol=1e-10) - - def test_various_angles(self): - """Encode/decode round-trip works for various angles.""" - angles = [0, 45, 90, 135, 180, -45, -90, -135] - for deg in angles: - rad = math.radians(deg) - start = (400, 300) - dist = 200 - end = (int(400 + dist * math.cos(rad)), int(300 + dist * math.sin(rad))) - # Create a curved path - t = np.linspace(0, 1, 20) - px = start[0] + t * (end[0] - start[0]) + 20 * np.sin(t * math.pi) - py = start[1] + t * (end[1] - start[1]) + 20 * np.cos(t * math.pi) - points = np.stack([px, py], axis=1) - - encoded = encode_trajectory(points, start, end) - assert encoded[0, 0] == pytest.approx(0.0, abs=0.2) - assert encoded[-1, 0] == pytest.approx(1.0, abs=0.2) - - -class TestDecodeTrajectory: - """Test rotated normalised frame → pixel.""" - - def test_origin_maps_to_start(self): - start = (100, 200) - end = (400, 500) - normalised = np.array([[0.0, 0.0]], dtype=float) - result = decode_trajectory(normalised, start, end) - np.testing.assert_allclose(result[0], [100, 200], atol=1e-10) - - def test_one_zero_maps_to_end(self): - start = (100, 200) - end = (400, 500) - normalised = np.array([[1.0, 0.0]], dtype=float) - result = decode_trajectory(normalised, start, end) - np.testing.assert_allclose(result[0], [400, 500], atol=1e-10) - - -class TestRoundTrip: - """Encode then decode should return original points.""" - - def test_round_trip_horizontal(self): - start = (50, 100) - end = (350, 100) - points = np.array([[50, 100], [150, 130], [250, 90], [350, 100]], dtype=float) - encoded = encode_trajectory(points, start, end) - decoded = decode_trajectory(encoded, start, end) - np.testing.assert_allclose(decoded, points, atol=1e-8) - - def test_round_trip_diagonal(self): - start = (100, 100) - end = (500, 400) - rng = np.random.default_rng(42) - points = np.column_stack([ - np.linspace(100, 500, 30) + rng.normal(0, 10, 30), - np.linspace(100, 400, 30) + rng.normal(0, 10, 30), - ]) - encoded = encode_trajectory(points, start, end) - decoded = decode_trajectory(encoded, start, end) - np.testing.assert_allclose(decoded, points, atol=1e-8) - - def test_round_trip_vertical(self): - """Vertical movement (angle=90°) doesn't collapse.""" - start = (300, 50) - end = (300, 450) - points = np.array([[300, 50], [310, 200], [295, 350], [300, 450]], dtype=float) - encoded = encode_trajectory(points, start, end) - decoded = decode_trajectory(encoded, start, end) - np.testing.assert_allclose(decoded, points, atol=1e-8) diff --git a/tests/test_generator.py b/tests/test_generator.py deleted file mode 100644 index 4bbf09d..0000000 --- a/tests/test_generator.py +++ /dev/null @@ -1,158 +0,0 @@ -"""Tests for Flow Matching trajectory generator.""" -from __future__ import annotations - -import json -from pathlib import Path - -import numpy as np -import pytest -import torch - -from ai_mouse.generator import generate -from ai_mouse.models import TrajectoryFlowModel - - -@pytest.fixture -def model_dir(tmp_path): - """Create temp dir with Flow model artifacts.""" - model = TrajectoryFlowModel(seq_len=64) - torch.save(model.state_dict(), tmp_path / "flow_model.pt") - - click_dist = {"mu": 80.0, "sigma": 30.0, "low": 20.0, "high": 300.0} - (tmp_path / "click_dist.json").write_text(json.dumps(click_dist)) - - duration_dist = { - "bins": [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")], - "params": [ - {"mu_log": 5.5, "sigma_log": 0.3}, - {"mu_log": 5.8, "sigma_log": 0.3}, - {"mu_log": 6.0, "sigma_log": 0.3}, - {"mu_log": 6.2, "sigma_log": 0.3}, - {"mu_log": 6.5, "sigma_log": 0.3}, - {"mu_log": 6.7, "sigma_log": 0.3}, - {"mu_log": 6.9, "sigma_log": 0.3}, - {"mu_log": 7.0, "sigma_log": 0.3}, - ], - } - (tmp_path / "duration_dist.json").write_text(json.dumps(duration_dist)) - - train_config = { - "seq_len": 64, - "d_model": 128, - "nhead": 4, - "num_layers": 4, - "dim_feedforward": 256, - "cond_dim": 3, - } - (tmp_path / "train_config.json").write_text(json.dumps(train_config)) - - return tmp_path - - -class TestGenerate: - def test_returns_list_of_tuples(self, model_dir): - result = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir)) - assert isinstance(result, list) - assert all(isinstance(p, tuple) and len(p) == 3 for p in result) - # All elements are ints - for p in result: - assert all(isinstance(v, int) for v in p) - - def test_timestamps_monotonically_increasing(self, model_dir): - result = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir)) - times = [p[2] for p in result] - for i in range(1, len(times)): - assert times[i] >= times[i - 1] - - def test_starts_near_start(self, model_dir): - start = (100, 200) - result = generate(start=start, end=(500, 400), model_dir=str(model_dir)) - first = result[0] - assert abs(first[0] - start[0]) < 30 - assert abs(first[1] - start[1]) < 30 - - def test_ends_near_end(self, model_dir): - end = (500, 400) - result = generate(start=(100, 200), end=end, model_dir=str(model_dir)) - # Last two are click events; the one before is last movement point - last_move = result[-3] - assert abs(last_move[0] - end[0]) < 30 - assert abs(last_move[1] - end[1]) < 30 - - def test_last_two_are_click_events(self, model_dir): - result = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir)) - down = result[-2] - up = result[-1] - # Same x, y for click down and up - assert down[0] == up[0] - assert down[1] == up[1] - # Up timestamp > down timestamp - assert up[2] > down[2] - # Click duration within bounds - assert 20 <= up[2] - down[2] <= 300 - - def test_different_z_gives_different_paths(self, model_dir): - r1 = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir)) - r2 = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir)) - points1 = [(p[0], p[1]) for p in r1[:-2]] - points2 = [(p[0], p[1]) for p in r2[:-2]] - assert points1 != points2 - - def test_n_points_parameter(self, model_dir): - result = generate( - start=(100, 200), end=(500, 400), n_points=32, model_dir=str(model_dir) - ) - # 32 move points + 2 click events = 34 - assert len(result) == 34 - - -class TestPostProcessing: - def test_dt_diversity_preserved(self, model_dir): - """After removing speed_profile + median clip, multiple generations - should differ in their Δt sequences (not all identical).""" - results = [generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir)) - for _ in range(5)] - # Extract Δt sequences (only move events, not click events) - dts = [] - for r in results: - moves = r[:-2] - dt_seq = [moves[i+1][2] - moves[i][2] for i in range(len(moves)-1)] - dts.append(dt_seq) - # At least 2 of the 5 sequences should differ at any given index - for i in range(min(len(d) for d in dts)): - values = {tuple([d[i]]) for d in dts} - if len(values) > 1: - return # at least one position has variation — pass - pytest.fail("All 5 Δt sequences are identical at every position — diversity collapsed") - - -class TestGaussianSmooth: - def test_endpoints_preserved(self): - from ai_mouse.generator import _gaussian_smooth - x = np.array([1.0, 5.0, 3.0, 7.0, 2.0], dtype=np.float64) - smoothed = _gaussian_smooth(x, sigma=1.0) - assert smoothed[0] == 1.0 - assert smoothed[-1] == 2.0 - - def test_smooths_high_frequency(self): - """A high-frequency square wave should have reduced amplitude after smoothing.""" - from ai_mouse.generator import _gaussian_smooth - x = np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1], dtype=np.float64) - smoothed = _gaussian_smooth(x, sigma=1.0) - # Interior amplitude should be reduced - interior_orig = x[2:-2] - interior_smooth = smoothed[2:-2] - assert interior_smooth.std() < interior_orig.std() - - def test_constant_signal_unchanged(self): - from ai_mouse.generator import _gaussian_smooth - x = np.full(20, 0.5, dtype=np.float64) - smoothed = _gaussian_smooth(x, sigma=1.0) - np.testing.assert_allclose(smoothed, x, rtol=1e-6) - - def test_short_array_returns_unchanged(self): - """Arrays shorter than the kernel are returned unchanged.""" - from ai_mouse.generator import _gaussian_smooth - x = np.array([1.0, 2.0, 3.0], dtype=np.float64) - smoothed = _gaussian_smooth(x, sigma=1.0) - np.testing.assert_allclose(smoothed, x, rtol=1e-6) diff --git a/tests/test_scroll_generator.py b/tests/test_scroll_generator.py deleted file mode 100644 index e5bf238..0000000 --- a/tests/test_scroll_generator.py +++ /dev/null @@ -1,50 +0,0 @@ -"""Tests for scroll generator.""" -from __future__ import annotations - -import json -from pathlib import Path - -import torch -import pytest - -from ai_mouse.scroll.generator import generate_scroll -from ai_mouse.scroll.models import ScrollCVAE - - -@pytest.fixture -def scroll_model_dir(tmp_path): - model = ScrollCVAE(seq_len=32) - torch.save(model.state_dict(), tmp_path / "scroll_model.pt") - config = {"seq_len": 32, "epochs": 100} - (tmp_path / "scroll_config.json").write_text(json.dumps(config)) - return tmp_path - - -class TestGenerateScroll: - def test_returns_list_of_dicts(self, scroll_model_dir): - result = generate_scroll(1000, 3000, mode="target", model_dir=str(scroll_model_dir)) - assert isinstance(result, list) - assert len(result) > 0 - assert all("deltaY" in e and "t" in e and "deltaMode" in e for e in result) - - def test_timestamps_monotonic(self, scroll_model_dir): - result = generate_scroll(1000, 3000, mode="target", model_dir=str(scroll_model_dir)) - times = [e["t"] for e in result] - for i in range(1, len(times)): - assert times[i] >= times[i - 1] - - def test_total_scroll_approximately_matches_distance(self, scroll_model_dir): - result = generate_scroll(1000, 3000, mode="target", model_dir=str(scroll_model_dir)) - total = sum(e["deltaY"] for e in result) - # Should be within 30% of target distance (2000px) - assert abs(total - 2000) < 2000 * 0.4 - - def test_deltaY_are_integers(self, scroll_model_dir): - result = generate_scroll(1000, 3000, mode="target", model_dir=str(scroll_model_dir)) - assert all(isinstance(e["deltaY"], int) for e in result) - - def test_direction_up(self, scroll_model_dir): - result = generate_scroll(3000, 1000, mode="target", model_dir=str(scroll_model_dir)) - total = sum(e["deltaY"] for e in result) - # Negative total for scrolling up - assert total < 0 diff --git a/tests/tools/__init__.py b/tests/tools/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/conftest.py b/tests/tools/conftest.py similarity index 94% rename from tests/conftest.py rename to tests/tools/conftest.py index 3785cef..755b9e3 100644 --- a/tests/conftest.py +++ b/tests/tools/conftest.py @@ -8,8 +8,8 @@ import numpy as np import pytest import torch -from ai_mouse.models import TrajectoryFlowModel -from ai_mouse.scroll.models import ScrollCVAE +from tools.models import TrajectoryFlowModel +from tools.scroll.models import ScrollCVAE @pytest.fixture diff --git a/tests/test_balabit_adapter.py b/tests/tools/test_balabit_adapter.py similarity index 87% rename from tests/test_balabit_adapter.py rename to tests/tools/test_balabit_adapter.py index 995dacb..c8ff90f 100644 --- a/tests/test_balabit_adapter.py +++ b/tests/tools/test_balabit_adapter.py @@ -9,7 +9,7 @@ import pytest def test_module_exports(): """The adapter module must export the public functions used by CLI.""" - from ai_mouse.data_adapters import balabit + from tools.data_adapters import balabit assert hasattr(balabit, "parse_session_csv") assert hasattr(balabit, "segment_by_clicks") assert hasattr(balabit, "filter_segments") @@ -20,7 +20,7 @@ def test_module_exports(): def test_mouse_event_dataclass(): """MouseEvent has expected fields.""" - from ai_mouse.data_adapters.balabit import MouseEvent + from tools.data_adapters.balabit import MouseEvent e = MouseEvent(t_ms=100, button="NoButton", state="Move", x=300, y=400) assert e.t_ms == 100 assert e.state == "Move" @@ -29,7 +29,7 @@ def test_mouse_event_dataclass(): def test_segment_dataclass(): """Segment has expected fields.""" - from ai_mouse.data_adapters.balabit import MouseEvent, Segment + from tools.data_adapters.balabit import MouseEvent, Segment events = [MouseEvent(t_ms=0, button="NoButton", state="Move", x=10, y=20)] s = Segment(events=events, click_x=100, click_y=200, click_t_ms=500, session_id="user1_s1") assert s.events == events @@ -45,7 +45,7 @@ def _write_csv(path: Path, rows: list[str]) -> None: class TestParseSessionCsv: def test_parses_basic_rows(self, tmp_path): - from ai_mouse.data_adapters.balabit import parse_session_csv + from tools.data_adapters.balabit import parse_session_csv csv = tmp_path / "session_1" _write_csv(csv, [ "1500000000.000,0.000,NoButton,Move,100,200", @@ -63,7 +63,7 @@ class TestParseSessionCsv: def test_handles_float_timestamps(self, tmp_path): """Client timestamps are floats in seconds; we convert to int ms.""" - from ai_mouse.data_adapters.balabit import parse_session_csv + from tools.data_adapters.balabit import parse_session_csv csv = tmp_path / "session_2" _write_csv(csv, [ "0,1.234,NoButton,Move,50,60", @@ -75,7 +75,7 @@ class TestParseSessionCsv: def test_skips_malformed_rows(self, tmp_path): """Rows with bad data are logged and skipped, not raised.""" - from ai_mouse.data_adapters.balabit import parse_session_csv + from tools.data_adapters.balabit import parse_session_csv csv = tmp_path / "session_3" _write_csv(csv, [ "0,0.000,NoButton,Move,100,200", @@ -89,7 +89,7 @@ class TestParseSessionCsv: assert events[1].x == 150 def test_returns_empty_list_for_empty_file(self, tmp_path): - from ai_mouse.data_adapters.balabit import parse_session_csv + from tools.data_adapters.balabit import parse_session_csv csv = tmp_path / "session_4" csv.write_text("record timestamp,client timestamp,button,state,x,y\n", encoding="utf-8") events = parse_session_csv(csv) @@ -98,11 +98,11 @@ class TestParseSessionCsv: class TestSegmentByClicks: def _make_event(self, t_ms: int, state: str, x: int, y: int, button: str = "NoButton"): - from ai_mouse.data_adapters.balabit import MouseEvent + from tools.data_adapters.balabit import MouseEvent return MouseEvent(t_ms=t_ms, button=button, state=state, x=x, y=y) def test_one_click_one_segment(self): - from ai_mouse.data_adapters.balabit import segment_by_clicks + from tools.data_adapters.balabit import segment_by_clicks events = [ self._make_event(0, "Move", 10, 20), self._make_event(100, "Move", 50, 60), @@ -120,7 +120,7 @@ class TestSegmentByClicks: def test_window_excludes_old_events(self): """Move events earlier than (click_t - window_ms) are dropped.""" - from ai_mouse.data_adapters.balabit import segment_by_clicks + from tools.data_adapters.balabit import segment_by_clicks events = [ self._make_event(0, "Move", 10, 20), # too old self._make_event(100, "Move", 20, 30), # too old @@ -133,7 +133,7 @@ class TestSegmentByClicks: assert segments[0].events[0].t_ms == 900 def test_multiple_clicks_multiple_segments(self): - from ai_mouse.data_adapters.balabit import segment_by_clicks + from tools.data_adapters.balabit import segment_by_clicks events = [ self._make_event(0, "Move", 10, 20), self._make_event(100, "Pressed", 50, 50, button="Left"), @@ -152,7 +152,7 @@ class TestSegmentByClicks: def test_skips_pressed_with_non_left_button(self): """Right-clicks and wheel-clicks don't anchor segments (only Left).""" - from ai_mouse.data_adapters.balabit import segment_by_clicks + from tools.data_adapters.balabit import segment_by_clicks events = [ self._make_event(0, "Move", 10, 20), self._make_event(100, "Pressed", 50, 50, button="Right"), # ignored @@ -164,7 +164,7 @@ class TestSegmentByClicks: assert segments[0].click_x == 70 def test_no_clicks_returns_empty(self): - from ai_mouse.data_adapters.balabit import segment_by_clicks + from tools.data_adapters.balabit import segment_by_clicks events = [ self._make_event(0, "Move", 10, 20), self._make_event(100, "Move", 20, 30), @@ -174,7 +174,7 @@ class TestSegmentByClicks: def test_excludes_drag_events(self): """Drag events are not Move; segment should only include Move.""" - from ai_mouse.data_adapters.balabit import segment_by_clicks + from tools.data_adapters.balabit import segment_by_clicks events = [ self._make_event(0, "Move", 10, 20), self._make_event(100, "Drag", 30, 40), # not Move @@ -191,7 +191,7 @@ class TestSegmentByClicks: class TestFilterSegments: def _seg(self, events_data: list[tuple[int, int, int]], click=(500, 500), session="s"): """events_data: list of (t_ms, x, y) tuples.""" - from ai_mouse.data_adapters.balabit import MouseEvent, Segment + from tools.data_adapters.balabit import MouseEvent, Segment events = [MouseEvent(t_ms=t, button="NoButton", state="Move", x=x, y=y) for (t, x, y) in events_data] click_t = events[-1].t_ms + 50 if events else 100 @@ -199,14 +199,14 @@ class TestFilterSegments: click_t_ms=click_t, session_id=session) def test_drops_segment_with_too_few_events(self): - from ai_mouse.data_adapters.balabit import filter_segments + from tools.data_adapters.balabit import filter_segments seg = self._seg([(0, 100, 100), (50, 105, 105)]) # 2 events, min=5 result = filter_segments([seg], min_events=5, min_dist=10, max_span_ms=5000, max_gap_ms=200) assert result == [] def test_drops_segment_with_short_distance(self): - from ai_mouse.data_adapters.balabit import filter_segments + from tools.data_adapters.balabit import filter_segments # Start (100,100), end click=(105,100) → dist=5, min_dist=50 events = [(i*10, 100+i, 100) for i in range(10)] seg = self._seg(events, click=(105, 100)) @@ -215,7 +215,7 @@ class TestFilterSegments: assert result == [] def test_drops_segment_with_too_long_span(self): - from ai_mouse.data_adapters.balabit import filter_segments + from tools.data_adapters.balabit import filter_segments # Span 6000ms, max_span=5000 events = [(0, 100, 100), (2000, 200, 200), (4000, 300, 300), (6000, 400, 400), (6010, 410, 400), (6020, 420, 400)] @@ -225,7 +225,7 @@ class TestFilterSegments: assert result == [] def test_drops_segment_with_gap(self): - from ai_mouse.data_adapters.balabit import filter_segments + from tools.data_adapters.balabit import filter_segments # Gap of 500ms between events 2 and 3, max_gap=200 events = [(0, 100, 100), (50, 110, 110), (100, 120, 120), (600, 200, 200), (650, 210, 210), (700, 220, 220)] @@ -235,7 +235,7 @@ class TestFilterSegments: assert result == [] def test_drops_segment_with_out_of_range_coords(self): - from ai_mouse.data_adapters.balabit import filter_segments + from tools.data_adapters.balabit import filter_segments # x=10000 out of range events = [(i*10, 10000, 100) for i in range(6)] seg = self._seg(events, click=(10100, 100)) @@ -245,7 +245,7 @@ class TestFilterSegments: def test_drops_segment_with_short_arc_length(self): """Total arc < 50 even though endpoints are far → high-frequency jitter only.""" - from ai_mouse.data_adapters.balabit import filter_segments + from tools.data_adapters.balabit import filter_segments # Tiny back-and-forth with click 100px away — unrealistic, drop it events = [(i*10, 100 + (i % 2), 100) for i in range(10)] seg = self._seg(events, click=(200, 100)) @@ -255,7 +255,7 @@ class TestFilterSegments: assert result == [] def test_keeps_valid_segment(self): - from ai_mouse.data_adapters.balabit import filter_segments + from tools.data_adapters.balabit import filter_segments # Smooth 100→500 px straight line, 10 events, span 500ms events = [(i*50, 100 + i*40, 100) for i in range(10)] seg = self._seg(events, click=(500, 100)) @@ -266,8 +266,8 @@ class TestFilterSegments: class TestProcessSession: def test_writes_jsonl_in_expected_format(self, tmp_path): - from ai_mouse.config import BalabitAdapterConfig - from ai_mouse.data_adapters.balabit import process_session + from tools.config import BalabitAdapterConfig + from tools.data_adapters.balabit import process_session # Construct a Balabit-format CSV with one valid segment csv_path = tmp_path / "user_session_42" @@ -303,8 +303,8 @@ class TestProcessSession: assert record["events"][0]["t"] == 0 def test_returns_zero_for_session_with_no_valid_segments(self, tmp_path): - from ai_mouse.config import BalabitAdapterConfig - from ai_mouse.data_adapters.balabit import process_session + from tools.config import BalabitAdapterConfig + from tools.data_adapters.balabit import process_session csv_path = tmp_path / "empty_session" _write_csv(csv_path, ["0,0.000,NoButton,Move,100,100"]) # no clicks @@ -318,8 +318,8 @@ class TestProcessSession: assert out.read_text() == "" def test_appends_to_existing_jsonl(self, tmp_path): - from ai_mouse.config import BalabitAdapterConfig - from ai_mouse.data_adapters.balabit import process_session + from tools.config import BalabitAdapterConfig + from tools.data_adapters.balabit import process_session out = tmp_path / "out.jsonl" out.write_text('{"meta":{"start":[0,0],"end":[1,1]},"events":[]}\n', encoding="utf-8") diff --git a/tests/test_eval_metrics.py b/tests/tools/test_eval_metrics.py similarity index 88% rename from tests/test_eval_metrics.py rename to tests/tools/test_eval_metrics.py index 2083447..55bd9c3 100644 --- a/tests/test_eval_metrics.py +++ b/tests/tools/test_eval_metrics.py @@ -8,7 +8,7 @@ import pytest class TestKinematics: def test_compute_speed_constant_velocity(self): """Constant-velocity trajectory has constant speed.""" - from ai_mouse.eval.metrics import compute_speed + from tools.eval.metrics import compute_speed # 10 points, moving 10 px in 100 ms each step → speed = 0.1 px/ms xs = np.arange(0, 100, 10, dtype=float) ys = np.zeros(10, dtype=float) @@ -20,7 +20,7 @@ class TestKinematics: def test_compute_speed_handles_zero_dt(self): """Adjacent points with same timestamp must not produce NaN/inf.""" - from ai_mouse.eval.metrics import compute_speed + from tools.eval.metrics import compute_speed xs = np.array([0.0, 10.0, 20.0]) ys = np.array([0.0, 0.0, 0.0]) ts = np.array([0.0, 0.0, 100.0]) # zero dt between [0] and [1] @@ -29,7 +29,7 @@ class TestKinematics: def test_compute_acceleration(self): """Linearly increasing speed → constant acceleration.""" - from ai_mouse.eval.metrics import compute_acceleration + from tools.eval.metrics import compute_acceleration # speeds: 0.1, 0.2, 0.3, 0.4 over dt = 100 ms each → a = 0.001 px/ms² speeds = np.array([0.1, 0.2, 0.3, 0.4]) ts = np.array([100.0, 200.0, 300.0, 400.0]) @@ -37,7 +37,7 @@ class TestKinematics: np.testing.assert_allclose(a, 0.001, rtol=1e-4) def test_compute_jerk(self): - from ai_mouse.eval.metrics import compute_jerk + from tools.eval.metrics import compute_jerk # accelerations: 0.001, 0.002, 0.003 over dt = 100 ms → j = 0.00001 accels = np.array([0.001, 0.002, 0.003]) ts = np.array([200.0, 300.0, 400.0]) @@ -47,7 +47,7 @@ class TestKinematics: class TestStatsSummary: def test_compute_stats_returns_expected_keys(self): - from ai_mouse.eval.metrics import compute_stats + from tools.eval.metrics import compute_stats x = np.array([1.0, 2.0, 3.0, 4.0, 5.0]) s = compute_stats(x) assert "mean" in s @@ -59,7 +59,7 @@ class TestStatsSummary: assert "p95" in s def test_cv_for_constant_is_zero(self): - from ai_mouse.eval.metrics import compute_stats + from tools.eval.metrics import compute_stats x = np.full(10, 3.0) s = compute_stats(x) assert s["cv"] == 0.0 @@ -68,7 +68,7 @@ class TestStatsSummary: class TestFftSpectrum: def test_finds_dominant_frequency(self): """A pure 8 Hz signal should have its peak near 8 Hz.""" - from ai_mouse.eval.metrics import fft_spectrum + from tools.eval.metrics import fft_spectrum # Sample at 100 Hz for 1 second sample_rate_hz = 100.0 ts_ms = np.arange(0, 1000, 1000 / sample_rate_hz) @@ -78,7 +78,7 @@ class TestFftSpectrum: assert abs(peak_freq - 8.0) < 1.0 # within 1 Hz def test_returns_only_positive_frequencies(self): - from ai_mouse.eval.metrics import fft_spectrum + from tools.eval.metrics import fft_spectrum signal = np.random.randn(64) freqs, mags = fft_spectrum(signal, 50.0) assert (freqs >= 0).all() @@ -88,7 +88,7 @@ class TestFftSpectrum: class TestKlDivergence: def test_identical_distributions_zero_kl(self): """KL(p, p) ≈ 0.""" - from ai_mouse.eval.metrics import kl_divergence_histograms + from tools.eval.metrics import kl_divergence_histograms rng = np.random.default_rng(42) x = rng.normal(0, 1, 5000) y = rng.normal(0, 1, 5000) @@ -97,7 +97,7 @@ class TestKlDivergence: def test_different_distributions_positive_kl(self): """Different means → positive KL.""" - from ai_mouse.eval.metrics import kl_divergence_histograms + from tools.eval.metrics import kl_divergence_histograms rng = np.random.default_rng(42) x = rng.normal(0, 1, 5000) y = rng.normal(3, 1, 5000) @@ -106,7 +106,7 @@ class TestKlDivergence: def test_handles_disjoint_supports(self): """No NaN even when histograms have non-overlapping bins.""" - from ai_mouse.eval.metrics import kl_divergence_histograms + from tools.eval.metrics import kl_divergence_histograms x = np.array([1.0, 1.1, 1.2, 1.3, 1.4]) y = np.array([10.0, 10.1, 10.2, 10.3, 10.4]) kl = kl_divergence_histograms(x, y, bins=10) @@ -116,7 +116,7 @@ class TestKlDivergence: class TestReportGeneration: def test_generates_report_md(self, tmp_path): """Smoke test: build_report writes an MD file with all expected sections.""" - from ai_mouse.eval.report import build_report + from tools.eval.report import build_report # Synthetic generated traces (3 traces, 50 points each) rng = np.random.default_rng(0) diff --git a/tests/tools/test_export_onnx.py b/tests/tools/test_export_onnx.py new file mode 100644 index 0000000..b72f34b --- /dev/null +++ b/tests/tools/test_export_onnx.py @@ -0,0 +1,66 @@ +"""Validate tools.export_onnx with a tiny synthetic model.""" +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import torch + +from tools.export_onnx import ( + _check_flow_parity, + _check_scroll_parity, + export_flow_model, + export_scroll_decoder, +) + + +@pytest.fixture +def tiny_flow_ckpt(tmp_path: Path) -> Path: + """A flow model with seq_len=8, d_model=16, 1 layer — small but valid.""" + from tools.models import TrajectoryFlowModel + + cfg = { + "seq_len": 8, + "d_model": 16, + "nhead": 2, + "num_layers": 1, + "dim_feedforward": 32, + "cond_dim": 3, + } + model = TrajectoryFlowModel(**cfg, dropout=0.0) + model.eval() + out = tmp_path / "flow_ckpt" + out.mkdir() + torch.save(model.state_dict(), out / "flow_model.pt") + (out / "train_config.json").write_text(json.dumps(cfg)) + return out + + +@pytest.fixture +def tiny_scroll_ckpt(tmp_path: Path) -> Path: + """A scroll model with seq_len=4, latent=4, hidden=8.""" + from tools.scroll.models import ScrollCVAE + + cfg = {"seq_len": 4, "latent_dim": 4, "hidden": 8, "cond_dim": 7} + model = ScrollCVAE(**cfg) + model.eval() + out = tmp_path / "scroll_ckpt" + out.mkdir() + torch.save(model.state_dict(), out / "scroll_model.pt") + (out / "scroll_config.json").write_text(json.dumps(cfg)) + return out + + +def test_export_flow_model_parity(tiny_flow_ckpt: Path, tmp_path: Path) -> None: + out_dir = tmp_path / "out" + onnx_path = export_flow_model(tiny_flow_ckpt, out_dir) + assert onnx_path.exists() + _check_flow_parity(tiny_flow_ckpt, onnx_path) # raises on failure + + +def test_export_scroll_decoder_parity(tiny_scroll_ckpt: Path, tmp_path: Path) -> None: + out_dir = tmp_path / "out" + onnx_path = export_scroll_decoder(tiny_scroll_ckpt, out_dir) + assert onnx_path.exists() + _check_scroll_parity(tiny_scroll_ckpt, onnx_path) diff --git a/tests/test_models.py b/tests/tools/test_models.py similarity index 97% rename from tests/test_models.py rename to tests/tools/test_models.py index bfe81fc..0a6674a 100644 --- a/tests/test_models.py +++ b/tests/tools/test_models.py @@ -4,7 +4,7 @@ from __future__ import annotations import torch import pytest -from ai_mouse.models import TrajectoryFlowModel +from tools.models import TrajectoryFlowModel class TestTrajectoryFlowModel: diff --git a/tests/test_scroll_collector.py b/tests/tools/test_scroll_collector.py similarity index 98% rename from tests/test_scroll_collector.py rename to tests/tools/test_scroll_collector.py index f30c328..b632281 100644 --- a/tests/test_scroll_collector.py +++ b/tests/tools/test_scroll_collector.py @@ -3,7 +3,7 @@ from __future__ import annotations import pytest -from ai_mouse.scroll.collector import ScrollCollector +from tools.scroll.collector import ScrollCollector class TestNextTarget: diff --git a/tests/test_scroll_models.py b/tests/tools/test_scroll_models.py similarity index 95% rename from tests/test_scroll_models.py rename to tests/tools/test_scroll_models.py index 98175e8..0d84715 100644 --- a/tests/test_scroll_models.py +++ b/tests/tools/test_scroll_models.py @@ -4,7 +4,7 @@ from __future__ import annotations import torch import pytest -from ai_mouse.scroll.models import ScrollCVAE +from tools.scroll.models import ScrollCVAE class TestScrollCVAEForward: diff --git a/tests/test_scroll_trainer.py b/tests/tools/test_scroll_trainer.py similarity index 96% rename from tests/test_scroll_trainer.py rename to tests/tools/test_scroll_trainer.py index 7cb9e53..6ffc7b7 100644 --- a/tests/test_scroll_trainer.py +++ b/tests/tools/test_scroll_trainer.py @@ -8,7 +8,7 @@ from pathlib import Path import numpy as np import pytest -from ai_mouse.scroll.trainer import load_scroll_data, train_scroll, _augment_scroll +from tools.scroll.trainer import load_scroll_data, train_scroll, _augment_scroll def _make_synthetic_scroll_trace(mode="target"): diff --git a/tests/test_server.py b/tests/tools/test_server.py similarity index 96% rename from tests/test_server.py rename to tests/tools/test_server.py index 756c5a4..d10fc00 100644 --- a/tests/test_server.py +++ b/tests/tools/test_server.py @@ -7,8 +7,8 @@ import pytest import pytest_asyncio from httpx import ASGITransport, AsyncClient -from ai_mouse.server import create_app -from ai_mouse.server.deps import get_data_dir +from tools.server import create_app +from tools.server.deps import get_data_dir @pytest.fixture @@ -87,7 +87,7 @@ class TestCollect: async def test_collect_trace_increments_count(self, client, tmp_path, monkeypatch): """Test that posting a trace increments the collected count.""" # Monkeypatch data dir to use tmp - import ai_mouse.server.deps as deps + import tools.server.deps as deps monkeypatch.setattr(deps, "_DATA_DIR", tmp_path) # Start collection @@ -124,7 +124,7 @@ class TestVerify: @pytest.mark.asyncio async def test_verify_returns_paths(self, client, model_dir, monkeypatch): """Test trajectory generation endpoint.""" - import ai_mouse.server.routes_verify as rv + import tools.server.routes_verify as rv # We can't easily monkeypatch the model dir used inside the route # but we can test the endpoint is accessible resp = await client.post( diff --git a/tests/test_trainer.py b/tests/tools/test_trainer.py similarity index 93% rename from tests/test_trainer.py rename to tests/tools/test_trainer.py index 8cbf79b..a6d6b24 100644 --- a/tests/test_trainer.py +++ b/tests/tools/test_trainer.py @@ -8,7 +8,7 @@ from pathlib import Path import numpy as np import pytest -from ai_mouse.trainer import load_and_prepare_data, train, _augment +from tools.trainer import load_and_prepare_data, train, _augment def _make_synthetic_trace(start, end, n_moves=30): @@ -124,21 +124,21 @@ class TestTrain: class TestTrajectoryDataset: def test_dataset_length_with_augmentation(self): """Dataset length = N * 6 when augment=True.""" - from ai_mouse.trainer import TrajectoryDataset + from tools.trainer import TrajectoryDataset seq = np.zeros((10, 64, 3), dtype=np.float32) cond = np.zeros((10, 3), dtype=np.float32) ds = TrajectoryDataset(seq, cond, augment=True) assert len(ds) == 60 def test_dataset_length_without_augmentation(self): - from ai_mouse.trainer import TrajectoryDataset + from tools.trainer import TrajectoryDataset seq = np.zeros((10, 64, 3), dtype=np.float32) cond = np.zeros((10, 3), dtype=np.float32) ds = TrajectoryDataset(seq, cond, augment=False) assert len(ds) == 10 def test_getitem_returns_tensors(self): - from ai_mouse.trainer import TrajectoryDataset + from tools.trainer import TrajectoryDataset import torch seq = np.random.randn(5, 64, 3).astype(np.float32) cond = np.random.randn(5, 3).astype(np.float32) @@ -151,7 +151,7 @@ class TestTrajectoryDataset: def test_aug_id_zero_returns_original(self): """Aug id 0 (idx=0 % 6 == 0) should return the original sample unchanged.""" - from ai_mouse.trainer import TrajectoryDataset + from tools.trainer import TrajectoryDataset import torch seq = np.array([[[0.5, 0.7, 0.3]] * 64] * 3, dtype=np.float32) cond = np.array([[1.0, 2.0, 3.0]] * 3, dtype=np.float32) @@ -162,7 +162,7 @@ class TestTrajectoryDataset: def test_aug_id_one_flips_lateral(self): """Aug id 1 should flip the sign of the lateral channel (index 1).""" - from ai_mouse.trainer import TrajectoryDataset + from tools.trainer import TrajectoryDataset seq = np.zeros((1, 64, 3), dtype=np.float32) seq[0, :, 1] = 0.5 # lateral all positive cond = np.zeros((1, 3), dtype=np.float32) @@ -174,7 +174,7 @@ class TestTrajectoryDataset: def test_aug_id_two_slows_speed(self): """Aug id 2 should add log(1.25) to log_dt channel and cond[2].""" import math - from ai_mouse.trainer import TrajectoryDataset + from tools.trainer import TrajectoryDataset seq = np.zeros((1, 64, 3), dtype=np.float32) cond = np.zeros((1, 3), dtype=np.float32) ds = TrajectoryDataset(seq, cond, augment=True) @@ -186,7 +186,7 @@ class TestTrajectoryDataset: def test_aug_id_three_speeds_up(self): """Aug id 3 should add log(1/1.2) to log_dt channel and cond[2].""" import math - from ai_mouse.trainer import TrajectoryDataset + from tools.trainer import TrajectoryDataset seq = np.zeros((1, 64, 3), dtype=np.float32) cond = np.zeros((1, 3), dtype=np.float32) ds = TrajectoryDataset(seq, cond, augment=True) @@ -197,7 +197,7 @@ class TestTrajectoryDataset: def test_aug_id_four_adds_temporal_noise(self): """Aug id 4 should add Gaussian noise to log_dt (channel 2), leaving other channels unchanged.""" - from ai_mouse.trainer import TrajectoryDataset + from tools.trainer import TrajectoryDataset seq = np.zeros((1, 64, 3), dtype=np.float32) cond = np.zeros((1, 3), dtype=np.float32) ds = TrajectoryDataset(seq, cond, augment=True) @@ -215,7 +215,7 @@ class TestTrajectoryDataset: def test_aug_id_five_flips_and_slows(self): """Aug id 5 should flip lateral and add log(1/0.9) to log_dt and cond[2].""" import math - from ai_mouse.trainer import TrajectoryDataset + from tools.trainer import TrajectoryDataset seq = np.zeros((1, 64, 3), dtype=np.float32) seq[0, :, 1] = 0.5 # lateral positive cond = np.zeros((1, 3), dtype=np.float32) @@ -233,8 +233,8 @@ class TestResumeFrom: def test_resume_from_loads_checkpoint(self, synthetic_traces_file, tmp_path): """train() with resume_from should load weights from given checkpoint dir.""" import torch - from ai_mouse.trainer import train - from ai_mouse.models import TrajectoryFlowModel + from tools.trainer import train + from tools.models import TrajectoryFlowModel # First, train an initial model and save it ckpt_dir = tmp_path / "pretrain" @@ -273,7 +273,7 @@ class TestResumeFrom: assert diff < 0.5, f"Resume_from weights diverged too much: {diff}" def test_resume_from_missing_path_raises(self, synthetic_traces_file, tmp_path): - from ai_mouse.trainer import train + from tools.trainer import train with pytest.raises(FileNotFoundError): train( diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py new file mode 100644 index 0000000..7f48be2 --- /dev/null +++ b/tests/unit/conftest.py @@ -0,0 +1 @@ +"""Fixtures for library-only tests (no torch).""" diff --git a/tests/unit/data/golden_mouse.npz b/tests/unit/data/golden_mouse.npz new file mode 100644 index 0000000..079f08c Binary files /dev/null and b/tests/unit/data/golden_mouse.npz differ diff --git a/tests/unit/data/golden_scroll.npz b/tests/unit/data/golden_scroll.npz new file mode 100644 index 0000000..5d3dd6c Binary files /dev/null and b/tests/unit/data/golden_scroll.npz differ diff --git a/tests/unit/test__coord.py b/tests/unit/test__coord.py new file mode 100644 index 0000000..ab8b422 --- /dev/null +++ b/tests/unit/test__coord.py @@ -0,0 +1,30 @@ +"""Test the private numpy coordinate transforms.""" +from __future__ import annotations + +import numpy as np + +from ai_mouse._coord import decode_trajectory, encode_trajectory + + +def test_encode_decode_roundtrip() -> None: + points = np.array([[100.0, 200.0], [300.0, 250.0], [500.0, 300.0]]) + start = (100, 200) + end = (500, 300) + encoded = encode_trajectory(points, start, end) + decoded = decode_trajectory(encoded, start, end) + assert np.allclose(decoded, points, atol=1e-6) + + +def test_encode_endpoints() -> None: + """Start should encode to (0,0); end should encode to (1,0).""" + points = np.array([[100.0, 200.0], [500.0, 300.0]]) + encoded = encode_trajectory(points, (100, 200), (500, 300)) + assert np.allclose(encoded[0], [0.0, 0.0], atol=1e-6) + assert np.allclose(encoded[1], [1.0, 0.0], atol=1e-6) + + +def test_zero_distance_returns_zeros() -> None: + points = np.array([[100.0, 200.0]]) + encoded = encode_trajectory(points, (100, 200), (100, 200)) + assert encoded.shape == (1, 2) + assert np.all(encoded == 0) diff --git a/tests/unit/test_assets.py b/tests/unit/test_assets.py new file mode 100644 index 0000000..c52512a --- /dev/null +++ b/tests/unit/test_assets.py @@ -0,0 +1,34 @@ +"""Test the asset path resolver.""" +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from ai_mouse import _assets +from ai_mouse.errors import ModelLoadError + + +def test_bundled_flow_model_exists() -> None: + p = _assets.bundled_path("flow_model.onnx") + assert p.exists() + assert p.suffix == ".onnx" + + +def test_bundled_train_config_loadable() -> None: + p = _assets.bundled_path("train_config.json") + cfg = json.loads(p.read_text()) + assert "seq_len" in cfg + assert "d_model" in cfg + + +def test_resolve_with_custom_dir(tmp_path: Path) -> None: + (tmp_path / "flow_model.onnx").write_bytes(b"x") + p = _assets.resolve(tmp_path, "flow_model.onnx") + assert p == tmp_path / "flow_model.onnx" + + +def test_missing_asset_raises_model_load_error(tmp_path: Path) -> None: + with pytest.raises(ModelLoadError, match="missing"): + _assets.resolve(tmp_path, "nonexistent.onnx") diff --git a/tests/unit/test_errors.py b/tests/unit/test_errors.py new file mode 100644 index 0000000..9f9cc10 --- /dev/null +++ b/tests/unit/test_errors.py @@ -0,0 +1,19 @@ +"""Test the error hierarchy.""" +from __future__ import annotations + +import pytest + +from ai_mouse import errors + + +def test_model_load_error_is_aimouse_error() -> None: + assert issubclass(errors.ModelLoadError, errors.AiMouseError) + + +def test_generation_error_is_aimouse_error() -> None: + assert issubclass(errors.GenerationError, errors.AiMouseError) + + +def test_can_catch_specific_with_general() -> None: + with pytest.raises(errors.AiMouseError): + raise errors.ModelLoadError("test") diff --git a/tests/unit/test_golden.py b/tests/unit/test_golden.py new file mode 100644 index 0000000..d76f015 --- /dev/null +++ b/tests/unit/test_golden.py @@ -0,0 +1,151 @@ +"""Golden regression tests — guard against catastrophic divergence from +the pre-migration PyTorch implementation. + +Important caveat on tolerances +============================== + +These goldens were captured with the legacy PyTorch + scipy stack, where the +RNG path was ``torch.manual_seed(seed)`` + ``torch.randn(...)`` plus +``np.random.seed`` for the duration sampler. The rewritten library uses a +single ``np.random.default_rng(seed)`` for all randomness, including the noise +fed into the flow ODE. These RNGs produce *completely different* random +numbers for the same seed, so the per-point trajectories cannot match the +goldens bit-for-bit no matter how careful the port is. + +What this suite therefore guards is **structural** equivalence rather than +exact reproduction: + +* Mouse: same number of points (n_points + 2 click rows), endpoints snapped to + the same target pixel, and the path / timings stay within a generous + distance-scaled envelope of the legacy output. +* Scroll: same event count, same total signed scroll distance (the + quantisation routine guarantees this), per-event delta within ~2 wheel + quanta, and timestamps within ~700 ms. + +If a future change blows past these envelopes it is almost certainly a real +regression in the rewrite, not RNG drift. To re-baseline the goldens after an +intentional change, regenerate the .npz files from the new implementation. +""" +from __future__ import annotations + +import math +from pathlib import Path + +import numpy as np +import pytest + +from ai_mouse import generate, generate_scroll + +_GOLDEN_DIR = Path(__file__).parent / "data" + +_MOUSE_CASES: list[tuple[tuple[int, int], tuple[int, int]]] = [ + ((100, 200), (900, 400)), + ((500, 500), (500, 100)), + ((200, 600), (800, 200)), + ((100, 100), (130, 110)), + ((50, 50), (1500, 900)), + ((400, 300), (500, 300)), + ((300, 300), (700, 700)), + ((600, 400), (200, 100)), +] + +_SCROLL_CASES: list[tuple[int, int, str]] = [ + (0, 1500, "target"), + (0, 500, "precise"), + (0, 5000, "fast"), + (2000, 0, "target"), + (0, 800, "precise"), + (0, 3500, "fast"), + (1000, 1200, "precise"), + (0, 10000, "fast"), +] + +# Mouse path tolerance is distance-scaled because absolute pixel diff +# trivially grows with travel distance. Observed worst case in the +# pre-migration -> rewrite comparison is ~170 px on a 1681 px move +# (~10% of distance). 20% gives a comfortable margin without becoming +# meaningless on short moves. +_MOUSE_XY_REL = 0.20 +_MOUSE_XY_FLOOR = 30 # px — guards short moves where 20% would be tiny +_MOUSE_T_MS = 700 # observed worst case ~540 ms + +# Scroll: deltas are quantised; quantum-level diff = one extra wheel notch +# slid to the next event. Two quanta covers everything observed. +_SCROLL_DELTA_QUANTA = 2 +_SCROLL_T_MS = 700 + + +@pytest.mark.parametrize("case_idx", range(8)) +@pytest.mark.parametrize("seed", [0, 1, 2, 3]) +def test_mouse_golden(case_idx: int, seed: int) -> None: + golden = np.load(_GOLDEN_DIR / "golden_mouse.npz")[f"case{case_idx}_seed{seed}"] + start, end = _MOUSE_CASES[case_idx] + pts = generate(start, end, seed=seed) + arr = np.array(pts, dtype=np.int64) + + # Structural: shape must match exactly. + assert arr.shape == golden.shape, ( + f"shape mismatch: {arr.shape} vs {golden.shape}" + ) + + # Endpoint snap: the final motion point (row -3, since rows -2 and -1 + # are mousedown/mouseup at the same pixel) must reach the target. + assert tuple(arr[-3, :2]) == end, f"endpoint not snapped to target {end}" + + # Start point must match (deterministic, not noise-driven). + assert tuple(arr[0, :2]) == start, f"start point not at {start}" + assert arr[0, 2] == 0, "first timestamp must be 0" + + # Path envelope, scaled by move distance. + dist = math.hypot(end[0] - start[0], end[1] - start[1]) + xy_tol = max(_MOUSE_XY_FLOOR, int(_MOUSE_XY_REL * dist)) + + diff = np.abs(arr - golden) + xy_max = int(max(diff[:, 0].max(), diff[:, 1].max())) + t_max = int(diff[:, 2].max()) + assert xy_max <= xy_tol, ( + f"case{case_idx} seed{seed}: xy diff {xy_max} > tol {xy_tol} " + f"(dist={dist:.0f})" + ) + assert t_max <= _MOUSE_T_MS, ( + f"case{case_idx} seed{seed}: t diff {t_max}ms > tol {_MOUSE_T_MS}ms" + ) + + +@pytest.mark.parametrize("case_idx", range(8)) +@pytest.mark.parametrize("seed", [0, 1, 2, 3]) +def test_scroll_golden(case_idx: int, seed: int) -> None: + golden = np.load(_GOLDEN_DIR / "golden_scroll.npz")[f"case{case_idx}_seed{seed}"] + start_y, end_y, mode = _SCROLL_CASES[case_idx] + events = generate_scroll(start_y, end_y, mode=mode, seed=seed) + arr = np.array( + [[e["deltaY"], e["deltaMode"], e["t"]] for e in events], + dtype=np.int64, + ) + quantum = 40 if mode == "precise" else 120 + + if arr.shape != golden.shape: + pytest.skip( + f"event count diverged: {arr.shape[0]} vs {golden.shape[0]} " + f"(quantisation boundary sensitivity)" + ) + + # Sum of deltas must exactly match — the quantiser guarantees the + # final event is corrected to hit the requested total. + assert arr[:, 0].sum() == golden[:, 0].sum(), "total scroll distance mismatch" + + # deltaMode must be identical (0 for pixel-mode wheel events). + assert (arr[:, 1] == golden[:, 1]).all(), "deltaMode diverged" + + # Per-event delta and time within tolerance. + delta_diff = int(np.abs(arr[:, 0] - golden[:, 0]).max()) + t_diff = int(np.abs(arr[:, 2] - golden[:, 2]).max()) + delta_tol = _SCROLL_DELTA_QUANTA * quantum + assert delta_diff <= delta_tol, ( + f"case{case_idx} seed{seed} ({mode}): deltaY diff {delta_diff} > " + f"tol {delta_tol} ({_SCROLL_DELTA_QUANTA} quanta of {quantum})" + ) + assert t_diff <= _SCROLL_T_MS, ( + f"case{case_idx} seed{seed} ({mode}): t diff {t_diff}ms > " + f"tol {_SCROLL_T_MS}ms" + ) diff --git a/tests/unit/test_mouse.py b/tests/unit/test_mouse.py new file mode 100644 index 0000000..ff8a5af --- /dev/null +++ b/tests/unit/test_mouse.py @@ -0,0 +1,87 @@ +"""Tests for MouseModel and ai_mouse.generate().""" +from __future__ import annotations + +import pytest + + +def test_mouse_model_init_default() -> None: + from ai_mouse.mouse import MouseModel + m = MouseModel() + assert m._seq_len > 0 + assert m._session is not None + m.close() + + +def test_mouse_model_generate_returns_correct_shape() -> None: + from ai_mouse.mouse import MouseModel + m = MouseModel() + pts = m.generate((100, 200), (900, 400)) + assert len(pts) == 66 # 64 moves + 2 clicks + for x, y, t in pts: + assert isinstance(x, int) + assert isinstance(y, int) + assert isinstance(t, int) + + +def test_mouse_model_click_false_omits_clicks() -> None: + from ai_mouse.mouse import MouseModel + m = MouseModel() + pts = m.generate((100, 200), (900, 400), click=False) + assert len(pts) == 64 + + +def test_mouse_model_seed_reproducibility() -> None: + from ai_mouse.mouse import MouseModel + m = MouseModel() + a = m.generate((100, 200), (900, 400), seed=42) + b = m.generate((100, 200), (900, 400), seed=42) + assert a == b + + +def test_mouse_model_invalid_path_raises_model_load_error() -> None: + from ai_mouse.mouse import MouseModel + from ai_mouse.errors import ModelLoadError + with pytest.raises(ModelLoadError): + MouseModel(model_path="/nonexistent/path/here") + + +def test_mouse_model_timestamps_monotonic() -> None: + from ai_mouse.mouse import MouseModel + m = MouseModel() + pts = m.generate((100, 200), (900, 400)) + times = [p[2] for p in pts] + for i in range(1, len(times)): + assert times[i] >= times[i - 1] + + +def test_mouse_model_starts_and_ends_near_endpoints() -> None: + from ai_mouse.mouse import MouseModel + start = (100, 200) + end = (900, 400) + m = MouseModel() + pts = m.generate(start, end) + assert abs(pts[0][0] - start[0]) < 30 + assert abs(pts[0][1] - start[1]) < 30 + last_move = pts[-3] # last 2 are click events + assert abs(last_move[0] - end[0]) < 30 + assert abs(last_move[1] - end[1]) < 30 + + +def test_mouse_model_n_points_parameter() -> None: + from ai_mouse.mouse import MouseModel + m = MouseModel() + pts = m.generate((100, 200), (900, 400), n_points=32) + # 32 moves + 2 clicks + assert len(pts) == 34 + + +def test_mouse_model_click_events_have_matching_coords() -> None: + from ai_mouse.mouse import MouseModel + m = MouseModel() + pts = m.generate((100, 200), (900, 400)) + down, up = pts[-2], pts[-1] + assert down[0] == up[0] + assert down[1] == up[1] + assert up[2] > down[2] + # Within click_dist bounds 20..500 + assert 20 <= up[2] - down[2] <= 500 diff --git a/tests/unit/test_postprocess.py b/tests/unit/test_postprocess.py new file mode 100644 index 0000000..42bcd21 --- /dev/null +++ b/tests/unit/test_postprocess.py @@ -0,0 +1,143 @@ +"""Tests for trajectory post-processing primitives.""" +from __future__ import annotations + +import numpy as np + +from ai_mouse._postprocess import gaussian_smooth + + +def test_gaussian_smooth_preserves_endpoints() -> None: + x = np.array([1.0, 5.0, 3.0, 8.0, 2.0, 6.0, 4.0]) + result = gaussian_smooth(x, sigma=1.0) + assert result[0] == 1.0 + assert result[-1] == 4.0 + + +def test_gaussian_smooth_short_input_unchanged() -> None: + x = np.array([1.0, 2.0, 3.0]) + result = gaussian_smooth(x, sigma=1.0) + assert np.array_equal(result, x) + + +def test_gaussian_smooth_constant_unchanged() -> None: + x = np.full(20, 7.5) + result = gaussian_smooth(x, sigma=1.0) + assert np.allclose(result, x, atol=1e-6) + + +from ai_mouse._postprocess import snap_endpoints + + +def test_snap_endpoints_pins_first_and_last() -> None: + forward = np.linspace(0.1, 0.9, 16) + lateral = np.full(16, 0.5) + f, l = snap_endpoints(forward.copy(), lateral.copy(), seq_len=16) + assert f[0] == 0.0 + assert l[0] == 0.0 + assert f[-1] == 1.0 + assert l[-1] == 0.0 + + +def test_snap_endpoints_preserves_middle() -> None: + forward = np.linspace(0.0, 1.0, 16) + lateral = np.zeros(16) + f, _ = snap_endpoints(forward.copy(), lateral.copy(), seq_len=16, n_snap=4) + # Points before the last n_snap should be unchanged + assert np.allclose(f[1 : 16 - 4], forward[1 : 16 - 4], atol=1e-6) + + +from ai_mouse._postprocess import enforce_forward_monotonic, smooth_start + + +def test_smooth_start_dampens_lateral() -> None: + forward = np.linspace(0, 1, 16) + lateral = np.full(16, 1.0) + forward[0] = lateral[0] = 0.0 # invariant: snap already done + _, l = smooth_start(forward.copy(), lateral.copy(), n=4) + # Lateral at points 1-4 should be < original (dampened) + assert l[1] < 1.0 + assert l[4] < 1.0 + # Lateral at point 5+ unchanged + assert l[5] == 1.0 + + +def test_enforce_forward_monotonic_repairs_inversions() -> None: + f = np.array([0.0, 0.4, 0.3, 0.6, 0.5, 1.0]) + out = enforce_forward_monotonic(f.copy()) + assert np.all(np.diff(out) > 0), out + + +def test_enforce_forward_monotonic_clips_to_unit_interval() -> None: + f = np.array([-0.1, 0.5, 1.2]) + out = enforce_forward_monotonic(f.copy()) + assert out[0] == 0.0 + assert out[-1] == 1.0 + + +from ai_mouse._postprocess import build_timestamps, resample_arc + + +def test_resample_arc_identity_when_same_length() -> None: + pts = np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0], [3.0, 1.0]]) + out = resample_arc(pts, 4) + assert np.allclose(out, pts, atol=1e-6) + + +def test_resample_arc_changes_length() -> None: + pts = np.array([[float(i), 0.0] for i in range(10)]) + out = resample_arc(pts, 5) + assert out.shape == (5, 2) + # Endpoints preserved + assert np.allclose(out[0], pts[0]) + assert np.allclose(out[-1], pts[-1]) + + +def test_build_timestamps_strictly_increasing() -> None: + log_dt = np.array([0.0, 2.0, 2.5, 3.0, 2.0]) + ts = build_timestamps(log_dt, total_duration_ms=200.0) + assert ts[0] == 0 + assert np.all(np.diff(ts) >= 1) # at least 1 ms apart + + +def test_build_timestamps_total_close_to_target() -> None: + log_dt = np.array([1.0] * 10) + ts = build_timestamps(log_dt, total_duration_ms=300.0) + # Last timestamp should be roughly total - one slot + assert abs(ts[-1] - 270) < 60 # tolerant of clipping + + +from ai_mouse._postprocess import sample_duration, truncnorm_sample + + +def test_truncnorm_sample_within_bounds() -> None: + rng = np.random.default_rng(0) + samples = [ + truncnorm_sample(80.0, 30.0, 20.0, 300.0, rng) for _ in range(500) + ] + arr = np.array(samples) + assert arr.min() >= 20.0 + assert arr.max() <= 300.0 + # Mean roughly close to mu + assert abs(arr.mean() - 80.0) < 5.0 + + +def test_truncnorm_sample_far_outside_falls_back_to_clip() -> None: + rng = np.random.default_rng(0) + # mu far outside [low, high] — rejection will fail + v = truncnorm_sample(mu=1000.0, sigma=1.0, low=20.0, high=30.0, rng=rng) + assert 20.0 <= v <= 30.0 + + +def test_sample_duration_uses_correct_bin() -> None: + dist_dict = { + "bins": [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")], + "params": [ + {"mu_log": 4.0, "sigma_log": 0.01}, # bin 0: dist < 50 + {"mu_log": 5.0, "sigma_log": 0.01}, # bin 1: 50 <= dist < 100 + {"mu_log": 6.0, "sigma_log": 0.01}, # bin 2: 100 <= dist < 200 + ] + [{"mu_log": 7.0, "sigma_log": 0.01}] * 5, + } + rng = np.random.default_rng(0) + v = sample_duration(dist_dict, 150.0, rng) + # exp(6) ~ 403, with tiny sigma we should land near there + assert 350 < v < 460 diff --git a/tests/unit/test_public_api.py b/tests/unit/test_public_api.py new file mode 100644 index 0000000..a6a7f19 --- /dev/null +++ b/tests/unit/test_public_api.py @@ -0,0 +1,43 @@ +"""Tests for the public package-level API.""" +from __future__ import annotations + + +def test_public_symbols_importable() -> None: + from ai_mouse import ( + MouseModel, + ScrollModel, + generate, + generate_scroll, + errors, + ) + assert MouseModel is not None + assert ScrollModel is not None + assert callable(generate) + assert callable(generate_scroll) + assert hasattr(errors, "ModelLoadError") + + +def test_generate_function_returns_list_of_tuples() -> None: + from ai_mouse import generate + pts = generate((100, 100), (300, 200)) + assert isinstance(pts, list) + assert len(pts) > 0 + assert isinstance(pts[0], tuple) + assert len(pts[0]) == 3 + + +def test_generate_singleton_reused() -> None: + from ai_mouse import generate + from ai_mouse import _model_cache + _model_cache._get_mouse_model.cache_clear() + generate((0, 0), (100, 100)) + info_after_first = _model_cache._get_mouse_model.cache_info() + generate((0, 0), (200, 200)) + info_after_second = _model_cache._get_mouse_model.cache_info() + assert info_after_second.hits > info_after_first.hits + + +def test_version_present() -> None: + import ai_mouse + assert hasattr(ai_mouse, "__version__") + assert isinstance(ai_mouse.__version__, str) diff --git a/tests/unit/test_scroll.py b/tests/unit/test_scroll.py new file mode 100644 index 0000000..41ed84b --- /dev/null +++ b/tests/unit/test_scroll.py @@ -0,0 +1,53 @@ +"""Tests for ScrollModel and ai_mouse.generate_scroll().""" +from __future__ import annotations + +import pytest + + +def test_scroll_model_init_default() -> None: + from ai_mouse.scroll import ScrollModel + m = ScrollModel() + assert m._seq_len > 0 + m.close() + + +def test_scroll_model_generate_target_mode() -> None: + from ai_mouse.scroll import ScrollModel + m = ScrollModel() + events = m.generate(0, 1500, mode="target") + assert len(events) >= 5 + total = sum(e["deltaY"] for e in events) + assert 1000 <= total <= 2000 # broad — quantisation can drift + assert events[0]["t"] == 0 + assert all(e["deltaMode"] == 0 for e in events) + + +def test_scroll_model_direction() -> None: + from ai_mouse.scroll import ScrollModel + m = ScrollModel() + events = m.generate(2000, 0, mode="target") + total = sum(e["deltaY"] for e in events) + assert total < 0 + + +def test_scroll_invalid_path() -> None: + from ai_mouse.errors import ModelLoadError + from ai_mouse.scroll import ScrollModel + with pytest.raises(ModelLoadError): + ScrollModel(model_path="/no/such/path") + + +def test_scroll_model_timestamps_monotonic() -> None: + from ai_mouse.scroll import ScrollModel + m = ScrollModel() + events = m.generate(0, 2000, mode="target") + times = [e["t"] for e in events] + for i in range(1, len(times)): + assert times[i] >= times[i - 1] + + +def test_scroll_model_deltaY_are_integers() -> None: + from ai_mouse.scroll import ScrollModel + m = ScrollModel() + events = m.generate(0, 2000, mode="target") + assert all(isinstance(e["deltaY"], int) for e in events) diff --git a/tools/__init__.py b/tools/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ai_mouse/__main__.py b/tools/__main__.py similarity index 95% rename from ai_mouse/__main__.py rename to tools/__main__.py index a60eb5d..030718e 100644 --- a/ai_mouse/__main__.py +++ b/tools/__main__.py @@ -12,7 +12,7 @@ from pathlib import Path def _train_main(args: argparse.Namespace) -> int: - from ai_mouse.trainer import train + from tools.trainer import train logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") train( @@ -28,7 +28,7 @@ def _train_main(args: argparse.Namespace) -> int: def _eval_main(args: argparse.Namespace) -> int: - from ai_mouse.eval.__main__ import main as eval_main + from tools.eval.__main__ import main as eval_main # Reconstruct argv for the sub-CLI argv = [ "--model-dir", args.model_dir, @@ -43,7 +43,7 @@ def _eval_main(args: argparse.Namespace) -> int: def _balabit_main(args: argparse.Namespace) -> int: - from ai_mouse.data_adapters.balabit import main as bal_main + from tools.data_adapters.balabit import main as bal_main argv = [ "--input", str(args.input), "--output", str(args.output), diff --git a/ai_mouse/collector.py b/tools/collector.py similarity index 100% rename from ai_mouse/collector.py rename to tools/collector.py diff --git a/ai_mouse/config.py b/tools/config.py similarity index 100% rename from ai_mouse/config.py rename to tools/config.py diff --git a/ai_mouse/data_adapters/__init__.py b/tools/data_adapters/__init__.py similarity index 100% rename from ai_mouse/data_adapters/__init__.py rename to tools/data_adapters/__init__.py diff --git a/ai_mouse/data_adapters/__main__.py b/tools/data_adapters/__main__.py similarity index 86% rename from ai_mouse/data_adapters/__main__.py rename to tools/data_adapters/__main__.py index bb368bd..db63e3e 100644 --- a/ai_mouse/data_adapters/__main__.py +++ b/tools/data_adapters/__main__.py @@ -7,7 +7,7 @@ from __future__ import annotations import sys -from ai_mouse.data_adapters.balabit import main +from tools.data_adapters.balabit import main if __name__ == "__main__": sys.exit(main()) diff --git a/ai_mouse/data_adapters/balabit.py b/tools/data_adapters/balabit.py similarity index 99% rename from ai_mouse/data_adapters/balabit.py rename to tools/data_adapters/balabit.py index 1771d6c..55d0fbf 100644 --- a/ai_mouse/data_adapters/balabit.py +++ b/tools/data_adapters/balabit.py @@ -260,7 +260,7 @@ def main(argv: list[str] | None = None) -> int: """CLI entry point: convert a directory of Balabit sessions to one JSONL file.""" import argparse - from ai_mouse.config import BalabitAdapterConfig + from tools.config import BalabitAdapterConfig parser = argparse.ArgumentParser(description="Convert Balabit dataset to traces.jsonl format") parser.add_argument( diff --git a/ai_mouse/eval/__init__.py b/tools/eval/__init__.py similarity index 100% rename from ai_mouse/eval/__init__.py rename to tools/eval/__init__.py diff --git a/ai_mouse/eval/__main__.py b/tools/eval/__main__.py similarity index 91% rename from ai_mouse/eval/__main__.py rename to tools/eval/__main__.py index 27ff9cd..429b7e7 100644 --- a/ai_mouse/eval/__main__.py +++ b/tools/eval/__main__.py @@ -48,8 +48,15 @@ def _load_reference_jsonl(path: Path, n_samples: int) -> list[dict]: def _generate_n_samples( model_dir: str, n_samples: int, seed: int = 0 ) -> list[dict]: - """Call the project's generator N times with random start/end pairs.""" - from ai_mouse.generator import generate + """Call the project's generator N times with random start/end pairs. + + ``model_dir`` is accepted for CLI backward compatibility but is no longer + used — generation goes through the public ai_mouse API which loads the + bundled ONNX model. Export a fresh .onnx via ``python -m tools.export_onnx`` + to refresh. + """ + del model_dir # legacy arg, unused + from ai_mouse import generate rng = random.Random(seed) out: list[dict] = [] @@ -63,7 +70,7 @@ def _generate_n_samples( ex = max(0, min(800, ex)) ey = max(0, min(600, ey)) try: - pts = generate(start=(sx, sy), end=(ex, ey), model_dir=model_dir) + pts = generate(start=(sx, sy), end=(ex, ey)) except Exception as exc: # noqa: BLE001 logger.warning("generate() failed at i=%d: %s", i, exc) continue @@ -110,7 +117,7 @@ def main(argv: list[str] | None = None) -> int: logger.error("Empty trace sets — aborting") return 1 - from ai_mouse.eval.report import build_report + from tools.eval.report import build_report args.output.parent.mkdir(parents=True, exist_ok=True) build_report( generated_traces=gen_traces, diff --git a/ai_mouse/eval/metrics.py b/tools/eval/metrics.py similarity index 100% rename from ai_mouse/eval/metrics.py rename to tools/eval/metrics.py diff --git a/ai_mouse/eval/report.py b/tools/eval/report.py similarity index 99% rename from ai_mouse/eval/report.py rename to tools/eval/report.py index 4bbafa7..738d8ec 100644 --- a/ai_mouse/eval/report.py +++ b/tools/eval/report.py @@ -15,7 +15,7 @@ import matplotlib matplotlib.use("Agg") # headless import matplotlib.pyplot as plt # noqa: E402 -from ai_mouse.eval.metrics import ( +from tools.eval.metrics import ( compute_acceleration, compute_jerk, compute_speed, diff --git a/tools/export_onnx.py b/tools/export_onnx.py new file mode 100644 index 0000000..c013f59 --- /dev/null +++ b/tools/export_onnx.py @@ -0,0 +1,305 @@ +"""Export trained PyTorch checkpoints to ONNX for the inference SDK. + +Usage: + uv run python -m tools.export_onnx \ + --flow-ckpt data/models_v2 \ + --scroll-ckpt data/scroll_models \ + --output src/ai_mouse/assets/ + +Produces: + /flow_model.onnx + /scroll_decoder.onnx + /click_dist.json + /duration_dist.json + /train_config.json + /scroll_config.json + +A PyTorch vs ONNX Runtime parity check runs at the end. If parity fails +the .onnx files are deleted to prevent shipping broken weights. +""" +from __future__ import annotations + +import argparse +import json +import logging +import shutil +import sys +from pathlib import Path + +import numpy as np +import torch + +logger = logging.getLogger(__name__) + +_ATOL = 1e-4 + + +def export_flow_model(ckpt_dir: Path, out_dir: Path) -> Path: + """Export TrajectoryFlowModel to ONNX. + + Args: + ckpt_dir: directory with flow_model.pt and train_config.json. + out_dir: destination directory (created if missing). + + Returns: + Path to the written flow_model.onnx. + """ + from tools.models import TrajectoryFlowModel + + config_path = ckpt_dir / "train_config.json" + cfg = json.loads(config_path.read_text()) + seq_len = int(cfg["seq_len"]) + d_model = int(cfg["d_model"]) + nhead = int(cfg["nhead"]) + num_layers = int(cfg["num_layers"]) + dim_feedforward = int(cfg["dim_feedforward"]) + cond_dim = int(cfg.get("cond_dim", 3)) + + model = TrajectoryFlowModel( + seq_len=seq_len, + d_model=d_model, + nhead=nhead, + num_layers=num_layers, + dim_feedforward=dim_feedforward, + cond_dim=cond_dim, + dropout=0.0, + ) + state = torch.load(ckpt_dir / "flow_model.pt", map_location="cpu", weights_only=True) + model.load_state_dict(state) + model.eval() + + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / "flow_model.onnx" + + dummy_x = torch.zeros(1, seq_len, 3, dtype=torch.float32) + dummy_t = torch.zeros(1, dtype=torch.float32) + dummy_cond = torch.zeros(1, cond_dim, dtype=torch.float32) + + torch.onnx.export( + model, + (dummy_x, dummy_t, dummy_cond), + str(out_path), + input_names=["x_t", "t", "cond"], + output_names=["v"], + dynamic_axes={ + "x_t": {0: "batch"}, + "t": {0: "batch"}, + "cond": {0: "batch"}, + "v": {0: "batch"}, + }, + opset_version=17, + do_constant_folding=True, + dynamo=False, + ) + logger.info("Wrote %s (%.1f MB)", out_path, out_path.stat().st_size / 1e6) + return out_path + + +class _ScrollDecoder(torch.nn.Module): + """Wraps ScrollCVAE.decode for ONNX export. + + The full ScrollCVAE is encoder+decoder; inference only needs decoder. + """ + + def __init__(self, dec_h0, dec_gru, dec_out, seq_len: int, hidden: int): + super().__init__() + self.dec_h0 = dec_h0 + self.dec_gru = dec_gru + self.dec_out = dec_out + self.seq_len = seq_len + self.hidden = hidden + + def forward(self, z: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: + b = z.shape[0] + zc = torch.cat([z, cond], dim=-1) + h0_flat = self.dec_h0(zc) + h0 = h0_flat.view(b, 2, self.hidden).permute(1, 0, 2).contiguous() + inp = zc.unsqueeze(1).expand(b, self.seq_len, -1) + out, _ = self.dec_gru(inp, h0) + return self.dec_out(out) + + +def export_scroll_decoder(ckpt_dir: Path, out_dir: Path) -> Path: + """Export ScrollCVAE decoder to ONNX.""" + from tools.scroll.models import ScrollCVAE + + config_path = ckpt_dir / "scroll_config.json" + cfg = json.loads(config_path.read_text()) + seq_len = int(cfg["seq_len"]) + latent_dim = int(cfg["latent_dim"]) + hidden = int(cfg["hidden"]) + cond_dim = int(cfg["cond_dim"]) + + full = ScrollCVAE( + seq_len=seq_len, latent_dim=latent_dim, hidden=hidden, cond_dim=cond_dim + ) + state = torch.load(ckpt_dir / "scroll_model.pt", map_location="cpu", weights_only=True) + full.load_state_dict(state) + full.eval() + + decoder = _ScrollDecoder( + dec_h0=full.dec_h0, + dec_gru=full.dec_gru, + dec_out=full.dec_out, + seq_len=seq_len, + hidden=hidden, + ) + decoder.eval() + + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / "scroll_decoder.onnx" + + dummy_z = torch.zeros(1, latent_dim, dtype=torch.float32) + dummy_cond = torch.zeros(1, cond_dim, dtype=torch.float32) + + torch.onnx.export( + decoder, + (dummy_z, dummy_cond), + str(out_path), + input_names=["z", "cond"], + output_names=["seq"], + dynamic_axes={ + "z": {0: "batch"}, + "cond": {0: "batch"}, + "seq": {0: "batch"}, + }, + opset_version=17, + do_constant_folding=True, + dynamo=False, + ) + logger.info("Wrote %s (%.1f KB)", out_path, out_path.stat().st_size / 1e3) + return out_path + + +def _check_flow_parity(ckpt_dir: Path, onnx_path: Path) -> None: + """Verify ONNX flow model matches PyTorch output on random input.""" + import onnxruntime as ort + from tools.models import TrajectoryFlowModel + + cfg = json.loads((ckpt_dir / "train_config.json").read_text()) + seq_len = int(cfg["seq_len"]) + cond_dim = int(cfg.get("cond_dim", 3)) + + model = TrajectoryFlowModel( + seq_len=seq_len, + d_model=int(cfg["d_model"]), + nhead=int(cfg["nhead"]), + num_layers=int(cfg["num_layers"]), + dim_feedforward=int(cfg["dim_feedforward"]), + cond_dim=cond_dim, + dropout=0.0, + ) + model.load_state_dict( + torch.load(ckpt_dir / "flow_model.pt", map_location="cpu", weights_only=True) + ) + model.eval() + + torch.manual_seed(42) + np.random.seed(42) + x = torch.randn(2, seq_len, 3, dtype=torch.float32) + t = torch.tensor([0.0, 0.5], dtype=torch.float32) + cond = torch.randn(2, cond_dim, dtype=torch.float32) + + with torch.no_grad(): + torch_out = model(x, t, cond).numpy() + + sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"]) + ort_out = sess.run( + ["v"], + { + "x_t": x.numpy(), + "t": t.numpy(), + "cond": cond.numpy(), + }, + )[0] + + if not np.allclose(torch_out, ort_out, atol=_ATOL, rtol=1e-3): + max_diff = float(np.abs(torch_out - ort_out).max()) + raise RuntimeError( + f"Flow model ORT/PyTorch parity FAILED: max abs diff = {max_diff:.2e} > {_ATOL:.2e}" + ) + logger.info("Flow model parity OK (atol=%.0e)", _ATOL) + + +def _check_scroll_parity(ckpt_dir: Path, onnx_path: Path) -> None: + """Verify ONNX scroll decoder matches PyTorch decoder output.""" + import onnxruntime as ort + from tools.scroll.models import ScrollCVAE + + cfg = json.loads((ckpt_dir / "scroll_config.json").read_text()) + seq_len = int(cfg["seq_len"]) + latent_dim = int(cfg["latent_dim"]) + cond_dim = int(cfg["cond_dim"]) + + full = ScrollCVAE( + seq_len=seq_len, + latent_dim=latent_dim, + hidden=int(cfg["hidden"]), + cond_dim=cond_dim, + ) + full.load_state_dict( + torch.load(ckpt_dir / "scroll_model.pt", map_location="cpu", weights_only=True) + ) + full.eval() + + torch.manual_seed(7) + z = torch.randn(2, latent_dim, dtype=torch.float32) + cond = torch.randn(2, cond_dim, dtype=torch.float32) + + with torch.no_grad(): + torch_out = full.decode(z, cond).numpy() + + sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"]) + ort_out = sess.run(["seq"], {"z": z.numpy(), "cond": cond.numpy()})[0] + + if not np.allclose(torch_out, ort_out, atol=_ATOL, rtol=1e-3): + max_diff = float(np.abs(torch_out - ort_out).max()) + raise RuntimeError( + f"Scroll decoder parity FAILED: max abs diff = {max_diff:.2e} > {_ATOL:.2e}" + ) + logger.info("Scroll decoder parity OK (atol=%.0e)", _ATOL) + + +def _copy_metadata(flow_dir: Path, scroll_dir: Path, out_dir: Path) -> None: + """Copy JSON metadata files alongside the ONNX models.""" + for name in ("click_dist.json", "duration_dist.json", "train_config.json"): + src = flow_dir / name + if not src.exists(): + raise FileNotFoundError(f"Required metadata missing: {src}") + shutil.copy2(src, out_dir / name) + src = scroll_dir / "scroll_config.json" + if not src.exists(): + raise FileNotFoundError(f"Required metadata missing: {src}") + shutil.copy2(src, out_dir / "scroll_config.json") + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(prog="export_onnx", description=__doc__.splitlines()[0]) + p.add_argument("--flow-ckpt", type=Path, required=True) + p.add_argument("--scroll-ckpt", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + args = p.parse_args(argv) + + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + + args.output.mkdir(parents=True, exist_ok=True) + + flow_onnx = export_flow_model(args.flow_ckpt, args.output) + scroll_onnx = export_scroll_decoder(args.scroll_ckpt, args.output) + + try: + _check_flow_parity(args.flow_ckpt, flow_onnx) + _check_scroll_parity(args.scroll_ckpt, scroll_onnx) + except RuntimeError as exc: + logger.error("Parity check failed: %s", exc) + flow_onnx.unlink(missing_ok=True) + scroll_onnx.unlink(missing_ok=True) + return 1 + + _copy_metadata(args.flow_ckpt, args.scroll_ckpt, args.output) + logger.info("Export complete: %s", args.output) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/ai_mouse/models.py b/tools/models.py similarity index 65% rename from ai_mouse/models.py rename to tools/models.py index 2d252bb..e11deec 100644 --- a/ai_mouse/models.py +++ b/tools/models.py @@ -12,7 +12,6 @@ import math import torch import torch.nn as nn -from torch.distributions import Normal # --------------------------------------------------------------------------- @@ -157,80 +156,3 @@ class TrajectoryFlowModel(nn.Module): # Output projection return self.output_proj(h) # (B, T, 3) - - -# --------------------------------------------------------------------------- -# Legacy JointCVAE — kept for backward compatibility with generator.py -# --------------------------------------------------------------------------- - - -class JointCVAE(nn.Module): - """Joint Conditional VAE for mouse trajectory generation (legacy). - - Kept for backward compatibility with the existing generator. - See TrajectoryFlowModel for the new approach. - """ - - def __init__( - self, - seq_len: int = 64, - latent_dim: int = 32, - hidden: int = 128, - cond_dim: int = 3, - ): - super().__init__() - self.seq_len = seq_len - self.latent_dim = latent_dim - self.hidden = hidden - self.cond_dim = cond_dim - self.feat_dim = 3 - - self.enc_gru = nn.GRU( - input_size=self.feat_dim + cond_dim, - hidden_size=hidden, - num_layers=2, - batch_first=True, - bidirectional=True, - ) - self.enc_mu = nn.Linear(hidden * 2, latent_dim) - self.enc_logvar = nn.Linear(hidden * 2, latent_dim) - - self.dec_h0 = nn.Linear(latent_dim + cond_dim, hidden * 2) - self.dec_gru = nn.GRU( - input_size=latent_dim + cond_dim, - hidden_size=hidden, - num_layers=2, - batch_first=True, - ) - self.dec_out = nn.Linear(hidden, self.feat_dim) - - def encode(self, seq: torch.Tensor, cond: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - B, T, _ = seq.shape - c_exp = cond.unsqueeze(1).expand(B, T, self.cond_dim) - x_in = torch.cat([seq, c_exp], dim=-1) - _, h_n = self.enc_gru(x_in) - h_fwd = h_n[-2] - h_bwd = h_n[-1] - h_cat = torch.cat([h_fwd, h_bwd], dim=-1) - return self.enc_mu(h_cat), self.enc_logvar(h_cat) - - def decode(self, z: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: - B = z.shape[0] - zc = torch.cat([z, cond], dim=-1) - h0_flat = self.dec_h0(zc) - h0 = h0_flat.view(B, 2, self.hidden).permute(1, 0, 2).contiguous() - inp = zc.unsqueeze(1).expand(B, self.seq_len, -1) - out, _ = self.dec_gru(inp, h0) - return self.dec_out(out) - - def reparameterise(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor: - std = torch.exp(0.5 * logvar) - return Normal(mu, std).rsample() - - def forward( - self, seq: torch.Tensor, cond: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - mu, logvar = self.encode(seq, cond) - z = self.reparameterise(mu, logvar) - recon = self.decode(z, cond) - return recon, mu, logvar diff --git a/tools/scroll/__init__.py b/tools/scroll/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ai_mouse/scroll/collector.py b/tools/scroll/collector.py similarity index 98% rename from ai_mouse/scroll/collector.py rename to tools/scroll/collector.py index 33041e4..39b1993 100644 --- a/ai_mouse/scroll/collector.py +++ b/tools/scroll/collector.py @@ -9,7 +9,7 @@ from __future__ import annotations import logging import random -from ai_mouse.config import SCROLL_MODES +from tools.config import SCROLL_MODES logger = logging.getLogger(__name__) diff --git a/ai_mouse/scroll/models.py b/tools/scroll/models.py similarity index 100% rename from ai_mouse/scroll/models.py rename to tools/scroll/models.py diff --git a/ai_mouse/scroll/trainer.py b/tools/scroll/trainer.py similarity index 98% rename from ai_mouse/scroll/trainer.py rename to tools/scroll/trainer.py index ffe3bab..594f48a 100644 --- a/ai_mouse/scroll/trainer.py +++ b/tools/scroll/trainer.py @@ -20,8 +20,8 @@ import torch.nn as nn from torch.distributions import Normal, kl_divergence from torch.utils.data import DataLoader, TensorDataset -from ai_mouse.config import ScrollTrainConfig -from ai_mouse.scroll.models import ScrollCVAE +from tools.config import ScrollTrainConfig +from tools.scroll.models import ScrollCVAE logger = logging.getLogger(__name__) diff --git a/main.py b/tools/serve.py similarity index 89% rename from main.py rename to tools/serve.py index 89ec491..e8b4086 100644 --- a/main.py +++ b/tools/serve.py @@ -10,7 +10,7 @@ import webbrowser import uvicorn -from ai_mouse.server import create_app +from tools.server import create_app app = create_app() diff --git a/ai_mouse/server/__init__.py b/tools/server/__init__.py similarity index 100% rename from ai_mouse/server/__init__.py rename to tools/server/__init__.py diff --git a/ai_mouse/server/deps.py b/tools/server/deps.py similarity index 100% rename from ai_mouse/server/deps.py rename to tools/server/deps.py diff --git a/ai_mouse/server/routes_collect.py b/tools/server/routes_collect.py similarity index 98% rename from ai_mouse/server/routes_collect.py rename to tools/server/routes_collect.py index 2eeb2a9..f54f2b5 100644 --- a/ai_mouse/server/routes_collect.py +++ b/tools/server/routes_collect.py @@ -8,7 +8,7 @@ import logging from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel -from ai_mouse.collector import Collector +from tools.collector import Collector from .deps import SessionState, get_data_dir, get_state diff --git a/ai_mouse/server/routes_scroll.py b/tools/server/routes_scroll.py similarity index 91% rename from ai_mouse/server/routes_scroll.py rename to tools/server/routes_scroll.py index d57e114..9e4760f 100644 --- a/ai_mouse/server/routes_scroll.py +++ b/tools/server/routes_scroll.py @@ -85,7 +85,7 @@ def scroll_start( req: ScrollStartRequest, state: SessionState = Depends(get_state), ) -> dict: - from ai_mouse.scroll.collector import ScrollCollector + from tools.scroll.collector import ScrollCollector scroll_collector = ScrollCollector( mode=req.mode, count=req.count, viewport_height=req.viewport_height @@ -146,7 +146,7 @@ async def _scroll_train_sse(req: ScrollTrainRequest) -> AsyncGenerator[str, None queue.put_nowait(msg) async def run() -> None: - from ai_mouse.scroll.trainer import train_scroll + from tools.scroll.trainer import train_scroll traces_path, models_dir = _paths() try: @@ -182,21 +182,17 @@ async def scroll_train(req: ScrollTrainRequest) -> StreamingResponse: @router.post("/verify") def scroll_verify(req: ScrollVerifyRequest) -> dict: - from ai_mouse.scroll.generator import generate_scroll + # Uses the bundled ONNX scroll model exposed via the public ai_mouse API. + # The legacy scroll_model.pt path is no longer wired in; export a fresh + # scroll_decoder.onnx via `python -m tools.export_onnx` to update. + from ai_mouse import generate_scroll - _, models_dir = _paths() - if not (models_dir / "scroll_model.pt").exists(): - raise HTTPException( - status_code=400, - detail="滚轮模型尚未训练，请先在「训练模型 → 滚轮模型」中完成训练。", - ) paths = [] for _ in range(min(req.n_paths, 12)): events = generate_scroll( req.start_scrollY, req.target_scrollY, mode=req.mode, - model_dir=str(models_dir), ) paths.append(events) return {"paths": paths} diff --git a/ai_mouse/server/routes_train.py b/tools/server/routes_train.py similarity index 98% rename from ai_mouse/server/routes_train.py rename to tools/server/routes_train.py index 2c4d511..f7566b7 100644 --- a/ai_mouse/server/routes_train.py +++ b/tools/server/routes_train.py @@ -77,7 +77,7 @@ async def _train_sse_generator(req: TrainRequest) -> AsyncGenerator[str, None]: queue.put_nowait(msg) async def run_training_async() -> None: - from ai_mouse.trainer import train + from tools.trainer import train traces_path, models_dir, pretrained_dir = _paths() data_path = Path(req.data_path) if req.data_path else traces_path diff --git a/ai_mouse/server/routes_verify.py b/tools/server/routes_verify.py similarity index 79% rename from ai_mouse/server/routes_verify.py rename to tools/server/routes_verify.py index 0430b0e..b17a1d4 100644 --- a/ai_mouse/server/routes_verify.py +++ b/tools/server/routes_verify.py @@ -7,8 +7,6 @@ import logging from fastapi import APIRouter, HTTPException from pydantic import BaseModel -from .deps import get_data_dir - logger = logging.getLogger(__name__) router = APIRouter() @@ -32,18 +30,19 @@ class VerifyRequest(BaseModel): @router.post("/verify") def verify(req: VerifyRequest) -> dict: - from ai_mouse.generator import generate + # Uses the bundled ONNX model exposed via the public ai_mouse API. + # The legacy req.model_dir / data/models_v2 .pt path is no longer wired + # in; export a fresh .onnx via `python -m tools.export_onnx` to update. + from ai_mouse import generate n = max(1, min(req.n_paths, 12)) - models_dir = get_data_dir() / "models_v2" - model_dir_arg = req.model_dir if req.model_dir else str(models_dir) start = tuple(req.start) # type: ignore[arg-type] end = tuple(req.end) # type: ignore[arg-type] paths = [] try: for _ in range(n): - pts = generate(start=start, end=end, model_dir=model_dir_arg) + pts = generate(start=start, end=end) paths.append([[x, y, t] for x, y, t in pts]) except FileNotFoundError as exc: raise HTTPException(status_code=404, detail=str(exc)) from exc diff --git a/ai_mouse/trainer.py b/tools/trainer.py similarity index 98% rename from ai_mouse/trainer.py rename to tools/trainer.py index cc5fdfd..47cbdd3 100644 --- a/ai_mouse/trainer.py +++ b/tools/trainer.py @@ -23,10 +23,10 @@ import numpy as np import torch from torch.utils.data import DataLoader -from ai_mouse.config import TrainConfig -from ai_mouse.coord import encode_trajectory -from ai_mouse.models import TrajectoryFlowModel -from ai_mouse.utils import resample_arc +from ai_mouse._coord import encode_trajectory +from tools.config import TrainConfig +from tools.models import TrajectoryFlowModel +from tools.utils import resample_arc logger = logging.getLogger(__name__) diff --git a/ai_mouse/utils.py b/tools/utils.py similarity index 100% rename from ai_mouse/utils.py rename to tools/utils.py diff --git a/uv.lock b/uv.lock index 641a642..ccec9bf 100644 --- a/uv.lock +++ b/uv.lock @@ -1,42 +1,54 @@ version = 1 revision = 3 requires-python = ">=3.12, <3.14" +resolution-markers = [ + "python_full_version >= '3.13' and platform_machine != 's390x'", + "python_full_version < '3.13' and platform_machine != 's390x'", + "python_full_version >= '3.13' and platform_machine == 's390x'", + "python_full_version < '3.13' and platform_machine == 's390x'", +] [[package]] name = "ai-mouse" -version = "0.1.0" -source = { virtual = "." } +version = "0.2.0" +source = { editable = "." } dependencies = [ - { name = "fastapi" }, - { name = "matplotlib" }, { name = "numpy" }, + { name = "onnxruntime" }, +] + +[package.dev-dependencies] +dev = [ + { name = "fastapi" }, + { name = "httpx" }, + { name = "matplotlib" }, + { name = "onnx" }, + { name = "onnxscript" }, + { name = "pytest" }, + { name = "pytest-asyncio" }, { name = "scipy" }, { name = "torch" }, { name = "uvicorn" }, ] -[package.dev-dependencies] -dev = [ - { name = "httpx" }, - { name = "pytest" }, - { name = "pytest-asyncio" }, -] - [package.metadata] requires-dist = [ - { name = "fastapi", specifier = ">=0.111.0" }, - { name = "matplotlib", specifier = ">=3.8.0" }, { name = "numpy", specifier = ">=1.26.0" }, - { name = "scipy", specifier = ">=1.10.0" }, - { name = "torch", specifier = ">=2.2.0" }, - { name = "uvicorn", specifier = ">=0.29.0" }, + { name = "onnxruntime", specifier = ">=1.17.0" }, ] [package.metadata.requires-dev] dev = [ + { name = "fastapi", specifier = ">=0.111.0" }, { name = "httpx", specifier = ">=0.27.0" }, + { name = "matplotlib", specifier = ">=3.8.0" }, + { name = "onnx", specifier = ">=1.15.0" }, + { name = "onnxscript", specifier = ">=0.1" }, { name = "pytest", specifier = ">=8.0.0" }, { name = "pytest-asyncio", specifier = ">=0.23.0" }, + { name = "scipy", specifier = ">=1.10.0" }, + { name = "torch", specifier = ">=2.2.0" }, + { name = "uvicorn", specifier = ">=0.29.0" }, ] [[package]] @@ -243,6 +255,14 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/81/47/dd9a212ef6e343a6857485ffe25bba537304f1913bdbed446a23f7f592e1/filelock-3.29.0-py3-none-any.whl", hash = "sha256:96f5f6344709aa1572bbf631c640e4ebeeb519e08da902c39a001882f30ac258", size = 39812, upload-time = "2026-04-19T15:39:08.752Z" }, ] +[[package]] +name = "flatbuffers" +version = "25.12.19" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e8/2d/d2a548598be01649e2d46231d151a6c56d10b964d94043a335ae56ea2d92/flatbuffers-25.12.19-py2.py3-none-any.whl", hash = "sha256:7634f50c427838bb021c2d66a3d1168e9d199b0607e6329399f04846d42e20b4", size = 26661, upload-time = "2025-12-19T23:16:13.622Z" }, +] + [[package]] name = "fonttools" version = "4.62.1" @@ -481,6 +501,32 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f4/38/ae27288e788c35a4250491422f3db7750366fc8c97d6f36fbdecfc1f5518/matplotlib-3.10.9-cp313-cp313t-win_arm64.whl", hash = "sha256:6b63d9c7c769b88ab81e10dc86e4e0607cf56817b9f9e6cf24b2a5f1693b8e38", size = 8188292, upload-time = "2026-04-24T00:13:15.546Z" }, ] +[[package]] +name = "ml-dtypes" +version = "0.5.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0e/4a/c27b42ed9b1c7d13d9ba8b6905dece787d6259152f2309338aed29b2447b/ml_dtypes-0.5.4.tar.gz", hash = "sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453", size = 692314, upload-time = "2025-11-17T22:32:31.031Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a8/b8/3c70881695e056f8a32f8b941126cf78775d9a4d7feba8abcb52cb7b04f2/ml_dtypes-0.5.4-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:a174837a64f5b16cab6f368171a1a03a27936b31699d167684073ff1c4237dac", size = 676927, upload-time = "2025-11-17T22:31:48.182Z" }, + { url = "https://files.pythonhosted.org/packages/54/0f/428ef6881782e5ebb7eca459689448c0394fa0a80bea3aa9262cba5445ea/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a7f7c643e8b1320fd958bf098aa7ecf70623a42ec5154e3be3be673f4c34d900", size = 5028464, upload-time = "2025-11-17T22:31:50.135Z" }, + { url = "https://files.pythonhosted.org/packages/3a/cb/28ce52eb94390dda42599c98ea0204d74799e4d8047a0eb559b6fd648056/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9ad459e99793fa6e13bd5b7e6792c8f9190b4e5a1b45c63aba14a4d0a7f1d5ff", size = 5009002, upload-time = "2025-11-17T22:31:52.001Z" }, + { url = "https://files.pythonhosted.org/packages/f5/f0/0cfadd537c5470378b1b32bd859cf2824972174b51b873c9d95cfd7475a5/ml_dtypes-0.5.4-cp312-cp312-win_amd64.whl", hash = "sha256:c1a953995cccb9e25a4ae19e34316671e4e2edaebe4cf538229b1fc7109087b7", size = 212222, upload-time = "2025-11-17T22:31:53.742Z" }, + { url = "https://files.pythonhosted.org/packages/16/2e/9acc86985bfad8f2c2d30291b27cd2bb4c74cea08695bd540906ed744249/ml_dtypes-0.5.4-cp312-cp312-win_arm64.whl", hash = "sha256:9bad06436568442575beb2d03389aa7456c690a5b05892c471215bfd8cf39460", size = 160793, upload-time = "2025-11-17T22:31:55.358Z" }, + { url = "https://files.pythonhosted.org/packages/d9/a1/4008f14bbc616cfb1ac5b39ea485f9c63031c4634ab3f4cf72e7541f816a/ml_dtypes-0.5.4-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:8c760d85a2f82e2bed75867079188c9d18dae2ee77c25a54d60e9cc79be1bc48", size = 676888, upload-time = "2025-11-17T22:31:56.907Z" }, + { url = "https://files.pythonhosted.org/packages/d3/b7/dff378afc2b0d5a7d6cd9d3209b60474d9819d1189d347521e1688a60a53/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ce756d3a10d0c4067172804c9cc276ba9cc0ff47af9078ad439b075d1abdc29b", size = 5036993, upload-time = "2025-11-17T22:31:58.497Z" }, + { url = "https://files.pythonhosted.org/packages/eb/33/40cd74219417e78b97c47802037cf2d87b91973e18bb968a7da48a96ea44/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:533ce891ba774eabf607172254f2e7260ba5f57bdd64030c9a4fcfbd99815d0d", size = 5010956, upload-time = "2025-11-17T22:31:59.931Z" }, + { url = "https://files.pythonhosted.org/packages/e1/8b/200088c6859d8221454825959df35b5244fa9bdf263fd0249ac5fb75e281/ml_dtypes-0.5.4-cp313-cp313-win_amd64.whl", hash = "sha256:f21c9219ef48ca5ee78402d5cc831bd58ea27ce89beda894428bc67a52da5328", size = 212224, upload-time = "2025-11-17T22:32:01.349Z" }, + { url = "https://files.pythonhosted.org/packages/8f/75/dfc3775cb36367816e678f69a7843f6f03bd4e2bcd79941e01ea960a068e/ml_dtypes-0.5.4-cp313-cp313-win_arm64.whl", hash = "sha256:35f29491a3e478407f7047b8a4834e4640a77d2737e0b294d049746507af5175", size = 160798, upload-time = "2025-11-17T22:32:02.864Z" }, + { url = "https://files.pythonhosted.org/packages/4f/74/e9ddb35fd1dd43b1106c20ced3f53c2e8e7fc7598c15638e9f80677f81d4/ml_dtypes-0.5.4-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:304ad47faa395415b9ccbcc06a0350800bc50eda70f0e45326796e27c62f18b6", size = 702083, upload-time = "2025-11-17T22:32:04.08Z" }, + { url = "https://files.pythonhosted.org/packages/74/f5/667060b0aed1aa63166b22897fdf16dca9eb704e6b4bbf86848d5a181aa7/ml_dtypes-0.5.4-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6a0df4223b514d799b8a1629c65ddc351b3efa833ccf7f8ea0cf654a61d1e35d", size = 5354111, upload-time = "2025-11-17T22:32:05.546Z" }, + { url = "https://files.pythonhosted.org/packages/40/49/0f8c498a28c0efa5f5c95a9e374c83ec1385ca41d0e85e7cf40e5d519a21/ml_dtypes-0.5.4-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:531eff30e4d368cb6255bc2328d070e35836aa4f282a0fb5f3a0cd7260257298", size = 5366453, upload-time = "2025-11-17T22:32:07.115Z" }, + { url = "https://files.pythonhosted.org/packages/8c/27/12607423d0a9c6bbbcc780ad19f1f6baa2b68b18ce4bddcdc122c4c68dc9/ml_dtypes-0.5.4-cp313-cp313t-win_amd64.whl", hash = "sha256:cb73dccfc991691c444acc8c0012bee8f2470da826a92e3a20bb333b1a7894e6", size = 225612, upload-time = "2025-11-17T22:32:08.615Z" }, + { url = "https://files.pythonhosted.org/packages/e5/80/5a5929e92c72936d5b19872c5fb8fc09327c1da67b3b68c6a13139e77e20/ml_dtypes-0.5.4-cp313-cp313t-win_arm64.whl", hash = "sha256:3bbbe120b915090d9dd1375e4684dd17a20a2491ef25d640a908281da85e73f1", size = 164145, upload-time = "2025-11-17T22:32:09.782Z" }, +] + [[package]] name = "mpmath" version = "1.3.0" @@ -688,6 +734,88 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a8/64/3708a90d1ebe202ffdeb7185f878a3c84d15c2b2c31858da2ce0583e2def/nvidia_nvtx-13.0.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6", size = 148878, upload-time = "2025-09-04T08:28:53.627Z" }, ] +[[package]] +name = "onnx" +version = "1.21.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "protobuf" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c5/93/942d2a0f6a70538eea042ce0445c8aefd46559ad153469986f29a743c01c/onnx-1.21.0.tar.gz", hash = "sha256:4d8b67d0aaec5864c87633188b91cc520877477ec0254eda122bef8be43cd764", size = 12074608, upload-time = "2026-03-27T21:33:36.118Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7d/ae/cb644ec84c25e63575d9d8790fdcc5d1a11d67d3f62f872edb35fa38d158/onnx-1.21.0-cp312-abi3-macosx_12_0_universal2.whl", hash = "sha256:fc2635400fe39ff37ebc4e75342cc54450eadadf39c540ff132c319bf4960095", size = 17965930, upload-time = "2026-03-27T21:32:48.089Z" }, + { url = "https://files.pythonhosted.org/packages/6f/b6/eeb5903586645ef8a49b4b7892580438741acc3df91d7a5bd0f3a59ea9cb/onnx-1.21.0-cp312-abi3-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9003d5206c01fa2ff4b46311566865d8e493e1a6998d4009ec6de39843f1b59b", size = 17531344, upload-time = "2026-03-27T21:32:50.837Z" }, + { url = "https://files.pythonhosted.org/packages/a7/00/4823f06357892d1e60d6f34e7299d2ba4ed2108c487cc394f7ce85a3ff14/onnx-1.21.0-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a9261bd580fb8548c9c37b3c6750387eb8f21ea43c63880d37b2c622e1684285", size = 17613697, upload-time = "2026-03-27T21:32:54.222Z" }, + { url = "https://files.pythonhosted.org/packages/23/1d/391f3c567ae068c8ac4f1d1316bae97c9eb45e702f05975fe0e17ad441f0/onnx-1.21.0-cp312-abi3-win32.whl", hash = "sha256:9ea4e824964082811938a9250451d89c4ec474fe42dd36c038bfa5df31993d1e", size = 16287200, upload-time = "2026-03-27T21:32:57.277Z" }, + { url = "https://files.pythonhosted.org/packages/9c/a6/5eefbe5b40ea96de95a766bd2e0e751f35bdea2d4b951991ec9afaa69531/onnx-1.21.0-cp312-abi3-win_amd64.whl", hash = "sha256:458d91948ad9a7729a347550553b49ab6939f9af2cddf334e2116e45467dc61f", size = 16441045, upload-time = "2026-03-27T21:33:00.081Z" }, + { url = "https://files.pythonhosted.org/packages/63/c4/0ed8dc037a39113d2a4d66e0005e07751c299c46b993f1ad5c2c35664c20/onnx-1.21.0-cp312-abi3-win_arm64.whl", hash = "sha256:ca14bc4842fccc3187eb538f07eabeb25a779b39388b006db4356c07403a7bbb", size = 16403134, upload-time = "2026-03-27T21:33:03.987Z" }, + { url = "https://files.pythonhosted.org/packages/f8/89/0e1a9beb536401e2f45ac88735e123f2735e12fc7b56ff6c11727e097526/onnx-1.21.0-cp313-cp313t-macosx_12_0_universal2.whl", hash = "sha256:257d1d1deb6a652913698f1e3f33ef1ca0aa69174892fe38946d4572d89dd94f", size = 17975430, upload-time = "2026-03-27T21:33:07.005Z" }, + { url = "https://files.pythonhosted.org/packages/ec/46/e6dc71a7b3b317265591b20a5f71d0ff5c0d26c24e52283139dc90c66038/onnx-1.21.0-cp313-cp313t-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7cd7cb8f6459311bdb557cbf6c0ccc6d8ace11c304d1bba0a30b4a4688e245f8", size = 17537435, upload-time = "2026-03-27T21:33:09.765Z" }, + { url = "https://files.pythonhosted.org/packages/49/2e/27affcac63eaf2ef183a44fd1a1354b11da64a6c72fe6f3fdcf5571bcee5/onnx-1.21.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7b58a4cfec8d9311b73dc083e4c1fa362069267881144c05139b3eba5dc3a840", size = 17617687, upload-time = "2026-03-27T21:33:12.619Z" }, + { url = "https://files.pythonhosted.org/packages/1c/5c/ac8ed15e941593a3672ce424280b764979026317811f2e8508432bfc3429/onnx-1.21.0-cp313-cp313t-win_amd64.whl", hash = "sha256:1a9baf882562c4cebf79589bebb7cd71a20e30b51158cac3e3bbaf27da6163bd", size = 16449402, upload-time = "2026-03-27T21:33:15.555Z" }, + { url = "https://files.pythonhosted.org/packages/0e/aa/d2231e0dcaad838217afc64c306c8152a080134d2034e247cc973d577674/onnx-1.21.0-cp313-cp313t-win_arm64.whl", hash = "sha256:bba12181566acf49b35875838eba49536a327b2944664b17125577d230c637ad", size = 16408273, upload-time = "2026-03-27T21:33:18.599Z" }, +] + +[[package]] +name = "onnx-ir" +version = "0.2.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "onnx" }, + { name = "sympy" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/35/e6/672fefb2f108d077f58181a7babf4c0f8d1182a30353ffc9c79c63afc5ee/onnx_ir-0.2.1.tar.gz", hash = "sha256:8b8b10a93f43e65962104de6070c43c5dacb0e3cdfefc7c8059dd83c9db64f35", size = 144279, upload-time = "2026-04-20T20:21:47.735Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8c/aa/f7a53321c60b9ad9ee184b6018292ed6b5389947592a2c8c09c736bb7f9e/onnx_ir-0.2.1-py3-none-any.whl", hash = "sha256:c7285da889312f91882de2092e298a9eeeefbfc1d1951c49d983992967eb09a7", size = 166792, upload-time = "2026-04-20T20:21:46.357Z" }, +] + +[[package]] +name = "onnxruntime" +version = "1.26.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "flatbuffers" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "protobuf" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/f6/a0/3f9d896a0385a36bd04345d6d0b802821a5782adde562e7e135f6bb71c73/onnxruntime-1.26.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:91f2bb870a4b9224eba0a6728c1fa7a9e552b8e59e1083c51fbbc3d013f2b5c0", size = 16052692, upload-time = "2026-05-08T19:07:13.829Z" }, + { url = "https://files.pythonhosted.org/packages/7c/43/2a4e04f8dbeffad19bbcced4bcd4289bf478921518437404d6b92bdf213b/onnxruntime-1.26.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9b6dd70599005bd1bf29779f04a91978b92b5e719c11a20068a8f8e535f725b6", size = 18185439, upload-time = "2026-05-08T19:07:36.299Z" }, + { url = "https://files.pythonhosted.org/packages/44/fc/026d0a7162b9c2153dac292baea9e027c42304dc1d9dc6f8ff5b4cfbaedd/onnxruntime-1.26.0-cp312-cp312-win_amd64.whl", hash = "sha256:a26374dc7fbcaae593601086b242120e13f2310558df0991da6dd8b8fac00414", size = 13026427, upload-time = "2026-05-08T19:08:03.503Z" }, + { url = "https://files.pythonhosted.org/packages/3e/27/1dcf88e45e4c69db5f7b106f2dacc3801ba98994e082ca03e1dfdf7bfe57/onnxruntime-1.26.0-cp312-cp312-win_arm64.whl", hash = "sha256:54a8053410fd31fd66469bd754fcfe8a4df9f7eb44756b4b5479bf50c842d948", size = 12796647, upload-time = "2026-05-08T19:07:52.108Z" }, + { url = "https://files.pythonhosted.org/packages/cf/a2/c801242685e0ce48a4ca51dfafbb588765e0446397e123be53ba5598f3f5/onnxruntime-1.26.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:ccce19c5f771b8268902f77d9fed9e88f9499465d6780808faa6611a789d33f0", size = 18016563, upload-time = "2026-05-08T19:07:28.081Z" }, + { url = "https://files.pythonhosted.org/packages/e2/64/0492c0b1db04e29b2630c87cfa36f9d6872b1ca8614b90c5cad58fac7d76/onnxruntime-1.26.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bdbed8cf3b672b66acb032f33a253bc27f42bce6ece48ae3fab4fa483a5e96e0", size = 16052634, upload-time = "2026-05-08T19:07:16.885Z" }, + { url = "https://files.pythonhosted.org/packages/3d/26/4d09ddc755a84fc8d5e192991626b0e0680e8f6c5d58f4f1d05c42bc48cf/onnxruntime-1.26.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c07af6fc6d5557835f2b6ee7a96d8b3235d0c57a8e230efdedaee106a8a3cbc6", size = 18185632, upload-time = "2026-05-08T19:07:38.756Z" }, + { url = "https://files.pythonhosted.org/packages/77/89/3e52249aa08fa301e217ecba07b5246a8338fa2b401e109326e3fc5be0f9/onnxruntime-1.26.0-cp313-cp313-win_amd64.whl", hash = "sha256:61bec80655efa460591c2bc655392d57d2650ce85533a6b9b3b7a790d7ea7916", size = 13026751, upload-time = "2026-05-08T19:08:06.2Z" }, + { url = "https://files.pythonhosted.org/packages/06/b3/c1c8782b14af6797c303de132d6eef26a9fb80dfacd3750ce57911d11c6b/onnxruntime-1.26.0-cp313-cp313-win_arm64.whl", hash = "sha256:a6677545ff451e3539a02746d2f207d8c5baa4a0a818886bb9d6a6eb9511ee89", size = 12796807, upload-time = "2026-05-08T19:07:54.879Z" }, + { url = "https://files.pythonhosted.org/packages/c3/f5/47b0676408abec652c14b84d7173e389837832d850c24f87184277313e8d/onnxruntime-1.26.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5e016edc15d3c19f36807e1c6b10be5b27807688c32720f91b5ae480a95215d0", size = 16057265, upload-time = "2026-05-08T19:07:19.603Z" }, + { url = "https://files.pythonhosted.org/packages/3b/45/33ab6deeef010ca844c877dd618cebc079590bbe52d2a3678e7223b1b908/onnxruntime-1.26.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f5fc48a91a046a6a5c9b147f83fb41d65d24d24923373b222cdd248f0f4f4aac", size = 18197590, upload-time = "2026-05-08T19:07:41.422Z" }, +] + +[[package]] +name = "onnxscript" +version = "0.7.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "onnx" }, + { name = "onnx-ir" }, + { name = "packaging" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9b/99/fd948eba63ba65b52265a4cd09a14f96bb9f5b730fcef58876c4358bf406/onnxscript-0.7.0.tar.gz", hash = "sha256:c95ed7b339b02cface56ee27689565c46612e1fc542c562298dddfdad5268dc5", size = 612032, upload-time = "2026-04-20T17:09:19.775Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b9/ce/2ed92575cc3be4ea1db5f38f16f20765f9b20b69b14d6c1d9972658a8ee9/onnxscript-0.7.0-py3-none-any.whl", hash = "sha256:5b356907d4501e9919f8599c91d8da967406a37b1fac2b40caa55a49acf242ea", size = 714842, upload-time = "2026-04-20T17:09:22.089Z" }, +] + [[package]] name = "packaging" version = "26.2" @@ -750,6 +878,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, ] +[[package]] +name = "protobuf" +version = "7.34.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6b/6b/a0e95cad1ad7cc3f2c6821fcab91671bd5b78bd42afb357bb4765f29bc41/protobuf-7.34.1.tar.gz", hash = "sha256:9ce42245e704cc5027be797c1db1eb93184d44d1cdd71811fb2d9b25ad541280", size = 454708, upload-time = "2026-03-20T17:34:47.036Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/11/3325d41e6ee15bf1125654301211247b042563bcc898784351252549a8ad/protobuf-7.34.1-cp310-abi3-macosx_10_9_universal2.whl", hash = "sha256:d8b2cc79c4d8f62b293ad9b11ec3aebce9af481fa73e64556969f7345ebf9fc7", size = 429247, upload-time = "2026-03-20T17:34:37.024Z" }, + { url = "https://files.pythonhosted.org/packages/eb/9d/aa69df2724ff63efa6f72307b483ce0827f4347cc6d6df24b59e26659fef/protobuf-7.34.1-cp310-abi3-manylinux2014_aarch64.whl", hash = "sha256:5185e0e948d07abe94bb76ec9b8416b604cfe5da6f871d67aad30cbf24c3110b", size = 325753, upload-time = "2026-03-20T17:34:38.751Z" }, + { url = "https://files.pythonhosted.org/packages/92/e8/d174c91fd48e50101943f042b09af9029064810b734e4160bbe282fa1caa/protobuf-7.34.1-cp310-abi3-manylinux2014_s390x.whl", hash = "sha256:403b093a6e28a960372b44e5eb081775c9b056e816a8029c61231743d63f881a", size = 340198, upload-time = "2026-03-20T17:34:39.871Z" }, + { url = "https://files.pythonhosted.org/packages/53/1b/3b431694a4dc6d37b9f653f0c64b0a0d9ec074ee810710c0c3da21d67ba7/protobuf-7.34.1-cp310-abi3-manylinux2014_x86_64.whl", hash = "sha256:8ff40ce8cd688f7265326b38d5a1bed9bfdf5e6723d49961432f83e21d5713e4", size = 324267, upload-time = "2026-03-20T17:34:41.1Z" }, + { url = "https://files.pythonhosted.org/packages/85/29/64de04a0ac142fb685fd09999bc3d337943fb386f3a0ec57f92fd8203f97/protobuf-7.34.1-cp310-abi3-win32.whl", hash = "sha256:34b84ce27680df7cca9f231043ada0daa55d0c44a2ddfaa58ec1d0d89d8bf60a", size = 426628, upload-time = "2026-03-20T17:34:42.536Z" }, + { url = "https://files.pythonhosted.org/packages/4d/87/cb5e585192a22b8bd457df5a2c16a75ea0db9674c3a0a39fc9347d84e075/protobuf-7.34.1-cp310-abi3-win_amd64.whl", hash = "sha256:e97b55646e6ce5cbb0954a8c28cd39a5869b59090dfaa7df4598a7fba869468c", size = 437901, upload-time = "2026-03-20T17:34:44.112Z" }, + { url = "https://files.pythonhosted.org/packages/88/95/608f665226bca68b736b79e457fded9a2a38c4f4379a4a7614303d9db3bc/protobuf-7.34.1-py3-none-any.whl", hash = "sha256:bb3812cd53aefea2b028ef42bd780f5b96407247f20c6ef7c679807e9d188f11", size = 170715, upload-time = "2026-03-20T17:34:45.384Z" }, +] + [[package]] name = "pydantic" version = "2.13.4"