Update app.py

app.py

@@ -68,7 +68,7 @@ def get_model_config(model_type, problem_type):
         }
     }
     return configs.get(model_type, {}).get(problem_type, {"model_class": None, "params": {}, "grid_params": {}})
-
+    
 def preprocess_data(X_train, X_test, numerical_features, categorical_features):
     # Define the numeric and categorical transformers
     numeric_transformer = Pipeline(steps=[
@@ -102,24 +102,45 @@ def preprocess_data(X_train, X_test, numerical_features, categorical_features):
         feature_names = numerical_features
 
     return X_train_processed, X_test_processed, feature_names, preprocessor
+def build_neural_network(input_shape, output_units, problem_type, layers_config, optimizer_name="Adam", loss_function="mse", metrics=["accuracy"]):
+    try:
+        model = keras.Sequential()
+        model.add(keras.layers.InputLayer(input_shape=input_shape))
+        for layer in layers_config:
+            if layer['type'] == 'dense':
+                model.add(keras.layers.Dense(layer['units'], activation=layer['activation']))
+            elif layer['type'] == 'dropout':
+                model.add(keras.layers.Dropout(layer['rate']))
+            elif layer['type'] == 'conv2d':
+                model.add(keras.layers.Conv2D(layer['filters'], tuple(layer['kernel_size']), activation=layer['activation']))
+            elif layer['type'] == 'lstm':
+                model.add(keras.layers.LSTM(layer['units'], activation=layer['activation']))
+            elif layer['type'] == 'maxpooling2d':
+                model.add(keras.layers.MaxPooling2D(pool_size=tuple(layer['pool_size'])))
+            elif layer['type'] == 'flatten':
+                model.add(keras.layers.Flatten())
+
+        if problem_type == "Regression":
+            model.add(keras.layers.Dense(1))
+        elif problem_type == "Binary Classification":
+            model.add(keras.layers.Dense(1, activation='sigmoid'))
+        elif problem_type == "Multi-Class":
+            model.add(keras.layers.Dense(output_units, activation='softmax'))
+
+        if optimizer_name == "Adam":
+            optimizer = keras.optimizers.Adam()
+        elif optimizer_name == "SGD":
+            optimizer = keras.optimizers.SGD()
+        elif optimizer_name == "RMSprop":
+            optimizer = keras.optimizers.RMSprop()
+
+        model.compile(optimizer=optimizer, loss=loss_function, metrics=metrics)
+        return model
+    except Exception as e:
+        st.error("Failed to build neural network. Check Debug Log for details.")
+        #log_error("Error in build_neural_network", e) # ADDED
+        raise
 
-def build_neural_network(input_shape, output_units, problem_type, layers_config, optimizer_name="Adam", learning_rate=0.001):
-    model = keras.Sequential([keras.layers.Input(shape=input_shape)])
-    for layer in layers_config:
-        if layer['type'] == 'dense':
-            model.add(keras.layers.Dense(layer['units'], activation=layer['activation']))
-        elif layer['type'] == 'dropout':
-            model.add(keras.layers.Dropout(layer['rate']))
-    if problem_type == "Regression":
-        model.add(keras.layers.Dense(1))
-    elif problem_type == "Binary Classification":
-        model.add(keras.layers.Dense(1, activation='sigmoid'))
-    elif problem_type == "Multi-Class":
-        model.add(keras.layers.Dense(output_units, activation='softmax'))
-    optimizer = {"Adam": keras.optimizers.Adam, "SGD": keras.optimizers.SGD, "RMSprop": keras.optimizers.RMSprop}[optimizer_name](learning_rate=learning_rate)
-    model.compile(optimizer=optimizer, loss={"Regression": "mse", "Binary Classification": "binary_crossentropy", "Multi-Class": "categorical_crossentropy"}[problem_type],
-                  metrics=["mae" if problem_type == "Regression" else "accuracy"])
-    return model
 
 def train_model(model, X_train, y_train, X_test, y_test, epochs, batch_size, problem_type, do_grid_search=False, params=None, grid_params=None, training_placeholder=None):
     start_time = time.time()