Update app.py

app.py

@@ -246,33 +246,69 @@ def build_autoencoder(input_shape, encoding_dim, layers_config, autoencoder_type
     autoencoder.compile(optimizer=optimizer, loss='mse', metrics=['mse'])
     return autoencoder, encoder, decoder
 
-def train_model(model, X_train, y_train, X_test, y_test, epochs, batch_size, problem_type, input_data=None, target_data=None, do_grid_search=False, params=None, grid_params=None, training_placeholder=None):
+class StreamlitCallback(keras.callbacks.Callback):
+    def __init__(self, placeholder):
+        super().__init__()
+        self.placeholder = placeholder
+        self.epoch_data = []
+
+    def on_epoch_end(self, epoch, logs=None):
+        # Append the logs for the current epoch
+        self.epoch_data.append(logs)
+        
+        # Create a DataFrame from the logs
+        df = pd.DataFrame(self.epoch_data)
+        
+        # Create the Plotly figure
+        fig = go.Figure()
+        
+        # Add training loss trace
+        fig.add_trace(go.Scatter(
+            x=df.index, y=df['loss'], mode='lines', name='Training Loss'
+        ))
+        
+        # Add validation loss trace (if available)
+        if 'val_loss' in df.columns:
+            fig.add_trace(go.Scatter(
+                x=df.index, y=df['val_loss'], mode='lines', name='Validation Loss'
+            ))
+        
+        # Add metric trace (e.g., accuracy or MSE)
+        metric_name = 'accuracy' if 'accuracy' in df.columns else 'mse'
+        if metric_name in df.columns:
+            fig.add_trace(go.Scatter(
+                x=df.index, y=df[metric_name], mode='lines', name=metric_name.capitalize()
+            ))
+        
+        # Update the layout
+        fig.update_layout(
+            title="Training Progress",
+            xaxis_title="Epoch",
+            yaxis_title="Value",
+            legend_title="Metrics"
+        )
+        
+        # Update the placeholder with the new figure
+        self.placeholder.plotly_chart(fig, use_container_width=True)
+
+def train_model(model, X_train, y_train, X_test, y_test, epochs, batch_size, problem_type, input_data=None, target_data=None, do_grid_search=False, params=None, grid_params=None):
     start_time = time.time()
     history = None
-    if input_data is None:
-        input_data = X_train
-    if target_data is None:
-        target_data = y_train if y_train is not None else X_train
+    
+    # Create a placeholder for the training progress chart
+    training_placeholder = st.empty()
+    
     if isinstance(model, keras.Model):
-        class StreamlitCallback(keras.callbacks.Callback):
-            def __init__(self, placeholder):
-                super().__init__()
-                self.placeholder = placeholder
-                self.epoch_data = []
-
-            def on_epoch_end(self, epoch, logs=None):
-                self.epoch_data.append(logs)
-                df = pd.DataFrame(self.epoch_data)
-                fig = px.line(df, x=df.index, y=['loss', 'val_loss'], labels={'index': 'Epoch', 'value': 'Loss'})
-                metric_name = 'mse' if 'mse' in logs else 'accuracy'
-                fig.add_trace(go.Scatter(x=df.index, y=df[metric_name], mode='lines', name=metric_name))
-                fig.add_trace(go.Scatter(x=df.index, y=df[f'val_{metric_name}'], mode='lines', name=f'val_{metric_name}'))
-                self.placeholder.plotly_chart(fig)
-
         streamlit_callback = StreamlitCallback(training_placeholder)
-        history = model.fit(input_data, target_data, epochs=epochs, batch_size=batch_size,
-                            validation_data=(X_test, y_test if y_test is not None else X_test), verbose=0,
-                            callbacks=[streamlit_callback])
+        history = model.fit(
+            input_data if input_data is not None else X_train,
+            target_data if target_data is not None else y_train,
+            epochs=epochs,
+            batch_size=batch_size,
+            validation_data=(X_test, y_test if y_test is not None else X_test),
+            verbose=0,
+            callbacks=[streamlit_callback]
+        )
     else:
         if do_grid_search and grid_params:
             grid_search = GridSearchCV(model, grid_params, cv=3, n_jobs=-1, scoring='accuracy' if problem_type in ["Binary Classification", "Multi-Class"] else 'neg_mean_squared_error')
@@ -282,6 +318,7 @@ def train_model(model, X_train, y_train, X_test, y_test, epochs, batch_size, pro
         else:
             model.set_params(**params)
             model.fit(X_train, y_train)
+    
     training_time = time.time() - start_time
     return history, model, training_time