57 lines
1.6 KiB
Python
57 lines
1.6 KiB
Python
"""Shared test fixtures for ai_mouse."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
from ai_mouse.scroll.models import ScrollCVAE
|
|
from tools.models import TrajectoryFlowModel
|
|
|
|
|
|
@pytest.fixture
|
|
def model_dir(tmp_path: Path) -> Path:
|
|
"""Create a temporary directory with trained Flow model artifacts."""
|
|
# Flow model
|
|
model = TrajectoryFlowModel(seq_len=64)
|
|
torch.save(model.state_dict(), tmp_path / "flow_model.pt")
|
|
|
|
# Click distribution
|
|
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 distribution
|
|
dur_dist = {
|
|
"bins": [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")],
|
|
"params": [{"mu_log": 5.5, "sigma_log": 0.5}] * 8,
|
|
}
|
|
(tmp_path / "duration_dist.json").write_text(json.dumps(dur_dist))
|
|
|
|
# Train config (architecture params)
|
|
train_cfg = {
|
|
"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_cfg))
|
|
|
|
return tmp_path
|
|
|
|
|
|
@pytest.fixture
|
|
def scroll_model_dir(tmp_path: Path) -> Path:
|
|
"""Create a temporary directory with trained scroll model artifacts."""
|
|
model = ScrollCVAE(seq_len=32)
|
|
torch.save(model.state_dict(), tmp_path / "scroll_model.pt")
|
|
|
|
scroll_cfg = {"seq_len": 32, "latent_dim": 16, "hidden": 64, "cond_dim": 7}
|
|
(tmp_path / "scroll_config.json").write_text(json.dumps(scroll_cfg))
|
|
|
|
return tmp_path
|