add visualisation elements

app.py

@@ -1,23 +1,22 @@
 import gradio as gr
 import numpy as np
+import matplotlib.pyplot as plt
 from sklearn.datasets import load_iris
 from sklearn.ensemble import GradientBoostingClassifier
 from sklearn.model_selection import train_test_split
 from sklearn.metrics import accuracy_score, confusion_matrix
 
-# 1. Load dataset
 iris = load_iris()
 X, y = iris.data, iris.target
 feature_names = iris.feature_names
 class_names = iris.target_names
 
-# Split into train/test
 X_train, X_test, y_train, y_test = train_test_split(
     X, y, test_size=0.3, random_state=42
 )
 
-# 2. Define a function that trains & evaluates a model given hyperparameters
 def train_and_evaluate(learning_rate, n_estimators, max_depth):
+    # Train model
     clf = GradientBoostingClassifier(
         learning_rate=learning_rate,
         n_estimators=n_estimators,
@@ -25,15 +24,31 @@ def train_and_evaluate(learning_rate, n_estimators, max_depth):
         random_state=42
     )
     clf.fit(X_train, y_train)
-    y_pred = clf.predict(X_test)
 
+    # Predict and compute metrics
+    y_pred = clf.predict(X_test)
     accuracy = accuracy_score(y_test, y_pred)
     cm = confusion_matrix(y_test, y_pred)
+
+    # Convert confusion matrix to readable string
     cm_display = "\n".join([str(row) for row in cm])
 
-    return f"Accuracy: {accuracy:.3f}\nConfusion Matrix:\n{cm_display}"
+    # Create a feature importance bar chart
+    importances = clf.feature_importances_
+    fig, ax = plt.subplots()
+    ax.barh(range(len(feature_names)), importances, color='skyblue')
+    ax.set_yticks(range(len(feature_names)))
+    ax.set_yticklabels(feature_names)
+    ax.set_xlabel("Importance")
+    ax.set_title("Feature Importances (Gradient Boosting)")
+
+    # Convert the Matplotlib figure to a Gradio-readable format
+    # (returns a temporary .png file path)
+    return (
+        f"Accuracy: {accuracy:.3f}\nConfusion Matrix:\n{cm_display}",
+        fig
+    )
 
-# 3. Define a prediction function for user-supplied feature values
 def predict_species(sepal_length, sepal_width, petal_length, petal_width,
                     learning_rate, n_estimators, max_depth):
     clf = GradientBoostingClassifier(
@@ -43,12 +58,10 @@ def predict_species(sepal_length, sepal_width, petal_length, petal_width,
         random_state=42
     )
     clf.fit(X_train, y_train)
-
     user_sample = np.array([[sepal_length, sepal_width, petal_length, petal_width]])
     prediction = clf.predict(user_sample)[0]
     return f"Predicted species: {class_names[prediction]}"
 
-# 4. Build the Gradio interface
 with gr.Blocks() as demo:
     with gr.Tab("Train & Evaluate"):
         gr.Markdown("## Train a GradientBoostingClassifier on the Iris dataset")
@@ -58,11 +71,12 @@ with gr.Blocks() as demo:
 
         train_button = gr.Button("Train & Evaluate")
         output_text = gr.Textbox(label="Results")
+        output_plot = gr.Plot(label="Feature Importance")
 
         train_button.click(
             fn=train_and_evaluate,
             inputs=[learning_rate_slider, n_estimators_slider, max_depth_slider],
-            outputs=output_text,
+            outputs=[output_text, output_plot],
         )
 
     with gr.Tab("Predict"):
@@ -71,8 +85,7 @@ with gr.Blocks() as demo:
         sepal_width_input = gr.Number(value=3.5, label=feature_names[1])
         petal_length_input = gr.Number(value=1.4, label=feature_names[2])
         petal_width_input = gr.Number(value=0.2, label=feature_names[3])
-        
-        # Hyperparams for the model that will do the prediction
+
         learning_rate_slider2 = gr.Slider(0.01, 1.0, value=0.1, step=0.01, label="learning_rate")
         n_estimators_slider2 = gr.Slider(50, 300, value=100, step=50, label="n_estimators")
         max_depth_slider2 = gr.Slider(1, 10, value=3, step=1, label="max_depth")