yellowdog/ai_mouse
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: move trainer/models/utils/config to tools/

Browse Source
This commit is contained in:
Huang Qi
2026-05-12 00:30:55 +08:00
parent c89025047c
commit ba52c49edf
13 changed files with 28 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
ai_mouse/data_adapters/balabit.py
View File

@@ -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(

6
ai_mouse/generator.py
View File

@@ -32,10 +32,10 @@ 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
from tools.config import GenerateConfig
from tools.models import TrajectoryFlowModel
from tools.utils import resample_arc
logger = logging.getLogger(__name__)

2
ai_mouse/scroll/collector.py
View File

@@ -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__)

2
ai_mouse/scroll/trainer.py
View File

@@ -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
logger = logging.getLogger(__name__)

2
tests/conftest.py
View File

@@ -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
@pytest.fixture

6
tests/test_balabit_adapter.py
View File

@@ -266,7 +266,7 @@ class TestFilterSegments:
class TestProcessSession:
def test_writes_jsonl_in_expected_format(self, tmp_path):
from ai_mouse.config import BalabitAdapterConfig
from tools.config import BalabitAdapterConfig
from ai_mouse.data_adapters.balabit import process_session
# Construct a Balabit-format CSV with one valid segment
@@ -303,7 +303,7 @@ 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 tools.config import BalabitAdapterConfig
from ai_mouse.data_adapters.balabit import process_session
csv_path = tmp_path / "empty_session"
@@ -318,7 +318,7 @@ class TestProcessSession:
assert out.read_text() == ""
def test_appends_to_existing_jsonl(self, tmp_path):
from ai_mouse.config import BalabitAdapterConfig
from tools.config import BalabitAdapterConfig
from ai_mouse.data_adapters.balabit import process_session
out = tmp_path / "out.jsonl"

2
tests/test_generator.py
View File

@@ -9,7 +9,7 @@ import pytest
import torch
from ai_mouse.generator import generate
from ai_mouse.models import TrajectoryFlowModel
from tools.models import TrajectoryFlowModel
@pytest.fixture

2
tests/test_models.py
View File

@@ -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:

26
tests/test_trainer.py
View File

@@ -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(

0
ai_mouse/config.py → tools/config.py
View File

0
ai_mouse/models.py → tools/models.py
View File

6
ai_mouse/trainer.py → tools/trainer.py
View File

@@ -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 tools.config import TrainConfig
from tools.models import TrajectoryFlowModel
from tools.utils import resample_arc
logger = logging.getLogger(__name__)

0
ai_mouse/utils.py → tools/utils.py
View File