Merge prepare_dataset.py from feature/ui-fix

data_preparation/prepare_dataset.py

@@ -5,22 +5,25 @@ import numpy as np
 from sklearn.preprocessing import StandardScaler
 from sklearn.model_selection import train_test_split
 
-try:
-    import torch
-    from torch.utils.data import Dataset, DataLoader
-except ImportError:  # pragma: no cover
-    torch = None
+torch = None
+Dataset = object  # type: ignore
+DataLoader = None
 
-    class Dataset:  # type: ignore
-        pass
 
-    class _MissingTorchDataLoader:  # type: ignore
-        def __init__(self, *args, **kwargs):
-            raise ImportError(
-                "PyTorch not installed"
-            )
+def _require_torch():
+    global torch, Dataset, DataLoader
+    if torch is None:
+        try:
+            import torch as _torch
+            from torch.utils.data import Dataset as _Dataset, DataLoader as _DataLoader
+        except ImportError as exc:  # pragma: no cover
+            raise ImportError("PyTorch not installed") from exc
 
-    DataLoader = _MissingTorchDataLoader  # type: ignore
+        torch = _torch
+        Dataset = _Dataset  # type: ignore
+        DataLoader = _DataLoader  # type: ignore
+
+    return torch, Dataset, DataLoader
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "data")
 
@@ -38,8 +41,9 @@ SELECTED_FEATURES = {
 
 class FeatureVectorDataset(Dataset):
     def __init__(self, features: np.ndarray, labels: np.ndarray):
-        self.features = torch.tensor(features, dtype=torch.float32)
-        self.labels = torch.tensor(labels, dtype=torch.long)
+        torch_mod, _, _ = _require_torch()
+        self.features = torch_mod.tensor(features, dtype=torch_mod.float32)
+        self.labels = torch_mod.tensor(labels, dtype=torch_mod.long)
 
     def __len__(self):
         return len(self.labels)
@@ -217,6 +221,7 @@ def get_numpy_splits(model_name: str, split_ratios=(0.7, 0.15, 0.15), seed: int
 
 def get_dataloaders(model_name: str, batch_size: int = 32, split_ratios=(0.7, 0.15, 0.15), seed: int = 42, scale: bool = True):
     """Return PyTorch DataLoaders for neural-network models."""
+    _, _, dataloader_cls = _require_torch()
     features, labels = _load_real_data(model_name)
     num_features = features.shape[1]
     num_classes = int(labels.max()) + 1
@@ -228,9 +233,9 @@ def get_dataloaders(model_name: str, batch_size: int = 32, split_ratios=(0.7, 0.
     val_ds   = FeatureVectorDataset(splits["X_val"],   splits["y_val"])
     test_ds  = FeatureVectorDataset(splits["X_test"],  splits["y_test"])
 
-    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
-    val_loader   = DataLoader(val_ds,   batch_size=batch_size, shuffle=False)
-    test_loader  = DataLoader(test_ds,  batch_size=batch_size, shuffle=False)
+    train_loader = dataloader_cls(train_ds, batch_size=batch_size, shuffle=True)
+    val_loader   = dataloader_cls(val_ds,   batch_size=batch_size, shuffle=False)
+    test_loader  = dataloader_cls(test_ds,  batch_size=batch_size, shuffle=False)
 
     return train_loader, val_loader, test_loader, num_features, num_classes, scaler