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.num_hidden_layers[i] < 8:
+        if st.session_state.hidden_layer_neurons[i] < 8:
             st.session_state.hidden_layer_neurons[i] += 1
 
     def decrease_neurons(i):
@@ -298,15 +298,15 @@ 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
+            self.model = None  # Store model instance explicitly
 
         def on_train_begin(self, logs=None):
-            self.model = self.model  # Ensure model is stored
+            self.model = self.model_  # Assign the model from the fit method
 
         def on_epoch_end(self, epoch, logs=None):
             self.losses["Epoch"].append(epoch + 1)
             self.losses["Train Loss"].append(logs["loss"])
-            self.losses["Val Loss"].append(logs["val_loss"])
+            self.losses["Val Loss"].append(logs.get("val_loss", logs["loss"]))  # Handle missing val_loss
             with self.placeholder.container():
                 col1, col2 = st.columns(2)
                 with col1:
@@ -314,24 +314,30 @@ with col_right:
                     if problem_type == "Classification":
                         # 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:
+                            y_pred_proba = self.model.predict(X_2d, verbose=0)
+                            y_pred = (y_pred_proba > 0.5).astype(int).ravel()
                             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)
+                            # Add decision boundary contour
+                            xx, yy = np.meshgrid(np.linspace(X_2d[:, 0].min(), X_2d[:, 0].max(), 100),
+                                                 np.linspace(X_2d[:, 1].min(), X_2d[:, 1].max(), 100))
+                            grid = np.c_[xx.ravel(), yy.ravel()]
+                            Z = self.model.predict(grid, verbose=0)
+                            Z = (Z > 0.5).astype(int).reshape(xx.shape)
+                            plt.contour(xx, yy, Z, 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),
                                                  np.linspace(X_2d[:, 1].min(), X_2d[:, 1].max(), 100))
                             grid = np.c_[xx.ravel(), yy.ravel()]
-                            Z = self.model.predict(grid, verbose=0)
+                            Z = self.model.predict(grid, verbose=0) if self.model else np.zeros_like(grid[:, 0])
                             Z = (Z > 0.5).astype(int).reshape(xx.shape)
-                            plt.contour(xx, yy, Z, levels=[0.5], colors='black', linewidths=2)  # Clear decision boundary
+                            plt.contour(xx, yy, Z, levels=[0.5], colors='black', linewidths=2)
                             plt.contourf(xx, yy, Z, alpha=0.3, cmap="coolwarm")
                             plt.scatter(X_2d[:, 0], X_2d[:, 1], c=self.y, cmap="coolwarm", edgecolors="k", alpha=0.7)
                     else:
-                        y_pred = self.model.predict(self.X, verbose=0)
+                        y_pred = self.model.predict(self.X, verbose=0) if self.model else np.zeros_like(self.X[:, 0])
                         plt.scatter(self.X[:, 0], self.y, c="blue", alpha=0.5)
                         plt.plot(self.X[:, 0], y_pred, "r-", linewidths=2)
                     plt.gca().set_facecolor("#333")
@@ -347,7 +353,14 @@ with col_right:
                     st.pyplot(fig)
 
     if st.session_state.training:
-        model = create_model(len(selected_features), st.session_state.hidden_layer_neurons)
-        callback = OutputCallback(selected_data, cv)
-        model.fit(selected_data, cv, epochs=999999, batch_size=batch_size, validation_split=1-train_ratio, callbacks=[callback], verbose=0)
+        try:
+            model = create_model(len(selected_features), st.session_state.hidden_layer_neurons)
+            callback = OutputCallback(selected_data, cv)
+            callback.model_ = model  # Explicitly set the model for the callback
+            model.fit(selected_data, cv, epochs=100,  # Reduced epochs for Spaces
+                      batch_size=batch_size, validation_split=1-train_ratio, 
+                      callbacks=[callback], verbose=0)
+        except Exception as e:
+            st.error(f"Training failed: {e}")
+
     st.markdown('</div>', unsafe_allow_html=True)