Created dataset.py

Up-directories the module
cpuguy96 · Dec 2, 2023 · 1de0b30 · 1de0b30
1 parent cafc860
commit 1de0b30
Show file tree

Hide file tree

Showing 6 changed files with 439 additions and 345 deletions.
diff --git a/stepcovnet/config.py b/stepcovnet/config.py
@@ -4,9 +4,9 @@
 import numpy as np
 from sklearn.model_selection import train_test_split
 
+from stepcovnet import dataset
 from stepcovnet.common.constants import NUM_ARROW_COMBS
 from stepcovnet.common.utils import get_channel_scalers
-from stepcovnet.dataset.ModelDataset import ModelDataset
 from stepcovnet.training.TrainingHyperparameters import TrainingHyperparameters
 
 
@@ -54,7 +54,7 @@ class TrainingConfig(AbstractConfig):
     def __init__(
         self,
         dataset_path: str,
-        dataset_type: Type[ModelDataset],
+        dataset_type: Type[dataset.ModelDataset],
         dataset_config,
         hyperparameters: TrainingHyperparameters,
         all_scalers=None,

diff --git a/stepcovnet/data.py b/stepcovnet/data.py
@@ -2,14 +2,13 @@
 
 from transformers import GPT2Tokenizer
 
-from stepcovnet.dataset.DistributedModelDataset import DistributedModelDataset
-from stepcovnet.dataset.ModelDataset import ModelDataset
+from stepcovnet import dataset
 
 
 class Tokenizers(Enum):
     GPT2 = GPT2Tokenizer.from_pretrained("gpt2")
 
 
 class ModelDatasetTypes(Enum):
-    SINGULAR_DATASET = ModelDataset
-    DISTRIBUTED_DATASET = DistributedModelDataset
+    SINGULAR_DATASET = dataset.ModelDataset
+    DISTRIBUTED_DATASET = dataset.DistributedModelDataset