diff --git a/ai_mouse/scroll/__init__.py b/ai_mouse/scroll/__init__.py index 3db526b..8dd7bfa 100644 --- a/ai_mouse/scroll/__init__.py +++ b/ai_mouse/scroll/__init__.py @@ -1,6 +1,4 @@ -"""Scroll wheel event generation subsystem.""" +"""Scroll wheel event generation (inference only).""" 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"] +__all__ = ["generate_scroll"] diff --git a/ai_mouse/scroll/generator.py b/ai_mouse/scroll/generator.py index ec23ff6..f066d1e 100644 --- a/ai_mouse/scroll/generator.py +++ b/ai_mouse/scroll/generator.py @@ -9,7 +9,7 @@ from pathlib import Path import numpy as np import torch -from ai_mouse.scroll.models import ScrollCVAE +from tools.scroll.models import ScrollCVAE logger = logging.getLogger(__name__) diff --git a/ai_mouse/server/routes_scroll.py b/ai_mouse/server/routes_scroll.py index d57e114..0c9af75 100644 --- a/ai_mouse/server/routes_scroll.py +++ b/ai_mouse/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: diff --git a/tests/conftest.py b/tests/conftest.py index 032e4ed..755b9e3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,8 +8,8 @@ import numpy as np import pytest import torch -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_scroll_collector.py b/tests/test_scroll_collector.py index f30c328..b632281 100644 --- a/tests/test_scroll_collector.py +++ b/tests/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_generator.py b/tests/test_scroll_generator.py index e5bf238..e4cc098 100644 --- a/tests/test_scroll_generator.py +++ b/tests/test_scroll_generator.py @@ -8,7 +8,7 @@ import torch import pytest from ai_mouse.scroll.generator import generate_scroll -from ai_mouse.scroll.models import ScrollCVAE +from tools.scroll.models import ScrollCVAE @pytest.fixture diff --git a/tests/test_scroll_models.py b/tests/test_scroll_models.py index 98175e8..0d84715 100644 --- a/tests/test_scroll_models.py +++ b/tests/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/test_scroll_trainer.py index 7cb9e53..6ffc7b7 100644 --- a/tests/test_scroll_trainer.py +++ b/tests/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/ai_mouse/scroll/collector.py b/tools/scroll/collector.py similarity index 100% rename from ai_mouse/scroll/collector.py rename to tools/scroll/collector.py 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 99% rename from ai_mouse/scroll/trainer.py rename to tools/scroll/trainer.py index 4ee2e10..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.scroll.models import ScrollCVAE from tools.config import ScrollTrainConfig +from tools.scroll.models import ScrollCVAE logger = logging.getLogger(__name__)