feat: unify folder structure of gtfs2vec with rest of the library

kraina-ai · Apr 19, 2023 · 554d9c8 · 554d9c8
1 parent 48cbb27
commit 554d9c8
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 7 deletions.
diff --git a/srai/embedders/__init__.py b/srai/embedders/__init__.py
@@ -2,7 +2,7 @@
 
 from ._base import Embedder
 from .count_embedder import CountEmbedder
-from .gtfs2vec_embedder import GTFS2VecEmbedder
+from .gtfs2vec import GTFS2VecEmbedder
 from .highway2vec import Highway2VecEmbedder
 
 __all__ = ["Embedder", "CountEmbedder", "GTFS2VecEmbedder", "Highway2VecEmbedder"]
diff --git a/srai/embedders/gtfs2vec/__init__.py b/srai/embedders/gtfs2vec/__init__.py
@@ -0,0 +1,6 @@
+"""Models."""
+
+from .embedder import GTFS2VecEmbedder
+from .model import GTFS2VecModel
+
+__all__ = ["GTFS2VecEmbedder", "GTFS2VecModel"]
diff --git a/srai/embedders/gtfs2vec_embedder.py → srai/embedders/gtfs2vec/embedder.py b/srai/embedders/gtfs2vec_embedder.py → srai/embedders/gtfs2vec/embedder.py
@@ -19,9 +19,9 @@
 from torch.utils.data import DataLoader
 
 from srai.embedders import Embedder
+from srai.embedders.gtfs2vec.model import GTFS2VecModel
 from srai.exceptions import ModelNotFitException
 from srai.loaders.gtfs_loader import GTFS2VEC_DIRECTIONS_PREFIX, GTFS2VEC_TRIPS_PREFIX
-from srai.models import GTFS2VecModel
 
 
 class GTFS2VecEmbedder(Embedder):

diff --git a/srai/models/gtfs2vec_model.py → srai/embedders/gtfs2vec/model.py b/srai/models/gtfs2vec_model.py → srai/embedders/gtfs2vec/model.py
diff --git a/srai/models/__init__.py b/srai/models/__init__.py