Update app.py

app.py

@@ -246,7 +246,7 @@ with col_center:
             st.session_state.hidden_layer_neurons.pop()
 
     def increase_neurons(i):
-        if st.session_state.hidden_layer_neurons[i] < 8:
+        if st.session_state.num_hidden_layers[i] < 8:
             st.session_state.hidden_layer_neurons[i] += 1
 
     def decrease_neurons(i):
@@ -258,7 +258,7 @@ with col_center:
         with col1:
             st.button("−", key=f"dec_{i}", on_click=decrease_neurons, args=(i,))
         with col2:
-            st.write(f"Layer {i+1}: {st.session_state.hidden_layer_neurons[i]}")
+            st.write(f"Layer {i+1}: {st.session_state.hidden_layer_neurons[i]} neurons")
         with col3:
             st.button("+", key=f"inc_{i}", on_click=increase_neurons, args=(i,))
     col_btn1, col_btn2 = st.columns(2)
@@ -298,6 +298,10 @@ with col_right:
             self.X, self.y = X, y
             self.losses = {"Epoch": [], "Train Loss": [], "Val Loss": []}
             self.placeholder = st.empty()
+            self.model = None  # Store the model instance
+
+        def on_train_begin(self, logs=None):
+            self.model = self.model  # Ensure model is stored
 
         def on_epoch_end(self, epoch, logs=None):
             self.losses["Epoch"].append(epoch + 1)
@@ -308,13 +312,14 @@ with col_right:
                 with col1:
                     plt.figure(figsize=(3, 3))
                     if problem_type == "Classification":
-                        # Ensure we use only the first two features for 2D plotting
-                        X_2d = self.X[:, :2]  # Use only X1 and X2 for decision boundary
+                        # Use only the first two features for 2D plotting
+                        X_2d = self.X[:, :2]  # Ensure 2D input
                         y_pred_proba = self.model.predict(X_2d, verbose=0)
                         y_pred = (y_pred_proba > 0.5).astype(int).ravel()
                         try:
                             plot_decision_regions(X_2d, self.y, clf=self.model, legend=2, colors='blue,red')
                             plt.scatter(X_2d[:, 0], X_2d[:, 1], c=self.y, cmap='coolwarm', edgecolors='k', alpha=0.7)
+                            plt.contour(X_2d[:, 0], X_2d[:, 1], y_pred.reshape(X_2d.shape[0], X_2d.shape[1]), levels=[0.5], colors='black', linewidths=2)
                         except Exception as e:
                             st.warning(f"Decision region plot failed: {e}")
                             xx, yy = np.meshgrid(np.linspace(X_2d[:, 0].min(), X_2d[:, 0].max(), 100),
@@ -328,7 +333,7 @@ with col_right:
                     else:
                         y_pred = self.model.predict(self.X, verbose=0)
                         plt.scatter(self.X[:, 0], self.y, c="blue", alpha=0.5)
-                        plt.plot(self.X[:, 0], y_pred, "r-", linewidth=2)
+                        plt.plot(self.X[:, 0], y_pred, "r-", linewidths=2)
                     plt.gca().set_facecolor("#333")
                     plt.xticks([])
                     plt.yticks([])