Replace SHAP with feature-distribution plots; compute Explore score from the deployed classifier (consistent with BYO); drop shap dep

app.py

@@ -13,7 +13,6 @@ predict memorization about as well as the full feature set.
 """
 import os, glob, json, joblib, numpy as np, pandas as pd, torch, gradio as gr
 import spaces  # Hugging Face ZeroGPU; @spaces.GPU is a no-op off ZeroGPU
-import shap
 import matplotlib
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
@@ -27,13 +26,18 @@ CLFS = {}
 for f in glob.glob(os.path.join(HERE, "models", "clf_*.joblib")):
     b = joblib.load(f)
     CLFS[b["dataset"]] = b
-EXPLAINERS = {}  # lazily built shap.TreeExplainer per classifier
 
+# Feature distributions per dataset, split by memorization outcome (for the "where does
+# this example fall" plots). Model-free: reads straight from the lookup table.
+FEAT_ORDER = ["base_ppl", "base_loss_variance", "one_step_rel_change",
+              "one_step_abs_change", "gradient_norm", "zlib_entropy"]
 FEAT_NAMES = {"zlib_entropy": "zlib entropy", "base_ppl": "base perplexity",
               "base_loss_variance": "base loss variance", "gradient_norm": "gradient norm",
               "one_step_abs_change": "one-step loss change", "one_step_rel_change": "one-step rel. loss change"}
-GRAD_FEATS = {"gradient_norm", "one_step_abs_change", "one_step_rel_change"}
-GREEN, RED, GREY = "#2e7d32", "#c62828", "#9aa0a6"
+DIST = {ds: {"mem": LOOKUP[(LOOKUP.dataset == ds) & LOOKUP.memorized],
+             "non": LOOKUP[(LOOKUP.dataset == ds) & ~LOOKUP.memorized]} for ds in DATASETS}
+GREEN, RED = "#2e7d32", "#c62828"          # generation highlight
+MEMCLR, NONCLR = "#f58518", "#4c78a8"      # memorized (orange) vs not-memorized (blue)
 
 # Headline ablation result (paper: dedup, mean +/- std over 3 training seeds), demo datasets only.
 HEADLINE = {
@@ -80,29 +84,38 @@ def headline_figure():
     return fig
 
 
-def shap_figure(feats, clf_name):
-    b = CLFS[clf_name]
-    if clf_name not in EXPLAINERS:
-        EXPLAINERS[clf_name] = shap.TreeExplainer(b["clf"])
-    expl = EXPLAINERS[clf_name]
-    order = b["feats"]
-    X = np.nan_to_num(np.array([[feats[k] for k in order]], dtype=float), posinf=1e10, neginf=-1e10)
-    Xs = b["scaler"].transform(X)
-    sv = expl.shap_values(Xs)
-    sv = np.asarray(sv[1] if isinstance(sv, list) else sv).reshape(-1)[:len(order)]
-
-    idx = np.argsort(np.abs(sv))               # ascending -> largest on top
-    labels = [f"{FEAT_NAMES.get(order[i], order[i])}  ({feats[order[i]]:.3g})" for i in idx]
-    vals = sv[idx]
-    colors = [GREEN if v > 0 else GREY for v in vals]
-    fig, ax = plt.subplots(figsize=(7.2, 3.2))
-    ax.barh(range(len(vals)), vals, color=colors)
-    ax.set_yticks(range(len(vals))); ax.set_yticklabels(labels, fontsize=9)
-    ax.axvline(0, color="k", lw=0.8)
-    ax.set_xlabel("push toward memorized  →   (log-odds contribution)")
-    grad_top = [order[i] in GRAD_FEATS for i in idx]
-    ax.set_title("Why this score" + ("  ·  gradient features dominate" if any(grad_top[-2:]) else ""))
-    fig.tight_layout()
+def dist_figure(feats, dataset):
+    """Six panels: where this example's feature values fall relative to the memorized
+    and not-memorized populations of the chosen dataset. Model-free."""
+    d = DIST.get(dataset) or DIST[next(iter(DIST))]
+    mem, non = d["mem"], d["non"]
+    fig, axes = plt.subplots(2, 3, figsize=(9.6, 5.4))
+    for ax, f in zip(axes.ravel(), FEAT_ORDER):
+        mv = mem[f].to_numpy(float); nv = non[f].to_numpy(float)
+        both = np.concatenate([mv, nv])
+        lo, hi = np.nanpercentile(both, [1, 99])
+        if not np.isfinite(lo) or not np.isfinite(hi) or lo == hi:
+            lo, hi = float(np.nanmin(both)), float(np.nanmax(both)) + 1e-9
+        use_log = lo > 0 and hi / lo > 50
+        bins = (np.logspace(np.log10(lo), np.log10(hi), 40) if use_log
+                else np.linspace(lo, hi, 40))
+        centers = 0.5 * (bins[:-1] + bins[1:])
+        for arr, color, label in [(nv, NONCLR, "not memorized"), (mv, MEMCLR, "memorized")]:
+            c, _ = np.histogram(np.clip(arr, lo, hi), bins=bins)
+            c = c / c.max() if c.max() > 0 else c   # normalize each class to its own peak
+            ax.fill_between(centers, c, step="mid", alpha=0.5, color=color, label=label)
+        if use_log:
+            ax.set_xscale("log")
+        v = float(feats[f]); vc = min(max(v, lo), hi)
+        ax.axvline(vc, color="k", lw=1.8)
+        ax.plot([vc], [ax.get_ylim()[1] * 0.9], marker="v", color="k", ms=9, clip_on=False)
+        off = "" if lo <= v <= hi else ("  (off ↓)" if v < lo else "  (off ↑)")
+        ax.set_title(f"{FEAT_NAMES[f]}\nthis example = {v:.3g}{off}", fontsize=8.5)
+        ax.set_yticks([]); ax.tick_params(labelsize=7)
+    axes.ravel()[0].legend(fontsize=7.5, loc="upper right", framealpha=0.9)
+    fig.suptitle(f"Where this example's base-model features fall — {dataset}  "
+                 f"(black ▼ = this example)", fontsize=10)
+    fig.tight_layout(rect=[0, 0, 1, 0.95])
     return fig
 
 
@@ -129,13 +142,16 @@ def draw(dataset, which):
                    "green tokens happen to match, red tokens do not.")
     if bool(r.base_reproduced):
         verdict += "\n\n_(Note: the base model already reproduces this suffix, so it is contamination rather than fine-tuning-induced memorization.)_"
-    verdict += (f"\n\n### Classifier risk score: **{float(r.score):.3f}**\n"
-                f"Computed from base-model features *before* fine-tuning. The chart below shows how each "
-                f"feature moved this score.")
     clf_name = dataset if dataset in CLFS else next(iter(CLFS))
-    feats = {k: float(r[k]) for k in CLFS[clf_name]["feats"]}
-    fig = shap_figure(feats, clf_name)
-    return r.prefix, gen, verdict, fig
+    b = CLFS[clf_name]
+    feats = {k: float(r[k]) for k in FEAT_ORDER}
+    X = np.nan_to_num(np.array([[feats[k] for k in b["feats"]]]), posinf=1e10, neginf=-1e10)
+    risk = float(b["clf"].predict_proba(b["scaler"].transform(X))[0, 1])
+    verdict += (f"\n\n### Classifier risk score: **{risk:.3f}**\n"
+                f"Computed by the memorization classifier from base-model features measured *before* "
+                f"fine-tuning (the same classifier used in the other tab). The plots below show where "
+                f"this example's feature values fall relative to memorized and not-memorized examples.")
+    return r.prefix, gen, verdict, dist_figure(feats, dataset)
 
 
 # ----------------------------- Predict (BYO) --------------------------------
@@ -156,8 +172,9 @@ def predict(text, clf_name):
     md = (f"### Predicted memorization risk: **{risk:.3f}** ({flag})\n"
           f"Forecast from base-model features using the classifier trained on **{clf_name}**. "
           f"There is no ground truth for arbitrary text, so treat it as a relative risk score. "
-          f"The chart shows which features drove it.")
-    return md, shap_figure(feats, clf_name)
+          f"The plots below show where your text's features fall relative to memorized and not-memorized "
+          f"examples in **{clf_name}**.")
+    return md, dist_figure(feats, clf_name)
 
 
 INTRO = """# Predicting Memorization Before Fine-Tuning
@@ -190,18 +207,18 @@ with gr.Blocks(title="Predicting Memorization Before Fine-Tuning") as demo:
             label="Model's actual 50-token generation (green = matches ground truth, red = differs)",
             color_map={"match": GREEN, "diff": RED}, combine_adjacent=True, show_legend=False)
         verdict = gr.Markdown()
-        feat = gr.Plot(label="Per-example feature contributions (SHAP)")
+        feat = gr.Plot(label="Where this example's features fall (memorized vs not-memorized)")
         go.click(draw, [d_in, f_in], [pref, gen, verdict, feat])
     with gr.Tab("Score your own text"):
         gr.Markdown("Paste a passage (at least ~100 tokens). We compute its base-model features and "
-                    "return the classifier's **predicted** memorization risk, with the per-example "
-                    "feature contributions below.")
+                    "return the classifier's **predicted** memorization risk, and show where your text's "
+                    "features fall relative to memorized and not-memorized examples.")
         txt = gr.Textbox(label="Your text", lines=6)
         clf_in = gr.Dropdown(sorted(CLFS.keys()), value=sorted(CLFS.keys())[0],
-                             label="Classifier (trained on)")
+                             label="Compare against dataset")
         pgo = gr.Button("Predict memorization risk", variant="primary")
         pmd = gr.Markdown()
-        pfeat = gr.Plot(label="Per-example feature contributions (SHAP)")
+        pfeat = gr.Plot(label="Where your text's features fall (memorized vs not-memorized)")
         pgo.click(predict, [txt, clf_in], [pmd, pfeat])
     gr.Markdown("_Datasets shown are public (FineWeb, PG-19, The Stack, OpenWebMath). "
                 "The classifier uses only base-model features and was trained on a separate fine-tuning run._")

requirements.txt

@@ -7,5 +7,4 @@ pandas>=2.0
 pyarrow>=14.0
 numpy>=2.0
 joblib>=1.5
-shap==0.49.1
 matplotlib>=3.7