Update app.py

app.py

@@ -359,21 +359,25 @@ with gr.Blocks() as demo:
             agent_btn = gr.Button("Run AI Agent (5 Insights + 5 Visualizations)")
 
     with gr.Row():
+        ab_dropdown = gr.Dropdown(choices=["A", "B"], label="Choose Model Variant", value="A")
         train_btn = gr.Button("Train Model with Optuna + WandB")
+
+    with gr.Row():
         metrics_output = gr.JSON(label="Performance Metrics")
         trials_output = gr.DataFrame(label="Top 7 Hyperparameter Trials")
 
+    with gr.Row():
+        conf_matrix_img = gr.Image(label="Confusion Matrix")
+        metric_plot_img = gr.Image(label="Metric Bar Plot")
+
     with gr.Row():
         explain_btn = gr.Button("SHAP + LIME Explainability")
         shap_img = gr.Image(label="SHAP Summary Plot")
         lime_img = gr.Image(label="LIME Explanation")
-        metrics_output = gr.JSON(label="Evaluation Metrics")
-        trials_output = gr.Dataframe(label="Classification Report")
-        conf_matrix_img = gr.Image(label="Confusion Matrix")
-        metric_plot_img = gr.Image(label="Metric Bar Plot")
 
+    # Button logic
     agent_btn.click(fn=analyze_data, inputs=[file_input], outputs=[insights_output, visual_output])
     train_btn.click(fn=train_model, inputs=[file_input, ab_dropdown], outputs=[metrics_output, trials_output, conf_matrix_img, metric_plot_img])
     explain_btn.click(fn=explainability, inputs=[], outputs=[shap_img, lime_img])
 
-demo.launch(debug=True)
+demo.launch(debug=True)