feat: added modular code generation and advanced preprocessing system

backend/app/pipeline/generators/code_generator.py

@@ -1,85 +1,25 @@
+from app.pipeline.generators.code_registry import CODE_GENERATOR_REGISTRY
+
 def generate_pipeline_code(pipeline):
-    code_lines = [
-        "from sklearn import datasets",
-        "from sklearn.model_selection import train_test_split",
-        "from sklearn.preprocessing import StandardScaler",
-        "from sklearn.linear_model import LogisticRegression, LinearRegression",
-        "from sklearn.tree import DecisionTreeClassifier",
-        "from sklearn.ensemble import RandomForestClassifier",
-        "",
-    ]
+    all_imports = set()
+    all_code = []
 
     for node in pipeline["nodes"]:
         node_type = node["type"]
         config = node.get("config", {})
 
-        if node_type == "dataset":
-            dataset_name = config.get("dataset", "iris")
-
-            code_lines.extend([
-                f"dataset = datasets.load_{dataset_name}()",
-                "X = dataset.data",
-                "y = dataset.target",
-                ""
-            ])
-
-        elif node_type == "train_test_split":
-            test_size = config.get("test_size", 0.2)
-            random_state = config.get("random_state", 42)
-
-            code_lines.extend([
-                "X_train, X_test, y_train, y_test = train_test_split(",
-                "    X,",
-                "    y,",
-                f"    test_size={test_size},",
-                f"    random_state={random_state}",
-                ")",
-                ""
-            ])
-
-        elif node_type == "preprocess":
-            code_lines.extend([
-                "scaler = StandardScaler()",
-                "X_train = scaler.fit_transform(X_train)",
-                "X_test = scaler.transform(X_test)",
-                ""
-            ])
-
-        elif node_type == "model":
-            algorithm = config.get("algorithm", "logistic_regression")
-
-            if algorithm == "logistic_regression":
-                code_lines.extend([
-                    "model = LogisticRegression(max_iter=1000)",
-                    "model.fit(X_train, y_train)",
-                    "predictions = model.predict(X_test)",
-                    ""
-                ])
+        generator = CODE_GENERATOR_REGISTRY.get(node_type)
 
-            elif algorithm == "linear_regression":
-                fit_intercept = config.get("fit_intercept", True)
+        if not generator:
+            continue
 
-                code_lines.extend([
-                    f"model = LinearRegression(fit_intercept={fit_intercept})",
-                    "model.fit(X_train, y_train)",
-                    "predictions = model.predict(X_test)",
-                    ""
-                ])
+        imports, code = generator(config)
 
-            elif algorithm == "decision_tree":
-                code_lines.extend([
-                    "model = DecisionTreeClassifier(random_state=42)",
-                    "model.fit(X_train, y_train)",
-                    "predictions = model.predict(X_test)",
-                    ""
-                ])
+        all_imports.update(imports)
+        all_code.extend(code)
 
-            elif algorithm == "random_forest":
-                code_lines.extend([
-                    "model = RandomForestClassifier(random_state=42)",
-                    "model.fit(X_train, y_train)",
-                    "predictions = model.predict(X_test)",
-                    ""
-                ])
+    final_code = sorted(all_imports)
+    final_code.append("")
+    final_code.extend(all_code)
 
-    return "\n".join(code_lines)
+    return "\n".join(final_code)

backend/app/pipeline/generators/code_registry.py

@@ -0,0 +1,11 @@
+from app.pipeline.generators.dataset_generator import generate_dataset_code
+from app.pipeline.generators.split_generator import generate_split_code
+from app.pipeline.generators.preprocess_generator import generate_preprocess_code
+from app.pipeline.generators.model_generator import generate_model_code
+
+CODE_GENERATOR_REGISTRY = {
+    "dataset": generate_dataset_code,
+    "train_test_split": generate_split_code,
+    "preprocess": generate_preprocess_code,
+    "model": generate_model_code
+}

backend/app/pipeline/generators/dataset_generator.py

@@ -0,0 +1,14 @@
+def generate_dataset_code(config):
+    dataset_name = config.get("dataset", "iris")
+
+    imports = {f"from sklearn.datasets import load_{dataset_name}"}
+
+    code = [
+        "# Load Dataset",
+        f"dataset = load_{dataset_name}()",
+        "X = dataset.data",
+        "y = dataset.target",
+        ""
+    ]
+
+    return imports, code

backend/app/pipeline/generators/model_generator.py

@@ -0,0 +1,58 @@
+def generate_model_code(config):
+    algorithm = config.get("algorithm", "logistic_regression")
+
+    imports = set()
+    code = ["# Model Training"]
+
+    if algorithm == "logistic_regression":
+        imports.add("from sklearn.linear_model import LogisticRegression")
+
+        C = config.get("C", 1.0)
+        solver = config.get("solver", "lbfgs")
+
+        code.extend([
+            f'model = LogisticRegression(C={C}, solver="{solver}", max_iter=1000)',
+            "model.fit(X_train, y_train)",
+            "predictions = model.predict(X_test)",
+            ""
+        ])
+
+    elif algorithm == "linear_regression":
+        imports.add("from sklearn.linear_model import LinearRegression")
+
+        fit_intercept = config.get("fit_intercept", True)
+
+        code.extend([
+            f"model = LinearRegression(fit_intercept={fit_intercept})",
+            "model.fit(X_train, y_train)",
+            "predictions = model.predict(X_test)",
+            ""
+        ])
+
+    elif algorithm == "decision_tree":
+        imports.add("from sklearn.tree import DecisionTreeClassifier")
+
+        max_depth = config.get("max_depth", 5)
+        criterion = config.get("criterion", "gini")
+
+        code.extend([
+            f'model = DecisionTreeClassifier(max_depth={max_depth}, criterion="{criterion}", random_state=42)',
+            "model.fit(X_train, y_train)",
+            "predictions = model.predict(X_test)",
+            ""
+        ])
+
+    elif algorithm == "random_forest":
+        imports.add("from sklearn.ensemble import RandomForestClassifier")
+
+        n_estimators = config.get("n_estimators", 100)
+        max_depth = config.get("max_depth", 5)
+
+        code.extend([
+            f"model = RandomForestClassifier(n_estimators={n_estimators}, max_depth={max_depth}, random_state=42)",
+            "model.fit(X_train, y_train)",
+            "predictions = model.predict(X_test)",
+            ""
+        ])
+
+    return imports, code

backend/app/pipeline/generators/preprocess_generator.py

@@ -0,0 +1,24 @@
+def generate_preprocess_code(config):
+    scaler_type = config.get("scaler_type", "standard")
+
+    scaler_map = {
+        "standard": "StandardScaler",
+        "minmax": "MinMaxScaler",
+        "robust": "RobustScaler"
+    }
+
+    scaler_class = scaler_map[scaler_type]
+
+    imports = {
+        f"from sklearn.preprocessing import {scaler_class}"
+    }
+
+    code = [
+        "# Preprocessing",
+        f"scaler = {scaler_class}()",
+        "X_train = scaler.fit_transform(X_train)",
+        "X_test = scaler.transform(X_test)",
+        ""
+    ]
+
+    return imports, code

backend/app/pipeline/generators/split_generator.py

@@ -0,0 +1,20 @@
+def generate_split_code(config):
+    test_size = config.get("test_size", 0.2)
+    random_state = config.get("random_state", 42)
+
+    imports = {
+        "from sklearn.model_selection import train_test_split"
+    }
+
+    code = [
+        "# Train Test Split",
+        "X_train, X_test, y_train, y_test = train_test_split(",
+        "    X,",
+        "    y,",
+        f"    test_size={test_size},",
+        f"    random_state={random_state}",
+        ")",
+        ""
+    ]
+
+    return imports, code

backend/app/pipeline/models/decision_tree.py

@@ -7,7 +7,11 @@ def train(input_data, config):
     y_train = input_data["y_train"]
     y_test = input_data["y_test"]
 
-    model = DecisionTreeClassifier(random_state=42)
+    model = DecisionTreeClassifier(
+        criterion=config.get("criterion", "gini"),
+        max_depth=config.get("max_depth", 5),
+        random_state=42
+    )
     model.fit(X_train, y_train)
 
     predictions = model.predict(X_test)

backend/app/pipeline/models/logistic_regression.py

@@ -7,7 +7,11 @@ def train(input_data, config):
     y_train = input_data["y_train"]
     y_test = input_data["y_test"]
 
-    model = LogisticRegression(max_iter=1000)
+    model = LogisticRegression(
+        C=config.get("C", 1.0),
+        solver=config.get("solver", "lbfgs"),
+        max_iter=1000
+    )
 
     model.fit(X_train, y_train)
 

backend/app/pipeline/models/random_forest.py

@@ -7,7 +7,11 @@ def train(input_data, config):
     y_train = input_data["y_train"]
     y_test = input_data["y_test"]
 
-    model = RandomForestClassifier(random_state=42)
+    model = RandomForestClassifier(
+        n_estimators=config.get("n_estimators", 100),
+        max_depth=config.get("max_depth", 5),
+        random_state=42
+    )
     model.fit(X_train, y_train)
 
     predictions = model.predict(X_test)

backend/app/pipeline/nodes/preprocess_node.py

@@ -1,15 +1,32 @@
+from sklearn.preprocessing import StandardScaler, MinMaxScaler, RobustScaler
+
+SCALER_REGISTRY = {
+    "standard": StandardScaler,
+    "minmax": MinMaxScaler,
+    "robust": RobustScaler
+}
+
 def run(input_data, config):
-    scale_factor = config.get("scale_factor", 1)
-    X_train = input_data["X_train"]
-    X_test = input_data["X_test"]
+    scaler_type = config.get("scaler_type", "standard")
+    if scaler_type not in SCALER_REGISTRY:
+        raise ValueError(f"Unknown scaler type: {scaler_type}")
+
+    X_train = input_data.get("X_train")
+    X_test = input_data.get("X_test")
+
+    scaler_class = SCALER_REGISTRY[scaler_type]
+    scaler = scaler_class()
+
+    scaled_X_train = scaler.fit_transform(X_train)
+    scaled_X_test = scaler.transform(X_test)
 
-    scaled_X_train = [[x * scale_factor for x in row] for row in X_train]
-    scaled_X_test = [[x * scale_factor for x in row] for row in X_test]
-
     return {
         "X_train": scaled_X_train,
         "X_test": scaled_X_test,
         "y_train": input_data["y_train"],
         "y_test": input_data["y_test"],
-        "task_type": input_data["task_type"]
+        "task_type": input_data["task_type"],
+        "preprocessing": {
+            "scaler_type": scaler_type
+        }
     }

backend/app/pipeline/registry/node_registry.py

@@ -46,9 +46,10 @@
             "inputs": ["X_train", "X_test", "y_train", "y_test"],
             "outputs": ["X_train", "X_test", "y_train", "y_test"],
             "config_schema": {
-                "scale_factor": {
-                    "type": "integer",
-                    "default": 1
+                "scaler_type": {
+                    "type": "string",
+                    "options": ["standard", "minmax", "robust"],
+                    "default": "standard"
                 }
             }
         }
@@ -65,10 +66,33 @@
                 "algorithm": {
                     "type": "string",
                     "options": ["linear_regression", "logistic_regression", "decision_tree", "random_forest"],
+                    "default": "logistic_regression"
                 },
                 "fit_intercept": {
                     "type": "boolean",
                     "default": True
+                },
+                "max_depth": {
+                    "type": "integer",
+                    "default": 5
+                },
+                "n_estimators": {
+                    "type": "integer",
+                    "default": 100
+                },
+                "criterion": {
+                    "type": "string",
+                    "options": ["gini", "entropy"],
+                    "default": "gini"
+                },
+                "C": {
+                    "type": "float",
+                    "default": 1.0
+                },
+                "solver": {
+                    "type": "string",
+                    "options": ["lbfgs", "liblinear"],
+                    "default": "lbfgs"
                 }
             }
         }