Update app.py

app.py

@@ -6,7 +6,6 @@ from tensorflow.keras.utils import plot_model
 import matplotlib.pyplot as plt
 from io import BytesIO
 
-# Main app title
 st.title("TensorFlow Neural Network Playground")
 
 # Sidebar for network configuration
@@ -29,17 +28,15 @@ def create_model(input_dim, hidden_dim, output_dim, lr):
 
 # Visualize the model architecture using Keras plot_model
 def plot_network(model):
-    # Generate the plot in memory
     img_data = BytesIO()
     plot_model(model, 
               to_file=img_data, 
               show_shapes=True, 
               show_layer_names=True,
-              rankdir='TB',  # Top to Bottom
+              rankdir='TB',
               expand_nested=True,
               dpi=96)
-    
-    # Display the image in Streamlit
+    img_data.seek(0)  # Reset buffer position to start
     st.image(img_data.getvalue(), caption="Neural Network Architecture")
 
 # Create and display the model structure
@@ -54,54 +51,36 @@ def generate_sample_data(samples=100):
     y = tf.keras.utils.to_categorical(y, output_nodes)
     return X, y
 
-# Train model button
 if st.button("Train Model"):
     X, y = generate_sample_data()
-    
-    # Training
     history = model.fit(X, y, epochs=10, verbose=0)
     
-    # Display results
     st.write("Training Complete!")
     fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
-    
-    # Plot accuracy
     ax1.plot(history.history['accuracy'])
     ax1.set_title('Model Accuracy')
     ax1.set_ylabel('Accuracy')
     ax1.set_xlabel('Epoch')
-    
-    # Plot loss
     ax2.plot(history.history['loss'])
     ax2.set_title('Model Loss')
     ax2.set_ylabel('Loss')
     ax2.set_xlabel('Epoch')
-    
     st.pyplot(fig)
 
-# Enhanced model summary with styling
 if st.checkbox("Show Model Summary"):
     st.subheader("Model Summary")
     summary_str = []
     model.summary(print_fn=lambda x: summary_str.append(x))
-    
-    # Format the summary in a more visual way
     st.markdown("### Model: sequential")
     st.markdown("""
     | Layer (type)         | Output Shape         | Param # |
     |----------------------|----------------------|---------|""")
-    
-    # Parse and display layers
-    for line in summary_str[1:-2]:  # Skip header and total params
+    for line in summary_str[1:-2]:
         if 'dense' in line.lower():
             parts = line.split()
             layer_name = parts[0] + " (Dense)"
             output_shape = parts[1]
             param_count = parts[2]
             st.markdown(f"| {layer_name:<20} | {output_shape:<20} | {param_count:>7} |")
-    
-    # Display total params
     total_params = summary_str[-1].split()[-1]
-    st.markdown(f"**Total params:** {total_params}")
-
-# Requirements.txt remains mostly the same, just ensure tensorflow is included
+    st.markdown(f"**Total params:** {total_params}")