fix: fixing DL pipeline

backend/app/pipeline/generators/dataset_generator.py

@@ -1,6 +1,7 @@
 def generate_dataset_code(config):
     dataset_name = config.get("dataset", "iris")
-    max_samples = config.get("max_samples", 2000)
+    data_dir = config.get("data_dir", "data")
+    max_samples = config.get("max_samples", 2000) if dataset_name in ["mnist", "fashion_mnist", "cifar10"] else None
 
     image_metadata = {
         "digits": (1, 8, 8),
@@ -9,32 +10,74 @@ def generate_dataset_code(config):
         "cifar10": (3, 32, 32),
     }
 
-    if dataset_name in ["mnist", "fashion_mnist", "cifar10"]:
-        channels, height, width = image_metadata[dataset_name]
-
-        imports = {
-            f"from app.services.modal_service import get_{dataset_name}"
-        }
+    torchvision_loaders = {
+        "mnist": "MNIST",
+        "fashion_mnist": "FashionMNIST",
+        "cifar10": "CIFAR10",
+    }
 
+    if dataset_name in torchvision_loaders:
+        loader_name = torchvision_loaders[dataset_name]
+        channels, height, width = image_metadata[dataset_name]
+        imports = {"from torchvision import datasets"}
+
         code = [
-            f"dataset = get_{dataset_name}({max_samples})",
-            "X = dataset['X']",
-            "y = dataset['y']",
+            f"dataset = datasets.{loader_name}(",
+            f"    root={data_dir!r},",
+            "    train=True,",
+            "    download=True",
+            ")",
+            "X = dataset.data",
+            "if hasattr(X, 'numpy'):",
+            "    X = X.numpy()",
+        ]
+        if max_samples is not None:
+            code.append(f"X = X[:{max_samples}]")
+
+        code.extend([
+            "X = X.tolist()",
+            "y = dataset.targets",
+            "if hasattr(y, 'numpy'):",
+            "    y = y.numpy()",
+        ])
+
+        if max_samples is not None:
+            code.append(f"y = y[:{max_samples}]")
+
+        code.extend([
+            "if hasattr(y, 'tolist'):",
+            "    y = y.tolist()",
             "data_format = 'image'",
             f"image_channels = {channels}",
             f"image_height = {height}",
             f"image_width = {width}",
-        ]
+            ""
+        ])
+        return imports, code
 
-    else:
-        imports = {
-            "from sklearn.datasets import load_iris"
-        }
+    loaders = {
+        "california_housing": "fetch_california_housing",
+        "digits": "load_digits"
+    }
+    loader_name = loaders.get(dataset_name, f"load_{dataset_name}")
+    imports = {f"from sklearn.datasets import {loader_name}"}
+    data_attribute = "images" if dataset_name == "digits" else "data"
 
-        code = [
-            "dataset = load_iris()",
-            "X = dataset.data",
-            "y = dataset.target",
-        ]
+    code = [
+        f"dataset = {loader_name}()",
+        f"X = dataset.{data_attribute}",
+        "y = dataset.target",
+    ]
+
+    if dataset_name == "digits":
+        channels, height, width = image_metadata[dataset_name]
+        code.extend([
+            "data_format = 'image'",
+            f"image_channels = {channels}",
+            f"image_height = {height}",
+            f"image_width = {width}",
+        ])
+
+    code.append("")
 
     return imports, code

backend/tests/conftest.py

@@ -0,0 +1,114 @@
+import pytest
+import app.services.modal_service
+import app.pipeline.nodes.neural_network_node
+import app.pipeline.nodes.dataset_node
+
+
+def _zeros_3d(n, h, w):
+    return [[[0.0] * w for _ in range(h)] for _ in range(n)]
+
+
+def _zeros_4d(n, h, w, c):
+    return [[[[0.0] * c for _ in range(w)] for _ in range(h)] for _ in range(n)]
+
+
+def _split(X, y, test_size=0.2, random_state=42):
+    split_idx = max(1, int(len(X) * (1 - test_size)))
+    return X[:split_idx], X[split_idx:], y[:split_idx], y[split_idx:]
+
+
+@pytest.fixture(autouse=True)
+def mock_modal_service(monkeypatch):
+    def mock_get_mnist(max_samples=2000):
+        return {
+            "X": _zeros_3d(max_samples, 28, 28),
+            "y": [i % 10 for i in range(max_samples)],
+            "dataset_name": "mnist",
+            "task_type": "classification",
+            "data_format": "image",
+            "image_channels": 1,
+            "image_height": 28,
+            "image_width": 28,
+        }
+
+    def mock_get_fashion_mnist(max_samples=2000):
+        return {
+            "X": _zeros_3d(max_samples, 28, 28),
+            "y": [i % 10 for i in range(max_samples)],
+            "dataset_name": "fashion_mnist",
+            "task_type": "classification",
+            "data_format": "image",
+            "image_channels": 1,
+            "image_height": 28,
+            "image_width": 28,
+        }
+
+    def mock_get_cifar10(max_samples=2000):
+        return {
+            "X": _zeros_4d(max_samples, 32, 32, 3),
+            "y": [i % 10 for i in range(max_samples)],
+            "dataset_name": "cifar10",
+            "task_type": "classification",
+            "data_format": "image",
+            "image_channels": 3,
+            "image_height": 32,
+            "image_width": 32,
+        }
+
+    def mock_run_mlp(input_data, config):
+        from modal_service.trainers.mlp_trainer import train
+        return train(input_data, config)
+
+    def mock_run_cnn(input_data, config):
+        from modal_service.trainers.cnn_trainer import train
+        return train(input_data, config)
+
+    def _build_split_input(dataset_name, max_samples, split_config):
+        test_size = split_config.get("test_size", 0.2)
+        if dataset_name == "mnist":
+            X = _zeros_3d(max_samples, 28, 28)
+            meta = {"data_format": "image", "image_channels": 1, "image_height": 28, "image_width": 28}
+        elif dataset_name == "fashion_mnist":
+            X = _zeros_3d(max_samples, 28, 28)
+            meta = {"data_format": "image", "image_channels": 1, "image_height": 28, "image_width": 28}
+        elif dataset_name == "cifar10":
+            X = _zeros_4d(max_samples, 32, 32, 3)
+            meta = {"data_format": "image", "image_channels": 3, "image_height": 32, "image_width": 32}
+        else:
+            raise ValueError(f"Unknown dataset: {dataset_name}")
+        y = [i % 10 for i in range(max_samples)]
+        X_train, X_test, y_train, y_test = _split(X, y, test_size=test_size)
+        return {
+            "X_train": X_train,
+            "X_test": X_test,
+            "y_train": y_train,
+            "y_test": y_test,
+            "task_type": "classification",
+            "dataset_name": dataset_name,
+            **meta
+        }
+
+    def mock_run_split_and_train_mlp(dataset_name, max_samples, split_config, train_config):
+        from modal_service.trainers.mlp_trainer import train
+        return train(_build_split_input(dataset_name, max_samples, split_config), train_config)
+
+    def mock_run_split_and_train_cnn(dataset_name, max_samples, split_config, train_config):
+        from modal_service.trainers.cnn_trainer import train
+        return train(_build_split_input(dataset_name, max_samples, split_config), train_config)
+
+    monkeypatch.setattr(app.services.modal_service, "get_mnist", mock_get_mnist)
+    monkeypatch.setattr(app.services.modal_service, "get_fashion_mnist", mock_get_fashion_mnist)
+    monkeypatch.setattr(app.services.modal_service, "get_cifar10", mock_get_cifar10)
+    monkeypatch.setattr(app.services.modal_service, "run_mlp", mock_run_mlp)
+    monkeypatch.setattr(app.services.modal_service, "run_cnn", mock_run_cnn)
+    monkeypatch.setattr(app.services.modal_service, "run_split_and_train_mlp", mock_run_split_and_train_mlp)
+    monkeypatch.setattr(app.services.modal_service, "run_split_and_train_cnn", mock_run_split_and_train_cnn)
+
+    monkeypatch.setattr(app.pipeline.nodes.neural_network_node, "run_mlp", mock_run_mlp)
+    monkeypatch.setattr(app.pipeline.nodes.neural_network_node, "run_cnn", mock_run_cnn)
+    monkeypatch.setattr(app.pipeline.nodes.neural_network_node, "run_split_and_train_mlp", mock_run_split_and_train_mlp)
+    monkeypatch.setattr(app.pipeline.nodes.neural_network_node, "run_split_and_train_cnn", mock_run_split_and_train_cnn)
+
+    monkeypatch.setattr(app.pipeline.nodes.dataset_node, "get_mnist", mock_get_mnist)
+    monkeypatch.setattr(app.pipeline.nodes.dataset_node, "get_fashion_mnist", mock_get_fashion_mnist)
+    monkeypatch.setattr(app.pipeline.nodes.dataset_node, "get_cifar10", mock_get_cifar10)