Add memorization-prediction demo (ZeroGPU): explore held-out 1M + score-your-own-text

README.md

@@ -1,15 +1,59 @@
 ---
-title: Predicting Memorization
-emoji: 🚀
-colorFrom: yellow
-colorTo: purple
+title: Predicting Memorization Before Fine-Tuning
+emoji: 🔎
+colorFrom: indigo
+colorTo: green
 sdk: gradio
-sdk_version: 6.22.0
-python_version: '3.12'
+sdk_version: 4.44.0
 app_file: app.py
 pinned: false
-license: apache-2.0
-short_description: Predicting memorization in language models before finetuning
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Predicting Memorization Before Fine-Tuning — demo
+
+Interactive companion to the paper. A classifier trained only on **pre-fine-tuning**
+base-model features predicts which sequences a fine-tuned model will memorize.
+
+- **Explore the held-out 1M** — pick a real Pythia-1.4B training sequence and see whether the
+  fine-tuned model actually memorized it (ground truth), next to the classifier's risk score and
+  the base-model feature values. Datasets: FineWeb, PG-19, The Stack, OpenWebMath (all public).
+- **Score your own text** — paste a passage and get the classifier's *predicted* memorization
+  risk from base-model features. This is a forecast; there is no ground truth for arbitrary text.
+
+## Files
+- `app.py` — Gradio app (two tabs).
+- `data/lookup.parquet` — precomputed table (prefix, suffix, memorized flag, classifier score,
+  base-model features) sampled from the held-out 1M run (all memorized examples up to a cap plus a
+  random sample of non-memorized).
+- `models/clf_*.joblib` — the trained GradientBoostingClassifier + scaler per dataset.
+- `byo_features.py` — computes the six base-model features for arbitrary text (loads
+  `EleutherAI/pythia-1.4b`; on ZeroGPU it is placed on `cuda` at startup per HF's guidance).
+- `assemble_lookup_data.py`, `train_classifiers.py` — offline scripts that produced the artifacts
+  above (not needed at runtime; kept for reproducibility).
+
+## Run locally
+```bash
+pip install -r requirements.txt
+python app.py            # opens a local Gradio URL
+```
+Locally the app loads Pythia-1.4B (~3 GB) at startup on a GPU if present, otherwise CPU; the
+`@spaces.GPU` decorator is a no-op off ZeroGPU.
+
+## Deploy on Hugging Face Spaces (free ZeroGPU)
+1. Use your **own, established HF account** (free ZeroGPU hosting requires a verified email and an
+   account older than 30 days; up to 2 ZeroGPU Spaces). Create a new Space: **SDK = Gradio**,
+   hardware = **ZeroGPU**.
+2. Push this folder to the Space repo (the YAML header above configures it):
+   ```bash
+   git clone https://huggingface.co/spaces/<user>/<space> && cd <space>
+   cp -r /path/to/demo_app/* . && git add -A && git commit -m "demo" && git push
+   ```
+3. The Space builds automatically. The Explore tab is instant. "Score your own text" runs the
+   base-model forward+backward on a GPU that ZeroGPU attaches on demand (via the `@spaces.GPU`
+   decorator in `app.py`), so it is fast; the free tier allows 5 GPU-minutes/day and each score
+   takes a couple of seconds.
+
+## Notes
+- The classifier uses only base-model (pre-fine-tuning) features and was trained on a separate
+  Run-1 fine-tuning run, then evaluated here on a disjoint Run-2 run.
+- Qwen and Llama can be added to the Explore tab once their run-2 token sequences are regenerated.

app.py

@@ -0,0 +1,132 @@
+"""Predicting Memorization Before Fine-Tuning — interactive demo.
+
+Two modes:
+  (1) Explore the held-out 1M   — pick a real training sequence, see whether the
+      fine-tuned Pythia-1.4B actually memorized it (ground truth), alongside the
+      classifier's pre-fine-tuning risk score and feature breakdown.
+  (2) Score your own text        — paste a passage, get the classifier's predicted
+      memorization risk from base-model features (a forecast; no ground truth).
+"""
+import os, glob, joblib, numpy as np, pandas as pd, torch, gradio as gr
+import spaces  # Hugging Face ZeroGPU; the @spaces.GPU decorator is a no-op off ZeroGPU
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+BASE_MODEL = os.environ.get("BASE_MODEL", "EleutherAI/pythia-1.4b")
+
+LOOKUP = pd.read_parquet(os.path.join(HERE, "data", "lookup.parquet"))
+DATASETS = sorted(LOOKUP["dataset"].unique())
+CLFS = {}
+for f in glob.glob(os.path.join(HERE, "models", "clf_*.joblib")):
+    b = joblib.load(f)
+    CLFS[b["dataset"]] = b
+FEATS = ["zlib_entropy", "base_ppl", "base_loss_variance",
+         "gradient_norm", "one_step_abs_change", "one_step_rel_change"]
+
+# ZeroGPU: place the model on cuda at module level (a CUDA-emulation mode makes this
+# valid at startup; a real GPU is attached only inside @spaces.GPU functions). Off
+# ZeroGPU we fall back to a real GPU if present, else CPU, so the app still runs locally.
+_ON_ZEROGPU = bool(os.environ.get("SPACES_ZERO_GPU"))
+DEVICE = "cuda" if (_ON_ZEROGPU or torch.cuda.is_available()) else "cpu"
+try:
+    from byo_features import BYOExtractor
+    EXTRACTOR = BYOExtractor(BASE_MODEL, device=DEVICE)
+    _LOAD_ERR = None
+except Exception as e:  # e.g. no model cached / offline local run
+    EXTRACTOR, _LOAD_ERR = None, str(e)
+
+
+@spaces.GPU(duration=120)
+def _compute_features(text):
+    """GPU-bound step of the BYO path (forward + backward on the base model)."""
+    return EXTRACTOR.features(text)
+
+
+def _feature_table(row):
+    names = {"zlib_entropy": "zlib entropy", "base_ppl": "base perplexity",
+             "base_loss_variance": "base loss variance", "gradient_norm": "gradient norm",
+             "one_step_rel_change": "one-step rel. loss change"}
+    return pd.DataFrame(
+        [(v, f"{float(row[k]):.4f}") for k, v in names.items() if k in row],
+        columns=["Base-model feature", "Value"])
+
+
+# ----------------------------- Explore (lookup) -----------------------------
+def draw(dataset, which):
+    df = LOOKUP[LOOKUP.dataset == dataset]
+    if which == "Memorized only":
+        df = df[df.memorized]
+    elif which == "Not memorized":
+        df = df[~df.memorized]
+    if len(df) == 0:
+        return "No examples for this filter.", "", "", None
+    r = df.sample(1).iloc[0]
+    if r.memorized:
+        verdict = "### 🔴 Memorized\nThe fine-tuned model reproduced the ground-truth suffix exactly."
+    else:
+        verdict = "### 🟢 Not memorized\nThe fine-tuned model did not reproduce the ground-truth suffix."
+    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.)_"
+    score_md = (f"### Classifier risk score: **{float(r.score):.3f}**\n"
+                f"Computed from base-model features *before* fine-tuning. "
+                f"Higher means more likely to be memorized. Ground truth above is what actually happened.")
+    return (r.prefix, r.suffix, verdict + "\n\n" + score_md, _feature_table(r))
+
+
+# ----------------------------- Predict (BYO) --------------------------------
+def predict(text, clf_name):
+    text = (text or "").strip()
+    if not text:
+        return "Paste a passage first.", None
+    if EXTRACTOR is None:
+        return f"⚠️ Base model unavailable: {_LOAD_ERR}", None
+    try:
+        feats = _compute_features(text)
+    except ValueError as e:
+        return f"⚠️ {e}", None
+    b = CLFS[clf_name]
+    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])
+    flag = "high" if risk >= 0.5 else "low"
+    md = (f"### Predicted memorization risk: **{risk:.3f}** ({flag})\n"
+          f"Forecast from base-model features using the classifier trained on **{clf_name}**. "
+          f"This is a prediction of whether a model *would* memorize this sequence if fine-tuned on it; "
+          f"there is no ground truth for arbitrary text. Treat it as a relative risk score, not a calibrated probability.")
+    return md, _feature_table(feats)
+
+
+INTRO = """# Predicting Memorization Before Fine-Tuning
+A classifier trained only on **pre-fine-tuning** base-model features predicts which sequences a
+fine-tuned model will memorize. **Explore** verifies this against ground truth on a held-out 1M-sequence
+run; **Score your own text** applies the same classifier to a passage you provide.
+"""
+
+with gr.Blocks(title="Predicting Memorization Before Fine-Tuning") as demo:
+    gr.Markdown(INTRO)
+    with gr.Tab("Explore the held-out 1M"):
+        gr.Markdown("Pick a real Pythia-1.4B training sequence and compare the classifier's "
+                    "pre-fine-tuning risk score with whether the fine-tuned model actually memorized it.")
+        with gr.Row():
+            d_in = gr.Dropdown(DATASETS, value=DATASETS[0], label="Dataset")
+            f_in = gr.Radio(["All", "Memorized only", "Not memorized"], value="All", label="Filter")
+            go = gr.Button("Show a random example", variant="primary")
+        pref = gr.Textbox(label="Prefix (50 tokens given to the model)", lines=3)
+        suf = gr.Textbox(label="Ground-truth suffix (next 50 tokens)", lines=3)
+        verdict = gr.Markdown()
+        feat = gr.Dataframe(label="Base-model features for this sequence", interactive=False)
+        go.click(draw, [d_in, f_in], [pref, suf, 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. No ground truth exists for "
+                    "arbitrary text, so this is a forecast (see the trust discussion in the paper).")
+        txt = gr.Textbox(label="Your text", lines=6)
+        clf_in = gr.Dropdown(sorted(CLFS.keys()), value=sorted(CLFS.keys())[0],
+                             label="Classifier (trained on)")
+        pgo = gr.Button("Predict memorization risk", variant="primary")
+        pmd = gr.Markdown()
+        pfeat = gr.Dataframe(label="Base-model features for your text", interactive=False)
+        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._")
+
+if __name__ == "__main__":
+    demo.launch()

assemble_lookup_data.py

@@ -0,0 +1,59 @@
+"""Assemble the precomputed lookup table for the demo.
+
+For each (Pythia, dataset) Run-2 set, join token sequences + memorized labels +
+classifier score + base-model features, decode the 50-token prefix/suffix to text,
+and keep a compact sample (all memorized up to a cap, plus a random sample of
+non-memorized). Public / non-PII datasets only.
+
+Output: demo_app/data/lookup.parquet
+"""
+import os, numpy as np, pandas as pd
+from transformers import AutoTokenizer
+
+ROOT = "/projects/nulab/borkar.j/predicting_memorization"
+TOK = f"{ROOT}/checkpoints/pythia-1.4b"
+DATASETS = {"fineweb": "FineWeb", "pg19": "PG-19", "code": "The Stack", "owm": "OpenWebMath"}
+PREFIX_LEN, WINDOW_LEN = 50, 100
+MEM_CAP, NONMEM_N, SEED = 3000, 5000, 42
+OUT = f"{ROOT}/demo_app/data"
+
+tok = AutoTokenizer.from_pretrained(TOK)
+rng = np.random.default_rng(SEED)
+rows = []
+
+for d, disp in DATASETS.items():
+    ck = f"{ROOT}/checkpoints/pythia_{d}_run2_seed42"
+    ds = np.load(f"{ck}/dataset.npy")                        # [1M,128]
+    lab = np.load(f"{ck}/memorization_labels.npy").astype(bool)
+    fe = pd.read_csv(f"{ROOT}/full_run2_results/features/{d}_run2_full_features.csv")
+    sc = pd.read_csv(f"{ROOT}/full_run2_results/cleaned/scores/{d}_run2_scores.csv")
+    assert len(ds) == len(lab) == len(fe) == len(sc), d
+
+    mem = np.where(lab)[0]
+    non = np.where(~lab)[0]
+    if len(mem) > MEM_CAP:
+        mem = rng.choice(mem, MEM_CAP, replace=False)
+    non = rng.choice(non, min(NONMEM_N, len(non)), replace=False)
+    idx = np.concatenate([mem, non])
+
+    pref = tok.batch_decode(ds[idx, :PREFIX_LEN].tolist(), skip_special_tokens=True)
+    suf = tok.batch_decode(ds[idx, PREFIX_LEN:WINDOW_LEN].tolist(), skip_special_tokens=True)
+    sub = pd.DataFrame({
+        "model": "Pythia-1.4B", "dataset": disp,
+        "prefix": pref, "suffix": suf,
+        "memorized": lab[idx],
+        "base_reproduced": sc["base_reproduced"].values[idx].astype(bool),
+        "score": sc["score"].values[idx],
+        "zlib_entropy": fe["zlib_entropy"].values[idx],
+        "base_ppl": fe["base_ppl"].values[idx],
+        "base_loss_variance": fe["base_loss_variance"].values[idx],
+        "gradient_norm": fe["gradient_norm"].values[idx],
+        "one_step_rel_change": fe["one_step_rel_change"].values[idx],
+    })
+    rows.append(sub)
+    print(f"{disp:12s} kept {len(sub)} rows ({int(lab[idx].sum())} memorized)")
+
+os.makedirs(OUT, exist_ok=True)
+df = pd.concat(rows, ignore_index=True)
+df.to_parquet(f"{OUT}/lookup.parquet", index=False)
+print(f"\nwrote {OUT}/lookup.parquet  ({len(df)} rows, {df.memory_usage(deep=True).sum()/1e6:.1f} MB in mem)")

byo_features.py

@@ -0,0 +1,81 @@
+"""Compute the six base-model features for an arbitrary user sequence, mirroring
+feature_extraction_{cpu,gpu}.py so scores are comparable to the paper's classifier.
+
+Definitions (base model, no fine-tuning):
+  zlib_entropy        = compressed byte length of the decoded 50-token suffix
+  base_ppl            = exp(mean per-token loss) over the 50-token continuation
+  base_loss_variance  = variance of that per-token loss
+  gradient_norm       = ||grad of the continuation loss w.r.t. all params||_2
+  one_step_abs_change = loss_before - loss_after, after one SGD step (lr=2e-5) on x
+  one_step_rel_change = one_step_abs_change / loss_before
+"""
+import zlib, numpy as np, torch, torch.nn.functional as F
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+PREFIX_LEN, WINDOW_LEN, LR = 50, 100, 2e-5
+
+
+class BYOExtractor:
+    def __init__(self, model_path, device=None):
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.tok = AutoTokenizer.from_pretrained(model_path)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_path, torch_dtype=torch.float32).to(self.device).eval()
+
+    def _cont_loss(self, ids):
+        logits = self.model(ids).logits
+        cl = logits[:, PREFIX_LEN - 1:WINDOW_LEN - 1, :]
+        ct = ids[:, PREFIX_LEN:WINDOW_LEN]
+        return -F.log_softmax(cl, -1).gather(2, ct.unsqueeze(-1)).squeeze(-1)  # [1,50] per-token loss
+
+    def features(self, text):
+        enc = self.tok(text, return_tensors="pt", truncation=True, max_length=WINDOW_LEN)
+        ids = enc["input_ids"]
+        if ids.shape[1] < WINDOW_LEN:
+            raise ValueError(
+                f"Need at least {WINDOW_LEN} tokens (got {ids.shape[1]}). Paste a longer passage.")
+        return self.features_from_ids(ids[:, :WINDOW_LEN])
+
+    def features_from_ids(self, ids):
+        """ids: LongTensor [1, WINDOW_LEN] of exact token ids (bypasses tokenization).
+        Used both by features(text) and by the offline validation against the pipeline."""
+        if not torch.is_tensor(ids):
+            ids = torch.tensor(ids, dtype=torch.long)
+        if ids.dim() == 1:
+            ids = ids.unsqueeze(0)
+        ids = ids[:, :WINDOW_LEN].long().to(self.device)
+
+        # zlib on the decoded suffix
+        suffix_txt = self.tok.decode(ids[0, PREFIX_LEN:WINDOW_LEN], skip_special_tokens=True)
+        zlib_entropy = float(len(zlib.compress(suffix_txt.encode("utf-8"))))
+
+        # loss-based features
+        with torch.no_grad():
+            ptl = self._cont_loss(ids)
+        base_ppl = float(torch.exp(ptl.mean()))
+        base_loss_variance = float(ptl.var(unbiased=True))
+
+        # gradient norm + one-step loss change
+        self.model.zero_grad(set_to_none=True)
+        loss_before = self._cont_loss(ids).mean()
+        loss_before.backward()
+        with torch.no_grad():
+            gsq = sum(float(p.grad.pow(2).sum()) for p in self.model.parameters() if p.grad is not None)
+            gradient_norm = gsq ** 0.5
+            for p in self.model.parameters():          # one SGD step
+                if p.grad is not None:
+                    p.data.add_(p.grad, alpha=-LR)
+            loss_after = float(self._cont_loss(ids).mean())
+            for p in self.model.parameters():          # restore
+                if p.grad is not None:
+                    p.data.add_(p.grad, alpha=LR)
+        lb = float(loss_before)
+        one_step_abs_change = lb - loss_after
+        one_step_rel_change = one_step_abs_change / lb if lb else 0.0
+        self.model.zero_grad(set_to_none=True)
+
+        return {
+            "zlib_entropy": zlib_entropy, "base_ppl": base_ppl,
+            "base_loss_variance": base_loss_variance, "gradient_norm": gradient_norm,
+            "one_step_abs_change": one_step_abs_change, "one_step_rel_change": one_step_rel_change,
+        }

data/lookup.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4bcf4588bd84e05e61578ff3d7794f505ec2917c3a88f2bedfe93cd1094ab114
+size 6934910

models/clf_code.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:40902845c23e3e82d781d33d9890ee9236bd5040a0456e2f95f9094a4b1803ff
+size 856310

models/clf_fineweb.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:367d9be1ebaabc7f7e673a516d998e0dd73356a219344ef6bacc71f92a6fb820
+size 447428

models/clf_owm.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:73df384220a1127674a6ce31bf35d03ef50746aaabaf94defd944ea0eb3334c7
+size 766728

models/clf_pg19.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d75a46e460e8511bea39e25187373cdd50d0d790dd2b43b38504fe19e5676c0f
+size 363618

requirements.txt

@@ -0,0 +1,9 @@
+spaces
+gradio>=4.44
+transformers>=4.44
+torch>=2.8
+scikit-learn==1.7.2
+pandas>=2.0
+pyarrow>=14.0
+numpy>=2.0
+joblib>=1.5

train_classifiers.py

@@ -0,0 +1,39 @@
+"""Train + save the memorization classifiers used by the demo's your-own-text mode.
+
+Mirrors score_1m_nodup.py: train on the Run-1 deduplicated features with base-model-
+reproduced positives removed, 1:3 rebalanced, on the six no-dup features. One classifier
+per dataset (they transfer, per the cross-domain result, but we keep all four). Saves the
+fitted StandardScaler + GradientBoostingClassifier per dataset via joblib.
+"""
+import os, numpy as np, pandas as pd, joblib
+from sklearn.ensemble import GradientBoostingClassifier
+from sklearn.preprocessing import StandardScaler
+
+ROOT = "/projects/nulab/borkar.j/predicting_memorization"
+OUT = f"{ROOT}/demo_app/models"
+FEATS = ["zlib_entropy", "base_ppl", "base_loss_variance",
+         "gradient_norm", "one_step_abs_change", "one_step_rel_change"]
+DATASETS = {"fineweb": "FineWeb", "pg19": "PG-19", "code": "The Stack", "owm": "OpenWebMath"}
+os.makedirs(OUT, exist_ok=True)
+
+for d, disp in DATASETS.items():
+    df = pd.read_csv(f"{ROOT}/{d}_features.csv")
+    m = np.load(f"{ROOT}/base_extractability/base_memorized_{d}.npy")
+    pi = df.index[df.label == 1].tolist()
+    df["bm"] = 0
+    df.loc[pi, "bm"] = m.astype(int)
+    df = df[df.approx_dup_count == 1]                       # dedup
+    df = df[~((df.label == 1) & (df.bm == 1))]             # remove base-reproduced positives
+    pos = df[df.label == 1]; neg = df[df.label == 0]; n = 3 * len(pos)
+    if len(neg) > n:
+        neg = neg.sample(n=n, random_state=42)
+    tr = pd.concat([pos, neg]).sample(frac=1, random_state=42).reset_index(drop=True)
+    X = np.nan_to_num(tr[FEATS].values, posinf=1e10, neginf=-1e10)
+    y = tr.label.values
+    sc = StandardScaler().fit(X)
+    clf = GradientBoostingClassifier(n_estimators=200, max_depth=5, learning_rate=0.1,
+                                     subsample=0.8, random_state=42).fit(sc.transform(X), y)
+    joblib.dump({"scaler": sc, "clf": clf, "feats": FEATS, "dataset": disp},
+                f"{OUT}/clf_{d}.joblib")
+    print(f"{disp:12s} trained on {len(tr)} rows ({int(y.sum())} pos); saved clf_{d}.joblib")
+print("done")