docs: design spec for ai_mouse library refactor (ONNX + src-layout)
Captures the decisions made in brainstorming for 0.2.0: - Split src/ai_mouse/ (pure-inference, numpy+onnxruntime only) from tools/ (training/server/eval, torch+fastapi+...) - Bundled ONNX weights via importlib.resources - Public API: MouseModel/ScrollModel classes + cached generate() helpers - ONNX export script with PyTorch parity check - Golden tests to lock semantics during NumPy rewrite of post-processing - 5-stage migration plan, git URL install, no PyPI Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
593
docs/superpowers/specs/2026-05-11-ai-mouse-library-design.md
Normal file
593
docs/superpowers/specs/2026-05-11-ai-mouse-library-design.md
Normal file
@@ -0,0 +1,593 @@
|
||||
# ai_mouse 库化与 ONNX Runtime 迁移 — 设计文档
|
||||
|
||||
**日期**: 2026-05-11
|
||||
**状态**: Draft (待用户审阅)
|
||||
**目标版本**: ai_mouse 0.2.0
|
||||
|
||||
## 背景与目标
|
||||
|
||||
当前 `ai_mouse` 仓库既是应用又是包:推理代码、训练代码、Web UI、采集、评估、数据适配器都混在 `ai_mouse/` 包目录下,运行时强依赖 `torch`、`fastapi`、`scipy`、`matplotlib`。这与"被其他项目直接引用的库"定位不符——任何只想调用 `generate()` 的下游必须装一整套 GB 级依赖。
|
||||
|
||||
本次重构的目标:
|
||||
|
||||
1. **物理边界**: 推理库代码与开发期代码完全分开,推理库不依赖 torch
|
||||
2. **ONNX 化**: 推理走 ONNX Runtime,预训练权重以 ONNX 格式打进 wheel
|
||||
3. **可消费**: 下游 `pip install git+...` 即可 `from ai_mouse import generate`
|
||||
4. **API 优化**: 同时提供函数式入口(向后兼容)和类入口(性能场景),底层共享 ORT session
|
||||
|
||||
非目标(显式排除):
|
||||
|
||||
- 不发布到 PyPI(本阶段仅 git URL 安装)
|
||||
- 不重写训练代码以脱离 torch
|
||||
- 不做 ONNX Runtime Training
|
||||
- 不为视觉真实性设硬阈值断言
|
||||
|
||||
## 约束(由用户决策确认)
|
||||
|
||||
| 维度 | 决策 |
|
||||
|---|---|
|
||||
| 使用场景 | 纯推理 SDK,零 torch 运行时依赖 |
|
||||
| 权重交付 | 打进 wheel(`importlib.resources`) |
|
||||
| 开发期代码位置 | 仓库顶层 `tools/`,不进 wheel |
|
||||
| 公开 API 形态 | 类 + 函数双轨;函数内部走进程级 lru_cache 单例 |
|
||||
| 发布渠道 | git URL 直装,不发 PyPI |
|
||||
|
||||
## 设计
|
||||
|
||||
### §1 公开 API
|
||||
|
||||
#### 顶层入口 (`src/ai_mouse/__init__.py`)
|
||||
|
||||
```python
|
||||
from ai_mouse import (
|
||||
MouseModel, # 类: 常驻 InferenceSession
|
||||
ScrollModel,
|
||||
generate, # 便利函数: 进程级单例缓存
|
||||
generate_scroll,
|
||||
)
|
||||
from ai_mouse import errors # ModelLoadError, GenerationError
|
||||
__version__ = "0.2.0"
|
||||
```
|
||||
|
||||
#### `MouseModel`
|
||||
|
||||
```python
|
||||
class MouseModel:
|
||||
def __init__(
|
||||
self,
|
||||
model_path: str | Path | None = None, # None ⇒ wheel 内置权重
|
||||
providers: Sequence[str] | None = None, # None ⇒ ORT 默认顺序 (CPU)
|
||||
seed: int | None = None, # None ⇒ 每次调用随机
|
||||
) -> None: ...
|
||||
|
||||
def generate(
|
||||
self,
|
||||
start: tuple[int, int],
|
||||
end: tuple[int, int],
|
||||
n_points: int = 64,
|
||||
speed: float | None = None,
|
||||
click: bool = True, # True ⇒ 末尾追加 mouse-down/-up
|
||||
seed: int | None = None, # 覆盖实例 seed
|
||||
) -> list[tuple[int, int, int]]: ...
|
||||
|
||||
def sample_click_duration_ms(self, seed: int | None = None) -> int: ...
|
||||
|
||||
def close(self) -> None: ...
|
||||
def __enter__(self) -> "MouseModel": ...
|
||||
def __exit__(self, *exc) -> None: ...
|
||||
```
|
||||
|
||||
#### `ScrollModel`(对称)
|
||||
|
||||
```python
|
||||
class ScrollModel:
|
||||
def __init__(self, model_path=None, providers=None, seed=None): ...
|
||||
|
||||
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]: ...
|
||||
|
||||
def close(self) -> None: ...
|
||||
def __enter__(self) -> "ScrollModel": ...
|
||||
def __exit__(self, *exc) -> None: ...
|
||||
```
|
||||
|
||||
#### 便利函数
|
||||
|
||||
```python
|
||||
def generate(start, end, *, n_points=64, speed=None, click=True, seed=None,
|
||||
model_path=None, providers=None) -> list[tuple[int, int, int]]: ...
|
||||
|
||||
def generate_scroll(start_scroll_y, target_scroll_y, *,
|
||||
mode="target", viewport_height=800, seed=None,
|
||||
model_path=None, providers=None) -> list[dict]: ...
|
||||
```
|
||||
|
||||
- 内部以 `(model_path or "__bundled__", tuple(providers or ()))` 为 key,`functools.lru_cache(maxsize=4)` 缓存模型实例
|
||||
- 多次调用同参数 ⇒ 复用同一 `InferenceSession`,不重复加载
|
||||
- `maxsize=4` 的选择: 典型场景只用 bundled 权重 + 默认 provider(size=1 够用);保留余量给"自定义权重 + CPU/GPU 双 provider"等极少数场景。再大没意义,大模型对象常驻内存反而浪费
|
||||
|
||||
#### 与旧 API 的关键差异
|
||||
|
||||
| 旧 | 新 | 原因 |
|
||||
|---|---|---|
|
||||
| `model_dir: str` | `model_path: str \| Path` | 接受 Path,语义更明确 |
|
||||
| `config: GenerateConfig` 参数 | 移除 | 内部物;可调项直接 kwarg |
|
||||
| 无显式 click 开关 | `click: bool = True` | 默认行为不变,但允许纯轨迹 |
|
||||
| 隐式 `viewport_norm` | `viewport_height: int = 800` | 暴露给调用者 |
|
||||
| 每次调用可能重载模型 | 类持久 / 函数走 lru_cache | SDK 不能"每次几百 ms 起步" |
|
||||
|
||||
#### 线程安全
|
||||
|
||||
ORT `InferenceSession.run` 本身线程安全。`MouseModel`/`ScrollModel` 共享一个 session,可在多线程并发 `.generate()`。
|
||||
|
||||
`seed=None` 时使用 `np.random.default_rng()` 局部实例,不污染全局 numpy 状态。
|
||||
|
||||
#### 错误层级
|
||||
|
||||
```python
|
||||
class AiMouseError(Exception): ...
|
||||
class ModelLoadError(AiMouseError): ... # 权重缺失 / shape 不符 / providers 不可用
|
||||
class GenerationError(AiMouseError): ... # 推理时数值异常 (NaN 等)
|
||||
```
|
||||
|
||||
### §2 ONNX 导出方案
|
||||
|
||||
#### 工具脚本 `tools/export_onnx.py`
|
||||
|
||||
一次性脚本(训练完成后跑),产出可烧进 wheel 的 `.onnx`。库代码不依赖 torch;只有这个脚本依赖。
|
||||
|
||||
```bash
|
||||
uv run python tools/export_onnx.py \
|
||||
--flow-ckpt data/models_v2 \
|
||||
--scroll-ckpt data/scroll_models \
|
||||
--output src/ai_mouse/assets/
|
||||
```
|
||||
|
||||
行为:
|
||||
- 读 `train_config.json` ⇒ 用一致超参实例化 `TrajectoryFlowModel`
|
||||
- `torch.load(flow_model.pt)` ⇒ `model.eval()` ⇒ 导出为 `flow_model.onnx`
|
||||
- 同样导出 scroll **decoder** 部分 ⇒ `scroll_decoder.onnx`(encoder 仅训练用)
|
||||
- 复制 `click_dist.json`、`duration_dist.json`、`train_config.json`、`scroll_config.json` 到 `assets/`
|
||||
|
||||
#### 轨迹模型导出
|
||||
|
||||
```python
|
||||
torch.onnx.export(
|
||||
model,
|
||||
args=(
|
||||
torch.zeros(1, seq_len, 3), # x_t
|
||||
torch.zeros(1), # t
|
||||
torch.zeros(1, 3), # cond
|
||||
),
|
||||
f="flow_model.onnx",
|
||||
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,
|
||||
)
|
||||
```
|
||||
|
||||
关键决策:
|
||||
- **`seq_len=64` 静态**: 位置嵌入是固定 shape 可学习参数;`n_points` API 若 ≠ 64,在后处理阶段线性插值重采样,**不进 ONNX**
|
||||
- **Batch 动态**: 支持下游批量生成
|
||||
- **Opset 17**: 支持 SDPA。如有问题脚本回落到 14
|
||||
- **常量折叠**: 位置嵌入、time/cond MLP 权重可折掉
|
||||
|
||||
#### Euler ODE 循环留在 Python
|
||||
|
||||
10 步 Euler 在 Python/NumPy 中循环,每步一次 `session.run()`:
|
||||
|
||||
```python
|
||||
x = rng.standard_normal((1, SEQ_LEN, 3), dtype=np.float32)
|
||||
dt = 1.0 / N_STEPS
|
||||
for i in range(N_STEPS):
|
||||
t = np.full((1,), (i + 0.5) * dt, dtype=np.float32)
|
||||
v = self.session.run(["v"], {"x_t": x, "t": t, "cond": cond_np})[0]
|
||||
x = x + dt * v
|
||||
```
|
||||
|
||||
权衡:10 次 ORT 调用开销 ≈ 1-3 ms;不引入 ONNX `Loop` op 复杂度;单文件权重;调试容易。未来如需性能可升级单图 Loop 版,本次不做。
|
||||
|
||||
#### 滚轮模型导出(只导 decoder)
|
||||
|
||||
```python
|
||||
class ScrollDecoder(nn.Module):
|
||||
"""Wrap ScrollCVAE.decode for ONNX export."""
|
||||
def forward(self, z, cond): ...
|
||||
|
||||
torch.onnx.export(
|
||||
ScrollDecoder(scroll_cvae),
|
||||
args=(torch.zeros(1, LATENT_DIM), torch.zeros(1, 7)),
|
||||
f="scroll_decoder.onnx",
|
||||
input_names=["z", "cond"],
|
||||
output_names=["seq"],
|
||||
dynamic_axes={"z": {0: "b"}, "cond": {0: "b"}, "seq": {0: "b"}},
|
||||
opset_version=17,
|
||||
)
|
||||
```
|
||||
|
||||
`z` 在 Python 端 `rng.standard_normal((1, LATENT_DIM))` 采样;ONNX 内不放采样算子。
|
||||
|
||||
#### 体积估算
|
||||
|
||||
| 文件 | 估算大小 |
|
||||
|---|---|
|
||||
| `flow_model.onnx` | ~2-3 MB (4 层 Transformer, d_model=128, ~600K 参数 FP32) |
|
||||
| `scroll_decoder.onnx` | <300 KB |
|
||||
| JSON 元数据 (3-4 个) | <20 KB |
|
||||
| **wheel 总增量** | **~3 MB** |
|
||||
|
||||
未来可选 `--quantize int8` 压到 ~1 MB,本次不做(YAGNI)。
|
||||
|
||||
#### 兼容性自检(导出脚本内置)
|
||||
|
||||
末尾跑烟雾测试:
|
||||
1. ORT 加载刚导出的 `.onnx`
|
||||
2. 与 PyTorch 模型同输入做前向
|
||||
3. `np.allclose(torch_out, ort_out, atol=1e-4)` 必须通过
|
||||
4. 失败 ⇒ 不写入 `assets/`,保留 `.pt`
|
||||
|
||||
### §3 后处理 / 坐标变换的纯 NumPy 重写
|
||||
|
||||
`coord.py`、`_gaussian_smooth`、`_sample_duration`、`utils.resample_arc`、所有空间/时间后处理已经是 numpy。需要替换的只有四处:
|
||||
|
||||
| 现在用 | 替换为 |
|
||||
|---|---|
|
||||
| `torch.randn(...)` | `rng.standard_normal((1, T, 3), dtype=np.float32)` |
|
||||
| `model(x_t, t, cond)` | `session.run(["v"], {...})[0]` |
|
||||
| `torch.load(...) / .eval()` | `onnxruntime.InferenceSession(asset_path, providers=...)` |
|
||||
| `scipy.stats.truncnorm.rvs(...)` | 纯 numpy 拒绝采样(见下) |
|
||||
|
||||
#### 库内部文件分布
|
||||
|
||||
```
|
||||
src/ai_mouse/
|
||||
├── __init__.py # 公开 API 拼装
|
||||
├── mouse.py # MouseModel: session + Euler 循环 + 后处理 + click
|
||||
├── scroll.py # ScrollModel: decoder run + 量化 + 时间戳
|
||||
├── _coord.py # encode_trajectory, decode_trajectory
|
||||
├── _postprocess.py # 平滑、单调性、重采样、截断正态、duration 采样 ...
|
||||
├── _assets.py # importlib.resources 加载 onnx + json
|
||||
├── errors.py # AiMouseError, ModelLoadError, GenerationError
|
||||
├── py.typed # 空文件
|
||||
└── assets/
|
||||
├── flow_model.onnx
|
||||
├── scroll_decoder.onnx
|
||||
├── click_dist.json
|
||||
├── duration_dist.json
|
||||
├── train_config.json
|
||||
└── scroll_config.json
|
||||
```
|
||||
|
||||
下划线开头的模块是私有实现,**不在 `__all__`**,版本间可自由破坏。
|
||||
|
||||
**`tools/` 与 `src/ai_mouse/` 的依赖方向**: `tools/` 是仓库内部的开发期代码,允许 `from ai_mouse._coord import ...` 反向引用库私有模块(它们在同一 repo 共同演化,只是 `tools/` 不进 wheel)。但 `src/ai_mouse/` **不允许** `import tools.*` —— wheel 必须自包含。CI 的 library job 会通过"无 torch 环境"间接守住这条规则。
|
||||
|
||||
#### `_truncnorm_sample`(替代 scipy)
|
||||
|
||||
```python
|
||||
def _truncnorm_sample(
|
||||
mu: float, sigma: float, low: float, high: float,
|
||||
rng: np.random.Generator, max_tries: int = 32,
|
||||
) -> float:
|
||||
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))
|
||||
```
|
||||
|
||||
默认参数 `mu=80, sigma=30, low=20, high=300`:接受率 ≈ 97%,无可测开销。
|
||||
|
||||
#### `_postprocess.py` 公开函数清单
|
||||
|
||||
```python
|
||||
def snap_endpoints(forward, lateral, seq_len, n_snap=6) -> tuple[np.ndarray, np.ndarray]: ...
|
||||
def smooth_start(forward, lateral, n=4) -> tuple[np.ndarray, np.ndarray]: ...
|
||||
def enforce_forward_monotonic(forward) -> np.ndarray: ...
|
||||
def gaussian_smooth(x, sigma=1.0) -> np.ndarray: ...
|
||||
def build_timestamps(log_dt, total_duration_ms, dt_clip=(2.0, 150.0)) -> np.ndarray: ...
|
||||
def sample_duration(duration_dist, dist, rng) -> float: ...
|
||||
def truncnorm_sample(mu, sigma, low, high, rng) -> float: ...
|
||||
def resample_arc(pixels, n_points) -> np.ndarray: ...
|
||||
```
|
||||
|
||||
每个函数纯函数,显式接收 `rng`,便于单测。
|
||||
|
||||
#### Scroll 后处理改动
|
||||
|
||||
`scroll/generator.py` 已完全 numpy,迁移只换两行:
|
||||
- `torch.randn(1, latent_dim)` ⇒ `rng.standard_normal((1, latent_dim), dtype=np.float32)`
|
||||
- `model.decode(z, cond_t)` ⇒ `session.run(["seq"], {"z": z, "cond": cond})[0]`
|
||||
|
||||
其余 softmax 归一化、量化、时间戳构建照搬。
|
||||
|
||||
### §4 仓库迁移
|
||||
|
||||
#### 终态布局
|
||||
|
||||
```
|
||||
ai_mouse/ (repo root)
|
||||
├── src/
|
||||
│ └── ai_mouse/ (唯一进 wheel 的目录)
|
||||
│ ├── __init__.py
|
||||
│ ├── mouse.py
|
||||
│ ├── scroll.py
|
||||
│ ├── _coord.py
|
||||
│ ├── _postprocess.py
|
||||
│ ├── _assets.py
|
||||
│ ├── errors.py
|
||||
│ ├── py.typed
|
||||
│ └── assets/
|
||||
│ ├── flow_model.onnx
|
||||
│ ├── scroll_decoder.onnx
|
||||
│ ├── click_dist.json
|
||||
│ ├── duration_dist.json
|
||||
│ ├── train_config.json
|
||||
│ └── scroll_config.json
|
||||
├── tools/ (开发期脚本,不进 wheel)
|
||||
│ ├── __init__.py
|
||||
│ ├── __main__.py (CLI dispatch, 从 ai_mouse/__main__.py 移来)
|
||||
│ ├── train.py
|
||||
│ ├── serve.py (从 ./main.py 移来)
|
||||
│ ├── export_onnx.py ★ 新
|
||||
│ ├── trainer.py
|
||||
│ ├── models.py (TrajectoryFlowModel — torch)
|
||||
│ ├── collector.py
|
||||
│ ├── config.py (TrainConfig 等训练侧 dataclass)
|
||||
│ ├── server/
|
||||
│ ├── eval/
|
||||
│ ├── data_adapters/
|
||||
│ └── scroll/
|
||||
│ ├── trainer.py
|
||||
│ ├── models.py
|
||||
│ └── collector.py
|
||||
├── data/ (不变)
|
||||
├── docs/ (不变)
|
||||
├── static/ (不变,tools/serve.py 引用)
|
||||
├── tests/
|
||||
│ ├── unit/ (库测试,只需 numpy + ort)
|
||||
│ │ ├── conftest.py
|
||||
│ │ ├── test_mouse.py
|
||||
│ │ ├── test_scroll.py
|
||||
│ │ ├── test_coord.py
|
||||
│ │ ├── test_postprocess.py
|
||||
│ │ ├── test_assets.py
|
||||
│ │ ├── test_errors.py
|
||||
│ │ └── data/
|
||||
│ │ ├── golden_mouse.npz
|
||||
│ │ ├── golden_scroll.npz
|
||||
│ │ └── tiny_flow.onnx
|
||||
│ └── tools/ (训练/服务器测试,需 torch+fastapi+...)
|
||||
│ ├── conftest.py
|
||||
│ ├── test_trainer.py
|
||||
│ ├── test_server.py
|
||||
│ ├── test_eval_metrics.py
|
||||
│ ├── test_balabit_adapter.py
|
||||
│ ├── test_export_onnx.py
|
||||
│ └── test_scroll_trainer.py
|
||||
├── examples/
|
||||
│ └── quickstart.py
|
||||
├── pyproject.toml (runtime: numpy+ort;[dev]: torch+fastapi+...)
|
||||
├── README.md (SDK 视角重写)
|
||||
├── CHANGELOG.md (新建,以 0.2.0 起锚)
|
||||
├── CLAUDE.md (更新)
|
||||
└── uv.lock
|
||||
```
|
||||
|
||||
被**删除**:
|
||||
- `ai_mouse/generator.py` ⇒ 被 `src/ai_mouse/mouse.py` 取代
|
||||
- `ai_mouse/scroll/generator.py` ⇒ 被 `src/ai_mouse/scroll.py` 取代
|
||||
- `ai_mouse/utils.py`(`resample_arc` 移入 `_postprocess.py`)
|
||||
- `JointCVAE`(legacy)
|
||||
- `_BUNDLED_MODELS_DIR = Path(__file__).parent.parent / "data" / ...` 这类源码树相对路径
|
||||
|
||||
#### 迁移分阶段
|
||||
|
||||
**阶段 1 — 物理移动,训练侧继续工作**
|
||||
- `git mv` 训练 / 服务器 / 采集 / 评估 / 数据适配器进 `tools/`
|
||||
- 修订 `tools/` 内部 import 为绝对 `from tools.X`
|
||||
- `main.py` ⇒ `tools/serve.py`
|
||||
- `tests/test_*` 按归属分进 `tests/unit/` 或 `tests/tools/`
|
||||
- 验收: `python -m tools train ...`、`python -m tools serve` 可跑;`from ai_mouse import generate` 仍能用(老 torch 实现)
|
||||
|
||||
**阶段 2 — src-layout 切换 + pyproject 更新**
|
||||
- `git mv ai_mouse src/ai_mouse`
|
||||
- pyproject 切到 hatchling,运行时依赖收紧到 `numpy + onnxruntime`;`[dev]` 群组收纳 torch+fastapi+scipy+matplotlib+pytest+...
|
||||
- 验收: `uv build` 产出 wheel,内容仅 `src/ai_mouse`
|
||||
|
||||
**阶段 3 — ONNX 导出脚本 + 权重塞进 assets/**
|
||||
- 写 `tools/export_onnx.py`
|
||||
- 运行导出,产出 `.onnx`+元数据 commit 进库
|
||||
- 验收: `pytest tests/tools/test_export_onnx.py` 验证 ORT vs PyTorch 数值一致
|
||||
|
||||
**阶段 4 — 库代码改写(纯 numpy + ORT)**
|
||||
- 抽 `_postprocess.py`、复制 `_coord.py`、写 `_assets.py`、`mouse.py`、`scroll.py`、`errors.py`
|
||||
- 重写 `__init__.py` 暴露 API + lru_cache
|
||||
- 删除阶段 1~2 残留的 torch 引用源文件
|
||||
- 验收: golden 测试通过;`uv pip install .` 不引入 torch;`from ai_mouse import generate` 直接出结果
|
||||
|
||||
**阶段 5 — 文档与清理**
|
||||
- 重写 README(SDK 视角)
|
||||
- 新建 CHANGELOG (0.2.0)
|
||||
- 更新 CLAUDE.md
|
||||
- 删 `JointCVAE`
|
||||
- 写 `examples/quickstart.py`
|
||||
|
||||
#### 风险与对策
|
||||
|
||||
| 风险 | 对策 |
|
||||
|---|---|
|
||||
| 权重二进制进 git (~3 MB) | 直接 commit;不引入 LFS |
|
||||
| 中间状态 CI 不绿 | 每个阶段独立 PR,自包含通过 |
|
||||
| `tools/` 内部 import 风格 | 一律绝对 `from tools.X`;`tools/__init__.py` 留空 |
|
||||
| 阶段 4 改 API 的"半破"中间态 | 阶段 4 是唯一签名变化节点,同 PR 内同步改 README/CLAUDE/examples |
|
||||
| `data/models_v2/` 路径引用混乱 | 库代码不读 `data/`(只读 `importlib.resources`);仅 `tools/export_onnx.py` 和 `tools/trainer.py` 关心 `data/` |
|
||||
| Golden 测试容差选择 | 起始 `atol=2`(像素 + ms),实测后可收紧 |
|
||||
| Windows 路径 | 全程 `pathlib.Path`;`importlib.resources` 跨平台 |
|
||||
| ORT CPU/GPU 包互斥 | README 明示 `pip install onnxruntime-gpu` 替换默认 `onnxruntime` |
|
||||
|
||||
### §5 测试策略
|
||||
|
||||
#### 两层测试,依赖边界硬隔离
|
||||
|
||||
`tests/unit/` 只装 `numpy + onnxruntime` 必须通过 — "纯推理 SDK"的硬保证。
|
||||
`tests/tools/` 装全套 `[dev]` 才跑训练/服务器/eval 测试。
|
||||
|
||||
#### Golden 测试(锁住语义不漂移)
|
||||
|
||||
迁移**开始之前**用旧 torch 实现跑 8 个 case × 4 个 seed = 32 个固定输出,存 `tests/unit/data/golden_mouse.npz`。
|
||||
|
||||
```python
|
||||
CASES = [
|
||||
((100, 200), (900, 400)), # 水平 800px
|
||||
((500, 500), (500, 100)), # 垂直 400px
|
||||
((200, 600), (800, 200)), # 720px 对角
|
||||
((100, 100), (130, 110)), # 极短 31px
|
||||
((50, 50), (1500, 900)), # 极长 1700px
|
||||
((400, 300), (500, 300)), # 水平 100px
|
||||
((300, 300), (700, 700)), # 45° 对角
|
||||
((600, 400), (200, 100)), # 反向对角
|
||||
]
|
||||
SEEDS = (0, 1, 2, 3)
|
||||
```
|
||||
|
||||
迁移后断言 `np.allclose(new_pts, golden_pts, atol=2)`(像素 + ms 容差)。
|
||||
|
||||
Scroll 同样做 golden;容差以 deltaY 量化粒度允许 ±1 个 quantum。
|
||||
|
||||
#### ONNX vs PyTorch 数值一致性
|
||||
|
||||
`tests/tools/test_export_onnx.py`:训练 1-epoch toy 模型(seq_len=8, d_model=16)⇒ 导出 ⇒ 同输入跑 PyTorch & ORT ⇒ `assert np.allclose(atol=1e-4)`。
|
||||
|
||||
真模型的一致性由 `tools/export_onnx.py` 末尾烟雾测试在导出时执行。
|
||||
|
||||
#### `_postprocess.py` 单元测试
|
||||
|
||||
每个纯函数单独测,边界条件覆盖。详细测试矩阵见原讨论;关键覆盖点:
|
||||
|
||||
| 函数 | 关键断言 |
|
||||
|---|---|
|
||||
| `snap_endpoints` | 输出端点严格 = (0,0) 和 (1,0);前 (seq_len - n_snap) 个点不变 |
|
||||
| `enforce_forward_monotonic` | 输出非递减 |
|
||||
| `gaussian_smooth` | 端点保持;权重和归一 |
|
||||
| `build_timestamps` | sum ≈ total_duration;严格单调;首项 0 |
|
||||
| `sample_duration` | 1000 次采样中位数落在 bin 对应 `[0.5×, 2×]exp(mu_log)` |
|
||||
| `truncnorm_sample` | 落在 `[low, high]`;均值 ≈ mu;兜底分支独立覆盖 |
|
||||
| `resample_arc` | 长度变换正确;N=M 恒等 |
|
||||
|
||||
#### `MouseModel` / `ScrollModel` 行为测试
|
||||
|
||||
- session 复用(同实例多次 generate 不重建 session)
|
||||
- 顶层 `generate()` 单例缓存命中
|
||||
- `seed` 复现性(同 seed 两次调用结果一致)
|
||||
- `click=False` 返回纯轨迹
|
||||
- 非法 `model_path` 抛 `ModelLoadError`
|
||||
|
||||
#### 线程安全测试
|
||||
|
||||
`ThreadPoolExecutor` 并发 32 次 `m.generate(...)`,断言全部完成、无异常。不保证并发下同 seed 严格一致(ORT 内部调度)。
|
||||
|
||||
#### CI 矩阵
|
||||
|
||||
```yaml
|
||||
jobs:
|
||||
library:
|
||||
# 只装运行时依赖 + pytest — 拦截库代码偷偷 import torch
|
||||
runs: |
|
||||
uv venv && uv pip install -e . pytest
|
||||
uv run pytest tests/unit -v
|
||||
dev:
|
||||
runs: |
|
||||
uv sync --group dev
|
||||
uv run pytest tests/ -v
|
||||
```
|
||||
|
||||
注: `pytest` 不能进运行时依赖(下游用户不需要),只能在 CI 安装行里临时加。
|
||||
`pytest-asyncio` 仅 `tests/tools/test_server.py` 需要,留在 `[dev]` 群组。
|
||||
|
||||
OS 矩阵: Linux + Windows;Python 3.12 + 3.13。
|
||||
|
||||
#### 不做的事
|
||||
|
||||
- ❌ 不测视觉真实性(主观,留给人眼审 eval report)
|
||||
- ❌ 不跑真 web E2E(httpx ASGI 集成已足够)
|
||||
- ❌ CI 不重训模型(用已 commit 的 ONNX)
|
||||
|
||||
### §6 文档与下游集成
|
||||
|
||||
#### README.md 重写(SDK 视角)
|
||||
|
||||
读者画像从"训练者"切换到"集成者"。结构:
|
||||
|
||||
1. 标题与一句话定位
|
||||
2. 安装(`pip install git+...` + GPU 可选说明)
|
||||
3. 快速上手(鼠标、滚轮、类用法、provider 切换、复现性)
|
||||
4. API 概览(表格)
|
||||
5. 线程安全说明
|
||||
6. 自训练模型 — 跳转 Development 小节
|
||||
7. Development(训练命令、Web UI、ONNX 导出 — 目标读者是 contributor)
|
||||
|
||||
#### `examples/quickstart.py`
|
||||
|
||||
最小可跑示例:`from ai_mouse import generate` ⇒ 打印 6 行结果 + 一段"实际使用模式"伪代码(时间戳节流 + 占位的鼠标移动调用)。
|
||||
|
||||
#### CHANGELOG.md(新建,Keep a Changelog 格式)
|
||||
|
||||
锚定 0.2.0:
|
||||
|
||||
```
|
||||
## [0.2.0] - YYYY-MM-DD
|
||||
### Changed (breaking)
|
||||
- Inference 不再依赖 PyTorch;运行时仅 numpy + onnxruntime
|
||||
- 公开 API 新增 MouseModel / ScrollModel 类
|
||||
- 函数式 generate / generate_scroll 签名: keyword-only;新增 click=, seed=;移除 config=
|
||||
- 训练 / 服务器 / 采集 / 评估代码移至 tools/,不再打入 wheel
|
||||
### Added
|
||||
- Wheel 内置 ONNX 预训练权重 (~3 MB)
|
||||
- ORT provider 自定义 (GPU/DirectML)
|
||||
### Removed
|
||||
- JointCVAE (legacy)
|
||||
- ai_mouse.config.GenerateConfig 顶层导出
|
||||
```
|
||||
|
||||
#### CLAUDE.md 更新清单
|
||||
|
||||
- Commands: `python -m ai_mouse <cmd>` ⇒ `python -m tools <cmd>`
|
||||
- 加一节"Library vs tools 边界": 库代码禁止 `import torch`
|
||||
- 移除 "Bundled weights live in `data/models_v2/`" ⇒ 改 `src/ai_mouse/assets/`
|
||||
- 测试章节区分 `tests/unit` / `tests/tools`
|
||||
|
||||
#### 下游契约 — 成功验收标准
|
||||
|
||||
1. `pip install git+...` ⇒ `from ai_mouse import generate` 直接出结果,无缺失文件、无 GB 级依赖
|
||||
2. `import ai_mouse` 时间 < 200 ms
|
||||
3. 单次 `generate()` < 50 ms (CPU,Euler 10 步 + 后处理)
|
||||
|
||||
#### 不做的事
|
||||
|
||||
- ❌ 不写 API reference 网站(docstring 足够)
|
||||
- ❌ 不写迁移 guide(破坏面有限,CHANGELOG 一段话)
|
||||
- ❌ 不录视频 / GIF
|
||||
- ❌ 不引入 docs i18n 流程
|
||||
|
||||
## 开放问题
|
||||
|
||||
无。所有决策点已经在与用户的对话中明确。
|
||||
|
||||
## 后续行动
|
||||
|
||||
本 spec 经用户批准后,进入 implementation plan 阶段(`writing-plans` skill),把 §4 的 5 个迁移阶段展开为可执行的 task list。
|
||||
Reference in New Issue
Block a user