Update app.py

app.py

@@ -14,18 +14,21 @@ if "epoch" not in st.session_state:
     st.session_state.epoch = 0
 if "running" not in st.session_state:
     st.session_state.running = False
-if "loss_history" not in st.session_state:
-    st.session_state.loss_history = []
+if "train_loss_history" not in st.session_state:
+    st.session_state.train_loss_history = []
+if "test_loss_history" not in st.session_state:
+    st.session_state.test_loss_history = []
 
 # ================= TRAINING CONTROL PANEL (Top) =================
 st.markdown("### Training Controls")
-col1, col2, col3, col4, col5, col6 = st.columns(6)
+col1, col2, col3, col4, col5, col6, col7 = st.columns(7)
 
 with col1:
     if st.button("↩️ Reset"):
         st.session_state.epoch = 0
         st.session_state.running = False
-        st.session_state.loss_history = []
+        st.session_state.train_loss_history = []
+        st.session_state.test_loss_history = []
 with col2:
     if st.button("▶️ Train"):
         st.session_state.running = True
@@ -38,11 +41,14 @@ with col5:
     problem_type = st.selectbox("Problem Type", ["Classification", "Regression"])
 with col6:
     learning_rate = st.selectbox("Learning Rate", [0.0001, 0.001, 0.01, 0.03, 0.1])
+with col7:
+    num_epochs = st.slider("Epochs", 1, 100, 10)
 
 if st.session_state.running:
     time.sleep(1)
     st.session_state.epoch += 1
-    st.session_state.loss_history.append(np.exp(-0.1 * st.session_state.epoch) + np.random.uniform(0, 0.05))
+    st.session_state.train_loss_history.append(np.exp(-0.1 * st.session_state.epoch) + np.random.uniform(0, 0.05))
+    st.session_state.test_loss_history.append(np.exp(-0.1 * st.session_state.epoch) + np.random.uniform(0, 0.05) + 0.02)
 
 # ================= MAIN LAYOUT =================
 col_features, col_hidden, col_plot = st.columns([2, 2, 3])
@@ -74,10 +80,12 @@ with col_plot:
 
     # Loss Plot
     fig, ax = plt.subplots()
-    ax.plot(range(len(st.session_state.loss_history)), st.session_state.loss_history, marker="o")
+    ax.plot(range(len(st.session_state.train_loss_history)), st.session_state.train_loss_history, marker="o", label="Train Loss")
+    ax.plot(range(len(st.session_state.test_loss_history)), st.session_state.test_loss_history, marker="s", label="Test Loss")
     ax.set_title("Epoch vs. Loss")
     ax.set_xlabel("Epoch")
     ax.set_ylabel("Loss")
+    ax.legend()
     st.pyplot(fig)
 
 # ================= DATASET SELECTION (Sidebar) =================