Fix HF Space graph rendering

app.py

@@ -277,7 +277,8 @@ FEATURES_HEADING = "#### Extracted Features"
 PROBS_HEADING = "#### Class Probabilities"
 SHAPLEY_HEADING = (
     "#### Shapley Attribution\n\n"
-    "*Waterfall: each bar is a feature's Shapley contribution toward the predicted class.*"
+    "*Each bar is a feature's Shapley contribution toward the predicted class "
+    "(positive = supports, negative = works against).*"
 )
 MEANS_HEADING = "#### What This Means"
 NOT_READY_MSG = (
@@ -432,7 +433,7 @@ def handle_own_wine(
         )
 
         progress(0.55, desc="Computing Shapley attributions (shapiq)…")
-        waterfall_fig, attributions, shapley_meta = safe_explain(features, pred_idx)
+        attribution_img, attributions, shapley_meta = safe_explain(features, pred_idx)
 
         progress(0.75, desc="Running what-if probes…")
         counterfactuals = run_counterfactuals(clf, features)
@@ -465,7 +466,7 @@ def handle_own_wine(
         SHOW,
         gr.update(value=probs_to_df(probabilities), visible=True),
         SHOW,
-        gr.update(value=waterfall_fig, visible=waterfall_fig is not None),
+        gr.update(value=attribution_img, visible=attribution_img is not None),
         gr.update(value=attributions, visible=True),
         SHOW,
         gr.update(value=explanation, visible=True),
@@ -544,7 +545,7 @@ def handle_sample_predict(
     )
 
     progress(0.5, desc="Computing Shapley attributions (shapiq)…")
-    waterfall_fig, attributions, shapley_meta = safe_explain(feat, pred_idx)
+    attribution_img, attributions, shapley_meta = safe_explain(feat, pred_idx)
 
     progress(0.85, desc="Writing plain-English explanation…")
     from agent import generate_layman_explanation
@@ -565,7 +566,7 @@ def handle_sample_predict(
         gr.update(value=probs_to_df(probabilities), visible=True),
         gr.update(value=result_md, visible=True),
         SHOW,
-        gr.update(value=waterfall_fig, visible=waterfall_fig is not None),
+        gr.update(value=attribution_img, visible=attribution_img is not None),
         gr.update(value=attributions, visible=True),
         SHOW,
         gr.update(value=explanation, visible=True),
@@ -670,7 +671,7 @@ Free-tier usage limits apply on Prior Labs; see their docs for current quotas.
                     with gr.Column():
                         sample_verdict_md = gr.Markdown(visible=False)
                 sample_shapley_heading = gr.Markdown(SHAPLEY_HEADING, visible=False)
-                sample_attributions_plot = gr.Plot(visible=False)
+                sample_attributions_plot = gr.Image(visible=False, type="pil")
                 sample_attributions_table = gr.DataFrame(visible=False)
                 sample_means_heading = gr.Markdown(MEANS_HEADING, visible=False)
                 sample_explanation_md = gr.Markdown(visible=False)
@@ -719,7 +720,7 @@ Free-tier usage limits apply on Prior Labs; see their docs for current quotas.
                         own_probs_heading = gr.Markdown(PROBS_HEADING, visible=False)
                         own_probs_table = gr.DataFrame(visible=False)
                 own_shapley_heading = gr.Markdown(SHAPLEY_HEADING, visible=False)
-                own_attributions_plot = gr.Plot(visible=False)
+                own_attributions_plot = gr.Image(visible=False, type="pil")
                 own_attributions_table = gr.DataFrame(visible=False)
                 own_means_heading = gr.Markdown(MEANS_HEADING, visible=False)
                 own_explanation_md = gr.Markdown(visible=False)

interpret.py

@@ -2,13 +2,17 @@
 
 from __future__ import annotations
 
+import io
 import warnings
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
+import matplotlib
+
+matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
-from matplotlib.figure import Figure
+from PIL import Image
 from tabpfn_client import TabPFNClassifier
 from tabpfn_extensions.interpretability.shapiq import get_tabpfn_imputation_explainer
 
@@ -19,6 +23,12 @@ _explainer_cache: Dict[int, object] = {}
 _BACKGROUND_ROWS = 256
 _SHAPIQ_BUDGET = 64
 
+_POS_COLOR = "#e85d75"
+_NEG_COLOR = "#4a9eff"
+_BG_COLOR = "#1e0f14"
+_TEXT_COLOR = "#f5e6d3"
+_GRID_COLOR = "#722f37"
+
 
 def clear_explainer_cache() -> None:
     """Drop cached explainers after reconnecting with a new TabPFN token."""
@@ -33,8 +43,6 @@ def _get_explainer(
     if class_index not in _explainer_cache:
         n = min(_BACKGROUND_ROWS, len(train_df))
         background = train_df[FEATURES].sample(n=n, random_state=42).values
-        # Permutation sampling avoids KernelSHAP's regression solve, which can emit
-        # divide-by-zero / overflow warnings with near-flat TabPFN coalitions.
         _explainer_cache[class_index] = get_tabpfn_imputation_explainer(
             clf,
             background,
@@ -48,49 +56,76 @@ def _get_explainer(
     return _explainer_cache[class_index]
 
 
-def _style_waterfall(fig: Figure) -> Figure:
-    """Dark-theme styling and extra margins so labels are not clipped in Gradio."""
-    fig.patch.set_facecolor("#1e0f14")
-    primary_ax = fig.axes[0] if fig.axes else None
-
-    for ax in fig.axes:
-        ax.set_facecolor("#1e0f14")
-        ax.tick_params(colors="#f5e6d3", labelcolor="#f5e6d3")
-        for spine in ax.spines.values():
-            spine.set_color("#722f37")
-        if ax.title:
-            ax.title.set_color("#f5e6d3")
-        if ax.xaxis.label:
-            ax.xaxis.label.set_color("#c9a87c")
-        for label in ax.get_yticklabels() + ax.get_xticklabels():
-            label.set_color("#f5e6d3")
-        xmin, xmax = ax.get_xlim()
-        pad = (xmax - xmin) * 0.14
-        ax.set_xlim(xmin - pad, xmax + pad)
-
-    if primary_ax is not None:
-        w, h = fig.get_size_inches()
-        fig.set_size_inches(max(w, 12.0), max(h, 6.5))
-        primary_ax.tick_params(axis="y", pad=10)
-        # Wide left margin for full feature names; padding on x for arrow heads.
-        fig.subplots_adjust(left=0.42, right=0.90, bottom=0.16, top=0.86)
-
-    return fig
-
-
-def _shapley_table(sv) -> pd.DataFrame:
-    """First-order Shapley values as a readable table."""
+def _shapley_pairs(sv) -> List[Tuple[str, float]]:
     rows = []
     for interaction, value in sv.dict_values.items():
         if len(interaction) != 1:
             continue
-        idx = interaction[0]
-        key = FEATURES[idx]
-        magnitude = abs(float(value))
+        key = FEATURES[interaction[0]]
+        rows.append((FEATURE_DISPLAY_NAMES[key], float(value)))
+    rows.sort(key=lambda item: abs(item[1]), reverse=True)
+    return rows
+
+
+def _shapley_bar_image(
+    pairs: List[Tuple[str, float]],
+    baseline: float,
+    final: float,
+) -> Image.Image:
+    """Diverging bar chart — rasterized for reliable display in Gradio on HF."""
+    if not pairs:
+        fig, ax = plt.subplots(figsize=(10, 2), facecolor=_BG_COLOR)
+        ax.set_facecolor(_BG_COLOR)
+        ax.text(0.5, 0.5, "No attributions available", ha="center", va="center", color=_TEXT_COLOR)
+        ax.axis("off")
+    else:
+        names = [name for name, _ in pairs]
+        values = [value for _, value in pairs]
+        height = max(5.0, len(names) * 0.42)
+        fig, ax = plt.subplots(figsize=(10, height), facecolor=_BG_COLOR)
+        ax.set_facecolor(_BG_COLOR)
+        colors = [_POS_COLOR if v > 0 else _NEG_COLOR for v in values]
+        y_pos = np.arange(len(names))
+        ax.barh(y_pos, values, color=colors, height=0.72)
+        ax.set_yticks(y_pos)
+        ax.set_yticklabels(names, color=_TEXT_COLOR, fontsize=10)
+        ax.axvline(0, color=_GRID_COLOR, linewidth=1)
+        ax.tick_params(axis="x", colors=_TEXT_COLOR, labelcolor=_TEXT_COLOR)
+        ax.set_xlabel("Shapley contribution (predicted-class probability)", color="#c9a87c")
+        ax.set_title(
+            f"Baseline {baseline:.3f}  →  prediction {final:.3f}",
+            color=_TEXT_COLOR,
+            fontsize=11,
+            pad=10,
+        )
+        for spine in ax.spines.values():
+            spine.set_color(_GRID_COLOR)
+        ax.grid(axis="x", color=_GRID_COLOR, alpha=0.35, linestyle="--", linewidth=0.6)
+        ax.invert_yaxis()
+        fig.tight_layout()
+
+    buf = io.BytesIO()
+    fig.savefig(
+        buf,
+        format="png",
+        dpi=150,
+        bbox_inches="tight",
+        facecolor=_BG_COLOR,
+        edgecolor="none",
+    )
+    plt.close(fig)
+    buf.seek(0)
+    return Image.open(buf)
+
+
+def _shapley_table(pairs: List[Tuple[str, float]]) -> pd.DataFrame:
+    rows = []
+    for name, value in pairs:
+        magnitude = abs(value)
         bar_len = max(1, round(magnitude * 40)) if magnitude > 0 else 0
         rows.append({
-            "Feature": FEATURE_DISPLAY_NAMES[key],
-            "Shapley": round(float(value), 4),
+            "Feature": name,
+            "Shapley": round(value, 4),
             "Effect": (
                 "supports predicted class" if value > 0
                 else "works against predicted class" if value < 0
@@ -98,10 +133,7 @@ def _shapley_table(sv) -> pd.DataFrame:
             ),
             "Bar": "█" * bar_len,
         })
-    df = pd.DataFrame(rows)
-    if df.empty:
-        return df
-    return df.reindex(df["Shapley"].abs().sort_values(ascending=False).index)
+    return pd.DataFrame(rows)
 
 
 def explain_prediction(
@@ -109,8 +141,8 @@ def explain_prediction(
     train_df: pd.DataFrame,
     features: dict,
     class_index: int,
-) -> Tuple[Optional[Figure], pd.DataFrame, Dict[str, Any]]:
-    """Run shapiq SV attribution and return waterfall figure, table, and metadata."""
+) -> Tuple[Optional[Image.Image], pd.DataFrame, Dict[str, Any]]:
+    """Run shapiq SV attribution and return a bar-chart image, table, and metadata."""
     explainer = _get_explainer(clf, train_df, class_index)
     x = np.array([[features[f] for f in FEATURES]], dtype=float)
     with warnings.catch_warnings():
@@ -121,23 +153,16 @@ def explain_prediction(
         )
         sv = explainer.explain(x, budget=_SHAPIQ_BUDGET)
 
-    display_names = np.array([FEATURE_DISPLAY_NAMES[f] for f in FEATURES])
-    ax = sv.plot_waterfall(
-        feature_names=display_names,
-        show=False,
-        max_display=len(FEATURES),
-        abbreviate=False,
-    )
-    fig = ax.figure if ax is not None else plt.gcf()
-    _style_waterfall(fig)
+    pairs = _shapley_pairs(sv)
+    shapley_sum = sum(value for _, value in pairs)
+    baseline = float(sv.baseline_value)
+    final = baseline + shapley_sum
 
-    shapley_sum = sum(
-        float(v) for k, v in sv.dict_values.items() if len(k) == 1
-    )
+    image = _shapley_bar_image(pairs, baseline, final)
     meta = {
-        "baseline_probability": round(float(sv.baseline_value), 4),
-        "final_probability": round(float(sv.baseline_value) + shapley_sum, 4),
+        "baseline_probability": round(baseline, 4),
+        "final_probability": round(final, 4),
         "predicted_class_index": class_index,
         "attribution_method": "shapiq SV via permutation sampling (first-order)",
     }
-    return fig, _shapley_table(sv), meta
+    return image, _shapley_table(pairs), meta