From 554d9c80bf390c6569e72f21e4b5cd1ae3dfabc9 Mon Sep 17 00:00:00 2001
From: Piotr <p.gramacki@gmail.com>
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"]