add docs + improve ux/ui

app.py

@@ -1,5 +1,3 @@
-# app.py
-
 import gradio as gr
 import numpy as np
 import pandas as pd
@@ -11,49 +9,48 @@ from sklearn.ensemble import GradientBoostingClassifier
 from sklearn.model_selection import train_test_split
 from sklearn.metrics import accuracy_score, confusion_matrix
 
-# In some remote environments, Matplotlib needs to be set to 'Agg' backend
 matplotlib.use('Agg')
 
 ################################################################################
-# SUGGESTED_DATASETS: Must actually exist on huggingface.co/datasets.
+# SUGGESTED_DATASETS: These must actually exist on huggingface.co/datasets
 #
-# "scikit-learn/iris" -> a tabular Iris dataset with a "train" split of 150 rows.
-# "uci/wine"          -> a tabular Wine dataset with a "train" split of 178 rows.
+# "scikit-learn/iris" -> A small, classic Iris dataset with a "train" split
+# "uci/wine"          -> Another small dataset with a "train" split
+# "SKIP/ENTER_CUSTOM" -> Placeholder to let the user enter a custom dataset ID
 ################################################################################
 SUGGESTED_DATASETS = [
     "scikit-learn/iris",
     "uci/wine",
-    "SKIP/ENTER_CUSTOM"  # a placeholder meaning "use custom_dataset_id"
+    "SKIP/ENTER_CUSTOM"
 ]
 
-
 def update_columns(dataset_id, custom_dataset_id):
     """
-    Loads the chosen dataset (train split) and returns its column names,
-    to populate the Label Column & Feature Columns selectors.
+    After the user chooses a dataset from the dropdown or enters their own,
+    this function loads the dataset's "train" split, converts it to a DataFrame,
+    and returns the columns. These columns are used to populate the Label and
+    Feature selectors in the UI.
     """
-    # If user picked a suggested dataset (not SKIP), use that
     if dataset_id != "SKIP/ENTER_CUSTOM":
         final_id = dataset_id
     else:
-        # Use the user-supplied dataset ID
         final_id = custom_dataset_id.strip()
 
     try:
-        # Load just the "train" split; many HF datasets have train/test/validation
         ds = load_dataset(final_id, split="train")
         df = pd.DataFrame(ds)
         cols = df.columns.tolist()
 
-        message = f"**Loaded dataset**: {final_id}\n\n**Columns found**: {cols}"
-        # Return list of columns for both label & features
+        message = (
+            f"**Loaded dataset**: `{final_id}`\n\n"
+            f"**Columns found**: {cols}"
+        )
         return (
             gr.update(choices=cols, value=None),   # label_col dropdown
             gr.update(choices=cols, value=[]),     # feature_cols checkbox group
             message
         )
     except Exception as e:
-        # If load fails or dataset doesn't exist
         err_msg = f"**Error loading** `{final_id}`: {e}"
         return (
             gr.update(choices=[], value=None),
@@ -61,43 +58,42 @@ def update_columns(dataset_id, custom_dataset_id):
             err_msg
         )
 
-
 def train_model(dataset_id, custom_dataset_id, label_column, feature_columns,
                 learning_rate, n_estimators, max_depth, test_size):
     """
-    1. Determine the final dataset ID (from dropdown or custom text).
-    2. Load the dataset -> create dataframe -> X, y.
-    3. Train GradientBoostingClassifier.
-    4. Return metrics (accuracy) and a Matplotlib figure with:
-       - Feature importance bar chart
-       - Confusion matrix heatmap
+    1. Decide which dataset ID to load (from dropdown or custom).
+    2. Load that dataset's 'train' split, turn into DataFrame, extract X (features) and y (label).
+    3. Train a GradientBoostingClassifier on X_train, y_train.
+    4. Compute accuracy and confusion matrix on X_test, y_test.
+    5. Plot and return feature importances + confusion matrix heatmap + textual summary.
     """
+    # Resolve final dataset ID
     if dataset_id != "SKIP/ENTER_CUSTOM":
         final_id = dataset_id
     else:
         final_id = custom_dataset_id.strip()
 
-    # Load dataset
+    # Load dataset -> df
     ds = load_dataset(final_id, split="train")
     df = pd.DataFrame(ds)
 
-    # Basic validation
+    # Validate columns
     if label_column not in df.columns:
         raise ValueError(f"Label column '{label_column}' not found in dataset columns.")
     for fc in feature_columns:
         if fc not in df.columns:
             raise ValueError(f"Feature column '{fc}' not found in dataset columns.")
 
-    # Build X, y arrays
+    # Convert to NumPy arrays
     X = df[feature_columns].values
     y = df[label_column].values
 
-    # Split
+    # Train/test split
     X_train, X_test, y_train, y_test = train_test_split(
         X, y, test_size=test_size, random_state=42
     )
 
-    # Train model
+    # Instantiate and train GradientBoostingClassifier
     clf = GradientBoostingClassifier(
         learning_rate=learning_rate,
         n_estimators=int(n_estimators),
@@ -106,14 +102,12 @@ def train_model(dataset_id, custom_dataset_id, label_column, feature_columns,
     )
     clf.fit(X_train, y_train)
 
-    # Predictions & metrics
+    # Evaluate
     y_pred = clf.predict(X_test)
     accuracy = accuracy_score(y_test, y_pred)
     cm = confusion_matrix(y_test, y_pred)
 
-    # Build a single figure with 2 subplots:
-    #   1) Feature importances
-    #   2) Confusion matrix heatmap
+    # Create Matplotlib figure with feature importances + confusion matrix
     fig, axs = plt.subplots(1, 2, figsize=(10, 4))
 
     # Subplot 1: Feature Importances
@@ -131,7 +125,7 @@ def train_model(dataset_id, custom_dataset_id, label_column, feature_columns,
     axs[1].set_xlabel("Predicted")
     axs[1].set_ylabel("True")
 
-    # Optionally annotate each cell with the count
+    # Optionally annotate each cell with numeric counts
     thresh = cm.max() / 2.0
     for i in range(cm.shape[0]):
         for j in range(cm.shape[1]):
@@ -140,7 +134,7 @@ def train_model(dataset_id, custom_dataset_id, label_column, feature_columns,
 
     plt.tight_layout()
 
-    # Build textual summary
+    # Textual summary
     text_summary = (
         f"**Dataset used**: `{final_id}`\n\n"
         f"**Label column**: `{label_column}`\n\n"
@@ -150,44 +144,75 @@ def train_model(dataset_id, custom_dataset_id, label_column, feature_columns,
 
     return text_summary, fig
 
-
-# Build the Gradio Blocks UI
+###############################################################################
+# Gradio UI
+###############################################################################
 with gr.Blocks() as demo:
-    gr.Markdown("# Train a GradientBoostingClassifier on any HF Dataset\n")
+
+    # High-level title and description
     gr.Markdown(
-        "1. Choose a suggested dataset from the dropdown **or** enter a custom dataset ID in the format `user/dataset`.\n"
-        "2. Click **Load Columns** to inspect the columns.\n"
-        "3. Pick a **Label column** and **Feature columns**.\n"
-        "4. Adjust hyperparameters and click **Train & Evaluate**.\n"
-        "5. Observe accuracy, feature importances, and a confusion matrix heatmap.\n\n"
-        "*(Note: the dataset must have a `train` split!)*"
+        """
+        # Interactive Gradient Boosting Demo
+
+        This Space demonstrates how to train a [GradientBoostingClassifier](https://scikit-learn.org/stable/modules/generated/sklearn.ensemble.GradientBoostingClassifier.html#gradientboostingclassifier) from **scikit-learn** on **tabular datasets** hosted on the [Hugging Face Hub](https://huggingface.co/datasets).
+
+        **Purpose**:
+        - Easy explore hyperparameters (_learning_rate, n_estimators, max_depth_) and quickly train an ML model on real data.
+        - Visualise model performance via confusion matrix heatmap and a feature importance plot.
+
+        **How to Use**:
+        1. Select one of the suggested datasets from the dropdown _or_ enter any valid dataset from the [Hugging Face Hub](https://huggingface.co/datasets).
+        2. Click **Load Columns** to retrieve the column names from the dataset's **train** split.
+        3. Choose exactly _one_ **Label column** (the target) and one or more **Feature columns** (the inputs).
+        4. Adjust hyperparameters (learning_rate, n_estimators, max_depth, test_size).
+        5. Click **Train & Evaluate** to train a Gradient Boosting model and see its accuracy, feature importances, and confusion matrix.
+
+        ---
+        **Please Note**:
+        - The dataset must have a **"train"** split with tabular columns (i.e., no nested structures).
+        - Large datasets may take time to download/train.
+        - The confusion matrix helps you see how predictions compare to ground-truth labels. The diagonal cells show correct predictions; off-diagonal cells indicate misclassifications.
+        - The feature importance plot shows which features the model relies on the most for its predictions.
+
+        You are now a machine learning engineer, congratulations 🤗 
+        """
     )
 
-    # Row 1: Dataset selection
     with gr.Row():
         dataset_dropdown = gr.Dropdown(
             label="Choose suggested dataset",
             choices=SUGGESTED_DATASETS,
-            value=SUGGESTED_DATASETS[0]  # default
+            value=SUGGESTED_DATASETS[0]
         )
         custom_dataset_id = gr.Textbox(
             label="Or enter a custom dataset ID",
-            placeholder="e.g. username/my_custom_dataset"
+            placeholder="e.g. user/my_custom_dataset"
         )
 
     load_cols_btn = gr.Button("Load Columns")
     load_cols_info = gr.Markdown()
 
-    # Row 2: label & feature columns
     with gr.Row():
         label_col = gr.Dropdown(choices=[], label="Label column (choose 1)")
         feature_cols = gr.CheckboxGroup(choices=[], label="Feature columns (choose 1 or more)")
 
-    # Hyperparameters
-    learning_rate_slider = gr.Slider(0.01, 1.0, value=0.1, step=0.01, label="learning_rate")
-    n_estimators_slider = gr.Slider(50, 300, value=100, step=50, label="n_estimators")
-    max_depth_slider = gr.Slider(1, 10, value=3, step=1, label="max_depth")
-    test_size_slider = gr.Slider(0.1, 0.9, value=0.3, step=0.1, label="test_size fraction (0.1-0.9)")
+    # Model Hyperparameters
+    learning_rate_slider = gr.Slider(
+        minimum=0.01, maximum=1.0, value=0.1, step=0.01, 
+        label="learning_rate"
+    )
+    n_estimators_slider = gr.Slider(
+        minimum=50, maximum=300, value=100, step=50, 
+        label="n_estimators"
+    )
+    max_depth_slider = gr.Slider(
+        minimum=1, maximum=10, value=3, step=1, 
+        label="max_depth"
+    )
+    test_size_slider = gr.Slider(
+        minimum=0.1, maximum=0.9, value=0.3, step=0.1, 
+        label="test_size fraction (0.1-0.9)"
+    )
 
     train_button = gr.Button("Train & Evaluate")