From ba52c49edfd24d429a63be199ebc4f30f5834351 Mon Sep 17 00:00:00 2001 From: Huang Qi Date: Tue, 12 May 2026 00:30:55 +0800 Subject: [PATCH] refactor: move trainer/models/utils/config to tools/ --- ai_mouse/data_adapters/balabit.py | 2 +- ai_mouse/generator.py | 6 +++--- ai_mouse/scroll/collector.py | 2 +- ai_mouse/scroll/trainer.py | 2 +- tests/conftest.py | 2 +- tests/test_balabit_adapter.py | 6 +++--- tests/test_generator.py | 2 +- tests/test_models.py | 2 +- tests/test_trainer.py | 26 +++++++++++++------------- {ai_mouse => tools}/config.py | 0 {ai_mouse => tools}/models.py | 0 {ai_mouse => tools}/trainer.py | 6 +++--- {ai_mouse => tools}/utils.py | 0 13 files changed, 28 insertions(+), 28 deletions(-) rename {ai_mouse => tools}/config.py (100%) rename {ai_mouse => tools}/models.py (100%) rename {ai_mouse => tools}/trainer.py (99%) rename {ai_mouse => tools}/utils.py (100%) diff --git a/ai_mouse/data_adapters/balabit.py b/ai_mouse/data_adapters/balabit.py index 1771d6c..55d0fbf 100644 --- a/ai_mouse/data_adapters/balabit.py +++ b/ai_mouse/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/generator.py b/ai_mouse/generator.py index 6b05055..2d6fefd 100644 --- a/ai_mouse/generator.py +++ b/ai_mouse/generator.py @@ -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__) diff --git a/ai_mouse/scroll/collector.py b/ai_mouse/scroll/collector.py index 33041e4..39b1993 100644 --- a/ai_mouse/scroll/collector.py +++ b/ai_mouse/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/trainer.py b/ai_mouse/scroll/trainer.py index ffe3bab..4ee2e10 100644 --- a/ai_mouse/scroll/trainer.py +++ b/ai_mouse/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 logger = logging.getLogger(__name__) diff --git a/tests/conftest.py b/tests/conftest.py index 3785cef..032e4ed 100644 --- a/tests/conftest.py +++ b/tests/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 @pytest.fixture diff --git a/tests/test_balabit_adapter.py b/tests/test_balabit_adapter.py index 995dacb..436e22c 100644 --- a/tests/test_balabit_adapter.py +++ b/tests/test_balabit_adapter.py @@ -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" diff --git a/tests/test_generator.py b/tests/test_generator.py index 4bbf09d..64ef61a 100644 --- a/tests/test_generator.py +++ b/tests/test_generator.py @@ -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 diff --git a/tests/test_models.py b/tests/test_models.py index bfe81fc..0a6674a 100644 --- a/tests/test_models.py +++ b/tests/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_trainer.py b/tests/test_trainer.py index 8cbf79b..a6d6b24 100644 --- a/tests/test_trainer.py +++ b/tests/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/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/models.py b/tools/models.py similarity index 100% rename from ai_mouse/models.py rename to tools/models.py diff --git a/ai_mouse/trainer.py b/tools/trainer.py similarity index 99% rename from ai_mouse/trainer.py rename to tools/trainer.py index cc5fdfd..65ecf5c 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 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