From 554d9c80bf390c6569e72f21e4b5cd1ae3dfabc9 Mon Sep 17 00:00:00 2001
From: Piotr
Date: Wed, 19 Apr 2023 18:39:00 +0200
Subject: [PATCH] feat: unify folder structure of gtfs2vec with rest of the
library
---
srai/embedders/__init__.py | 2 +-
srai/embedders/gtfs2vec/__init__.py | 6 ++++++
.../{gtfs2vec_embedder.py => gtfs2vec/embedder.py} | 2 +-
.../gtfs2vec_model.py => embedders/gtfs2vec/model.py} | 0
srai/models/__init__.py | 5 -----
5 files changed, 8 insertions(+), 7 deletions(-)
create mode 100644 srai/embedders/gtfs2vec/__init__.py
rename srai/embedders/{gtfs2vec_embedder.py => gtfs2vec/embedder.py} (99%)
rename srai/{models/gtfs2vec_model.py => embedders/gtfs2vec/model.py} (100%)
delete mode 100644 srai/models/__init__.py
diff --git a/srai/embedders/__init__.py b/srai/embedders/__init__.py
index 71947425..eeaf5ebc 100644
--- 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
new file mode 100644
index 00000000..36ae82d6
--- /dev/null
+++ 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 b/srai/embedders/gtfs2vec/embedder.py
similarity index 99%
rename from srai/embedders/gtfs2vec_embedder.py
rename to srai/embedders/gtfs2vec/embedder.py
index 876a7919..d25d59a5 100644
--- a/srai/embedders/gtfs2vec_embedder.py
+++ b/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 b/srai/embedders/gtfs2vec/model.py
similarity index 100%
rename from srai/models/gtfs2vec_model.py
rename to srai/embedders/gtfs2vec/model.py
diff --git a/srai/models/__init__.py b/srai/models/__init__.py
deleted file mode 100644
index b5e07f24..00000000
--- a/srai/models/__init__.py
+++ /dev/null
@@ -1,5 +0,0 @@
-"""Models."""
-
-from .gtfs2vec_model import GTFS2VecModel
-
-__all__ = ["GTFS2VecModel"]