Update app.py

app.py

@@ -6,10 +6,12 @@ import graphviz
 import time
 import tensorflow as tf
 from tensorflow import keras
-from tensorflow.keras import layers
+from tensorflow.keras import layers, Sequential
+from tensorflow.keras.regularizers import l2
 from sklearn.model_selection import train_test_split
 from sklearn.datasets import make_moons, make_circles, make_classification, make_blobs
 from mlxtend.plotting import plot_decision_regions
+from sklearn.base import BaseEstimator, ClassifierMixin
 
 # Set Streamlit page title
 st.set_page_config(page_title="Neural Network Trainer", layout="wide")
@@ -41,7 +43,7 @@ with col3:
     if st.button("⏸️ Pause"):
         st.session_state.running = False
 with col4:
-    activation = st.selectbox("Activation", ["ReLU", "Sigmoid", "Tanh"])
+    activation = st.selectbox("Activation", ["ReLU", "Sigmoid", "Tanh"], index=2)
 with col5:
     problem_type = st.selectbox("Problem Type", ["Classification", "Regression"])
 with col6:
@@ -87,19 +89,29 @@ X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_
 
 # ========== TRAINING ANN ========== #
 def build_ann():
-    model = keras.Sequential()
+    model = Sequential()
     model.add(layers.Input(shape=(X.shape[1],)))
     for units in neurons:
-        model.add(layers.Dense(units, activation=activation.lower()))
-    model.add(layers.Dense(1, activation="sigmoid" if problem_type == "Classification" else "linear"))
+        model.add(layers.Dense(units=units, activation='tanh'))
+    model.add(layers.Dense(units=1, activation='sigmoid', kernel_regularizer=l2(0.1)))
     model.compile(optimizer=keras.optimizers.Adam(learning_rate), loss="binary_crossentropy" if problem_type == "Classification" else "mse")
     return model
 
+class KerasClassifierWrapper(BaseEstimator, ClassifierMixin):
+    def __init__(self, model):
+        self.model = model
+    def fit(self, X, y):
+        self.model.fit(X, y, epochs=num_epochs, batch_size=batch_size, verbose=0)
+        return self
+    def predict(self, X):
+        return (self.model.predict(X) > 0.5).astype(int).flatten()
+
 if st.session_state.running:
     model = build_ann()
     history = model.fit(X_train, y_train, epochs=num_epochs, batch_size=batch_size, validation_data=(X_test, y_test), verbose=0)
     st.session_state.train_loss_history = history.history["loss"]
     st.session_state.test_loss_history = history.history["val_loss"]
+    wrapper_model = KerasClassifierWrapper(model)
 
 # ========== LOSS PLOT ========== #
 with col_plot:
@@ -116,7 +128,8 @@ with col_plot:
 # =================== DECISION REGION =================== #
 if problem_type == "Classification":
     fig, ax = plt.subplots()
-    plot_decision_regions(X_train, y_train, clf=model, ax=ax)
+    plot_decision_regions(X_train, y_train, clf=wrapper_model, ax=ax)
+    ax.scatter(X_train[:, 0], X_train[:, 1], c=y_train, edgecolor='k')
     ax.set_title("Decision Region")
     st.pyplot(fig)