feat: Choose dataset with options

2025-11-30 20:19:39 +01:00 · 2025-11-30 20:19:39 +01:00 · 81c767371e
commit 81c767371e
parent 20bdd4f566
5 changed files with 67 additions and 60 deletions
--- a/CNN-model/dataset_loaders/init.py
+++ b/CNN-model/dataset_loaders/init.py
@ -1,3 +1,8 @@
+from .Dataset import Dataset
 from .EnWik9 import EnWik9DataSet
 from .LoremIpsumDataset import LoremIpsumDataset
-from .Dataset import Dataset
+
+dataset_called: dict[str, type[Dataset]] = {
+    'enwik9': EnWik9DataSet,
+    'lorem_ipsum': LoremIpsumDataset
+}