Spaces:
Running on Zero
Running on Zero
File size: 2,622 Bytes
1dbc4b5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | """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)")
|