diff --git a/backend/app/pipeline/generators/code_generator.py b/backend/app/pipeline/generators/code_generator.py index d9d8926..6ed63fa 100644 --- a/backend/app/pipeline/generators/code_generator.py +++ b/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) \ No newline at end of file + return "\n".join(final_code) \ No newline at end of file diff --git a/backend/app/pipeline/generators/code_registry.py b/backend/app/pipeline/generators/code_registry.py new file mode 100644 index 0000000..3742fea --- /dev/null +++ b/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 +} \ No newline at end of file diff --git a/backend/app/pipeline/generators/dataset_generator.py b/backend/app/pipeline/generators/dataset_generator.py new file mode 100644 index 0000000..3cbd689 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/backend/app/pipeline/generators/model_generator.py b/backend/app/pipeline/generators/model_generator.py new file mode 100644 index 0000000..e2b7fec --- /dev/null +++ b/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 \ No newline at end of file diff --git a/backend/app/pipeline/generators/preprocess_generator.py b/backend/app/pipeline/generators/preprocess_generator.py new file mode 100644 index 0000000..6f780f9 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/backend/app/pipeline/generators/split_generator.py b/backend/app/pipeline/generators/split_generator.py new file mode 100644 index 0000000..feb093a --- /dev/null +++ b/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 \ No newline at end of file diff --git a/backend/app/pipeline/models/decision_tree.py b/backend/app/pipeline/models/decision_tree.py index 8a8f8a1..7c6a7fb 100644 --- a/backend/app/pipeline/models/decision_tree.py +++ b/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) diff --git a/backend/app/pipeline/models/logistic_regression.py b/backend/app/pipeline/models/logistic_regression.py index 5729951..a83d5d6 100644 --- a/backend/app/pipeline/models/logistic_regression.py +++ b/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) diff --git a/backend/app/pipeline/models/random_forest.py b/backend/app/pipeline/models/random_forest.py index 24a8ef6..c314de3 100644 --- a/backend/app/pipeline/models/random_forest.py +++ b/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) diff --git a/backend/app/pipeline/nodes/preprocess_node.py b/backend/app/pipeline/nodes/preprocess_node.py index 5304996..556525c 100644 --- a/backend/app/pipeline/nodes/preprocess_node.py +++ b/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 + } } \ No newline at end of file diff --git a/backend/app/pipeline/registry/node_registry.py b/backend/app/pipeline/registry/node_registry.py index df6b05f..7dd7fc0 100644 --- a/backend/app/pipeline/registry/node_registry.py +++ b/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" } } } diff --git a/frontend/src/components/ConfigPanel.tsx b/frontend/src/components/ConfigPanel.tsx index 1b8ff37..21902b4 100644 --- a/frontend/src/components/ConfigPanel.tsx +++ b/frontend/src/components/ConfigPanel.tsx @@ -1,5 +1,6 @@ "use client"; import type { Node } from "@xyflow/react"; +import { useEffect } from "react"; import { usePipelineStore } from "../store/pipelineStore"; interface ConfigPanelProps { @@ -14,6 +15,48 @@ type ConfigFieldSchema = { type ConfigSchemaMap = Record; +const MODEL_FIELD_MAP: Record = { + linear_regression: ["fit_intercept"], + logistic_regression: ["C", "solver"], + decision_tree: ["criterion", "max_depth"], + random_forest: ["n_estimators", "max_depth"], +}; + +const MODEL_DEFAULT_CONFIGS: Record> = { + linear_regression: { fit_intercept: true }, + logistic_regression: { C: 1.0, solver: "lbfgs" }, + decision_tree: { criterion: "gini", max_depth: 5 }, + random_forest: { n_estimators: 100, max_depth: 5 }, +}; + +const buildModelConfig = ( + algorithm: string, + existingConfig: Record, + options?: { useExistingValues?: boolean }, +) => { + const allowedFields = MODEL_FIELD_MAP[algorithm] ?? []; + const defaults = MODEL_DEFAULT_CONFIGS[algorithm] ?? {}; + const nextConfig: Record = { algorithm }; + + allowedFields.forEach((field) => { + if (options?.useExistingValues && field in existingConfig) { + nextConfig[field] = existingConfig[field]; + return; + } + if (field in defaults) { + nextConfig[field] = defaults[field]; + } + }); + + return nextConfig; +}; + +const areConfigsEqual = (left: Record, right: Record) => { + const leftKeys = Object.keys(left); + if (leftKeys.length !== Object.keys(right).length) return false; + return leftKeys.every((key) => Object.is(left[key], right[key])); +}; + export default function ConfigPanel({ nodeMetadata }: ConfigPanelProps) { const { nodes, selectedNodeId, updateNodeConfig } = usePipelineStore(); const selectedNode = nodes.find((node: Node) => node.id === selectedNodeId); @@ -23,17 +66,49 @@ export default function ConfigPanel({ nodeMetadata }: ConfigPanelProps) { const configSchema = (selectedNodeMetadata as { config_schema?: ConfigSchemaMap } | null)?.config_schema; const nodeConfig = ((selectedNode?.data as { config?: Record } | undefined)?.config as Record) ?? {}; + const isModelNode = selectedNode?.type === "model"; + const modelAlgorithm = isModelNode + ? (nodeConfig.algorithm as string | undefined) ?? + ((configSchema?.algorithm as ConfigFieldSchema | undefined)?.default as string | undefined) + : undefined; + const filteredConfigEntries = configSchema + ? Object.entries(configSchema).filter(([key]) => { + if (!isModelNode) return true; + if (key === "algorithm") return true; + const allowedFields = modelAlgorithm ? MODEL_FIELD_MAP[modelAlgorithm] ?? [] : []; + return allowedFields.includes(key); + }) + : []; + const orderedConfigEntries = isModelNode + ? [...filteredConfigEntries].sort(([leftKey], [rightKey]) => { + if (leftKey === "algorithm") return -1; + if (rightKey === "algorithm") return 1; + return 0; + }) + : filteredConfigEntries; const resolvedConfig = configSchema ? Object.fromEntries( - Object.entries(configSchema).map(([key, field]) => { + orderedConfigEntries.map(([key, field]) => { const typedField = field as ConfigFieldSchema; return [key, nodeConfig[key] ?? typedField.default]; }) ) : nodeConfig; + useEffect(() => { + if (!selectedNode || !isModelNode || !modelAlgorithm) return; + const nextConfig = buildModelConfig(modelAlgorithm, nodeConfig, { useExistingValues: true }); + if (!areConfigsEqual(nextConfig, nodeConfig)) { + updateNodeConfig(selectedNode.id, nextConfig, { replace: true }); + } + }, [selectedNode, isModelNode, modelAlgorithm, nodeConfig, updateNodeConfig]); + const handleConfigChange = (key: string, value: unknown) => { if (!selectedNode) return; + if (isModelNode && key === "algorithm" && typeof value === "string") { + updateNodeConfig(selectedNode.id, buildModelConfig(value, {}, { useExistingValues: false }), { replace: true }); + return; + } updateNodeConfig(selectedNode.id, { [key]: value }); }; @@ -58,7 +133,7 @@ export default function ConfigPanel({ nodeMetadata }: ConfigPanelProps) {

Configuration Fields

- {configSchema && Object.entries(configSchema).map(([key, field]) => { + {configSchema && orderedConfigEntries.map(([key, field]) => { const typedField = field as ConfigFieldSchema; const resolvedValue = nodeConfig[key] ?? typedField.default; const fieldType = typedField.type ?? "string"; diff --git a/frontend/src/store/pipelineStore.ts b/frontend/src/store/pipelineStore.ts index 1a5b121..ec53767 100644 --- a/frontend/src/store/pipelineStore.ts +++ b/frontend/src/store/pipelineStore.ts @@ -27,7 +27,11 @@ interface PipelineStore { addNode: (type: string, position: { x: number; y: number }) => void; setSelectedNode: (nodeId: string | null) => void; setSelectedEdge: (edgeId: string | null) => void; - updateNodeConfig: (nodeId: string, config: Record) => void; + updateNodeConfig: ( + nodeId: string, + config: Record, + options?: { replace?: boolean }, + ) => void; deleteNode: (nodeId: string) => void; deleteEdge: (edgeId: string) => void; onNodeDragStop: ( @@ -89,7 +93,7 @@ export const usePipelineStore = create((set) => ({ setSelectedEdge: (edgeId) => set(() => ({ selectedEdgeId: edgeId, selectedNodeId: null })), - updateNodeConfig: (nodeId, config) => + updateNodeConfig: (nodeId, config, options) => set((state) => ({ nodes: state.nodes.map((node) => node.id === nodeId @@ -97,13 +101,15 @@ export const usePipelineStore = create((set) => ({ ...node, data: { ...node.data, - config: { - ...(((node.data as Record).config as Record< - string, - unknown - >) ?? {}), - ...config, - }, + config: options?.replace + ? config + : { + ...(((node.data as Record).config as Record< + string, + unknown + >) ?? {}), + ...config, + }, }, } : node,