From a42477ff44d137073dff53bd9cbdd9cd3ab56d1e Mon Sep 17 00:00:00 2001 From: Lucia Quirke Date: Tue, 28 Jul 2026 08:18:14 +0000 Subject: [PATCH] Add non-gradient attribution baselines on a leave-k-out bank Two content-similarity baselines scored against a re-train bank and evaluated by per-query LDS, as a floor for the gradient methods: - activation similarity over the attribution target modules (uses the bank model's activations) - semantic search with jina-embeddings-v3 (pure text similarity) Point either at any bank written with save_retrained_models=true via --bank; the model and dataset are read from the bank's config.yaml, and LDS reuses evaluate_retrained's cached per-subset losses. Omit --bank to build the default GPT-2/WikiText bank first. On that bank they score mean per-query Spearman 0.13 and 0.09, below TrackStar (0.18), SOURCE (~0.39) and MAGIC (~0.92). Co-Authored-By: Claude Opus 4.8 --- examples/bank_baselines/README.md | 41 +++++ examples/bank_baselines/__init__.py | 0 .../bank_baselines/activation_baseline.py | 153 +++++++++++++++++ examples/bank_baselines/bm25_baseline.py | 69 ++++++++ examples/bank_baselines/common.py | 156 ++++++++++++++++++ examples/bank_baselines/gradient_baseline.py | 94 +++++++++++ examples/bank_baselines/qwen3_baseline.py | 76 +++++++++ examples/bank_baselines/semantic_baseline.py | 102 ++++++++++++ examples/magic/gpt2_wikitext_bank.yaml | 43 +++++ 9 files changed, 734 insertions(+) create mode 100644 examples/bank_baselines/README.md create mode 100644 examples/bank_baselines/__init__.py create mode 100644 examples/bank_baselines/activation_baseline.py create mode 100644 examples/bank_baselines/bm25_baseline.py create mode 100644 examples/bank_baselines/common.py create mode 100644 examples/bank_baselines/gradient_baseline.py create mode 100644 examples/bank_baselines/qwen3_baseline.py create mode 100644 examples/bank_baselines/semantic_baseline.py create mode 100644 examples/magic/gpt2_wikitext_bank.yaml diff --git a/examples/bank_baselines/README.md b/examples/bank_baselines/README.md new file mode 100644 index 00000000..98765241 --- /dev/null +++ b/examples/bank_baselines/README.md @@ -0,0 +1,41 @@ +# Non-gradient attribution baselines + +Text similarity baselines evaluated by mean LDS over 50 test queries. + +- `bm25_baseline.py` — BM25 lexical overlap: pure surface-form term overlap, no model or embedding. +- `gradient_baseline.py` — gradient cosine similarity: cosine of the full per-example loss gradients on the bank's model (TracIn-style, unpreconditioned). +- `activation_baseline.py` — activation similarity: each doc is the mean-pooled input activation to every linear matrix of the model, L2-normalized per matrix and concatenated, then cosine similarity. +- `semantic_baseline.py` — semantic search with `jinaai/jina-embeddings-v3` (asymmetric `retrieval.query`/`retrieval.passage`). +- `qwen3_baseline.py` — semantic search with `Qwen/Qwen3-Embedding-8B`, a SOTA decoder embedder (`--model` to swap). + +Each produces a `[num_train_docs, num_queries]` score matrix. + +## Running + +Point a baseline at a re-train bank (written with `save_retrained_models=true`); it reads the model and dataset from the bank's `config.yaml`: + +```bash +python -m examples.bank_baselines.bm25_baseline --bank runs/retrain_bank_path +python -m examples.bank_baselines.gradient_baseline --bank runs/retrain_bank_path +python -m examples.bank_baselines.activation_baseline --bank runs/retrain_bank_path +python -m examples.bank_baselines.semantic_baseline --bank runs/retrain_bank_path +python -m examples.bank_baselines.qwen3_baseline --bank runs/retrain_bank_path +``` + +Omit `--bank` to build the default GPT-2/WikiText bank (`examples/magic/gpt2_wikitext_bank.yaml`) first. `--query_split` sets the query set (default `test[1:51]`), `--out` the output dir (default `runs/bank_baselines/`). + +## Results + +GPT-2 / WikiText bank (100 subsets, 1% leave-out, `eps_root 1e-8`): + +| Method | mean ρ | +| --- | --- | +| BM25 lexical overlap | 0.16 | +| Qwen3-Embedding-8B | 0.11 | +| activation similarity | 0.09 | +| Jina v3 semantic search | 0.06 | +| gradient cosine similarity | 0.05 | + +## Notes + +`jina-embeddings-v3`'s custom code predates transformers 5.x; `semantic_baseline.load_model` sets an `all_tied_weights_keys` default and resets the NaN LoRA `lora_dropout_mask` buffers to ones so it loads and runs. NVIDIA's NV-Embed-v2 is a comparable SOTA embedder but its custom code is incompatible with transformers 5.x, so `qwen3_baseline.py` uses Qwen3-Embedding instead. diff --git a/examples/bank_baselines/__init__.py b/examples/bank_baselines/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/bank_baselines/activation_baseline.py b/examples/bank_baselines/activation_baseline.py new file mode 100644 index 00000000..ac99fb0e --- /dev/null +++ b/examples/bank_baselines/activation_baseline.py @@ -0,0 +1,153 @@ +"""Baseline: activation similarity over the attribution target modules. + +Represents each document by the mean-pooled input activation to every target +linear module of the bank's fully-trained model (the same modules the gradient +methods attribute through: all Conv1D/Linear except ``lm_head``). Each +per-module pooled vector is L2-normalized, then concatenated, so cosine +similarity between two docs equals the average per-module cosine similarity. + +A training doc that is activation-similar to a query is predicted to be +influential (removing it should raise query loss), so the loss-diff-convention +score is ``-cosine`` -- larger => larger predicted loss reduction from keeping +the doc. + +Run with (builds the default bank if --bank is omitted): + python -m examples.bank_baselines.activation_baseline --bank runs/retrain_bank_path +""" + +import argparse +from pathlib import Path + +import numpy as np +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.pytorch_utils import Conv1D + +from . import common + + +def target_module_names(model: torch.nn.Module) -> list[str]: + """Linear/Conv1D modules the gradient methods attribute through (not lm_head).""" + names = [] + for name, mod in model.named_modules(): + if name.endswith("lm_head"): + continue + if isinstance(mod, (Conv1D, torch.nn.Linear)): + names.append(name) + return names + + +@torch.no_grad() +def embed( + texts: list[str], + model: torch.nn.Module, + tokenizer, + module_names: list[str], + device: str, + max_len: int, + batch_size: int, +) -> np.ndarray: + """Per-module-normalized, concatenated mean-pooled activation embeddings.""" + modules = dict(model.named_modules()) + captured: dict[str, torch.Tensor] = {} + handles = [] + for mn in module_names: + + def hook(_m, inp, _out, _mn=mn): + captured[_mn] = inp[0].detach() + + handles.append(modules[mn].register_forward_hook(hook)) + + out_rows = [] + try: + for start in range(0, len(texts), batch_size): + batch = texts[start : start + batch_size] + enc = tokenizer( + batch, + return_tensors="pt", + padding=True, + truncation=True, + max_length=max_len, + ).to(device) + mask = enc["attention_mask"].unsqueeze(-1).float() # [B, T, 1] + captured.clear() + model(**enc) + + per_module = [] + denom = mask.sum(dim=1).clamp_min(1.0) # [B, 1] + for mn in module_names: + act = captured[mn].float() # [B, T, d] + pooled = (act * mask).sum(dim=1) / denom # [B, d] + pooled = torch.nn.functional.normalize(pooled, dim=-1) + per_module.append(pooled) + emb = torch.cat(per_module, dim=-1) # [B, sum_d] + out_rows.append(emb.cpu().numpy()) + finally: + for h in handles: + h.remove() + return np.concatenate(out_rows, axis=0) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--bank", default=None, help="re-train bank dir; built if omitted") + ap.add_argument("--query_split", default=common.DEFAULT_QUERY_SPLIT) + ap.add_argument("--out", default=str(common.REPO / "runs" / "bank_baselines")) + ap.add_argument("--device", default="cuda:0") + ap.add_argument("--max_len", type=int, default=1024) + ap.add_argument("--batch_size", type=int, default=32) + args = ap.parse_args() + + bank = common.ensure_bank(args.bank) + spec = common.read_bank_spec(bank) + out_dir = Path(args.out) + + tokenizer = AutoTokenizer.from_pretrained(spec.model) + tokenizer.pad_token = tokenizer.eos_token + model = AutoModelForCausalLM.from_pretrained( + spec.base_model, dtype=torch.float32, attn_implementation="eager" + ).to(args.device) + model.eval() + + module_names = target_module_names(model) + print(f"Embedding over {len(module_names)} modules, e.g. {module_names[:3]} ...") + + train_texts, query_texts = common.load_texts(spec, args.query_split) + print(f"Embedding {len(train_texts)} train docs ...") + train_emb = embed( + train_texts, + model, + tokenizer, + module_names, + args.device, + args.max_len, + args.batch_size, + ) + print(f"Embedding {len(query_texts)} query docs ...") + query_emb = embed( + query_texts, + model, + tokenizer, + module_names, + args.device, + args.max_len, + args.batch_size, + ) + + # Cosine similarity (rows already unit-norm per module; renormalize the + # concatenation so cosine == mean per-module cosine). + tn = train_emb / np.linalg.norm(train_emb, axis=1, keepdims=True) + qn = query_emb / np.linalg.norm(query_emb, axis=1, keepdims=True) + cosine = tn @ qn.T # [num_train, num_query] + + scores = -cosine # loss-diff convention (similar => influential) + score_path = common.save_scores(scores, out_dir, "activation_scores") + + rhos = common.evaluate_lds( + bank, score_path, out_dir / "activation_validate", spec, args.query_split + ) + common.report("activation similarity", rhos) + + +if __name__ == "__main__": + main() diff --git a/examples/bank_baselines/bm25_baseline.py b/examples/bank_baselines/bm25_baseline.py new file mode 100644 index 00000000..cf7f8928 --- /dev/null +++ b/examples/bank_baselines/bm25_baseline.py @@ -0,0 +1,69 @@ +"""Baseline: BM25 lexical overlap. + +Scores each training doc against each query by Okapi BM25 -- pure surface-form +term overlap, no model and no learned embedding. A lexical proxy for influence: +docs sharing query terms are predicted influential, so the loss-diff-convention +score is ``-bm25``. Often a stronger influence proxy than deep-semantic +embedders for small LMs, and essentially free. + +Run with (builds the default bank if --bank is omitted): + python -m examples.bank_baselines.bm25_baseline --bank runs/retrain_bank_path +""" + +import argparse +from pathlib import Path + +import numpy as np +from sklearn.feature_extraction.text import CountVectorizer + +from . import common + + +def bm25_scores( + train_texts: list[str], query_texts: list[str], k1: float = 1.5, b: float = 0.75 +) -> np.ndarray: + """Okapi BM25 score of every train doc against every query -> [docs, queries].""" + vectorizer = CountVectorizer() + tf = vectorizer.fit_transform(train_texts).astype(np.float64).tocsr() # [D, V] + n_docs, _ = tf.shape + + doc_len = np.asarray(tf.sum(axis=1)).ravel() + avgdl = doc_len.mean() + df = np.asarray((tf > 0).sum(axis=0)).ravel() + idf = np.log((n_docs - df + 0.5) / (df + 0.5) + 1.0) + + # Per-nonzero BM25 term weight: idf * tf*(k1+1) / (tf + k1*(1-b+b*dl/avgdl)). + rows, cols = tf.nonzero() + denom_bias = k1 * (1 - b + b * doc_len[rows] / avgdl) + weighted = tf.copy() + weighted.data = idf[cols] * tf.data * (k1 + 1.0) / (tf.data + denom_bias) + + # Query terms as a binary term set (Okapi BM25 ignores query term frequency). + q_bin = (vectorizer.transform(query_texts) > 0).astype(np.float64) # [Q, V] + return np.asarray((weighted @ q_bin.T).todense()) # [D, Q] + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--bank", default=None, help="re-train bank dir; built if omitted") + ap.add_argument("--query_split", default=common.DEFAULT_QUERY_SPLIT) + ap.add_argument("--out", default=str(common.REPO / "runs" / "bank_baselines")) + args = ap.parse_args() + + bank = common.ensure_bank(args.bank) + spec = common.read_bank_spec(bank) + out_dir = Path(args.out) + + train_texts, query_texts = common.load_texts(spec, args.query_split) + print(f"BM25 over {len(train_texts)} train docs, {len(query_texts)} queries ...") + scores = -bm25_scores(train_texts, query_texts) # loss-diff convention + score_path = common.save_scores(scores, out_dir, "bm25_scores") + + rhos = common.evaluate_lds( + bank, score_path, out_dir / "bm25_validate", spec, args.query_split + ) + common.report("BM25 lexical overlap", rhos) + + +if __name__ == "__main__": + main() diff --git a/examples/bank_baselines/common.py b/examples/bank_baselines/common.py new file mode 100644 index 00000000..2da7038b --- /dev/null +++ b/examples/bank_baselines/common.py @@ -0,0 +1,156 @@ +"""Shared helpers for non-gradient attribution baselines on a re-train bank. + +A "re-train bank" is any directory written with ``save_retrained_models=true`` +(e.g. by ``examples/magic/gpt2_wikitext_bank.yaml``): it holds ``subsets.json``, +``retrained/base`` and ``retrained/subset_*`` checkpoints, and the ``config.yaml`` +that built it. Each baseline reads the model + dataset straight from the bank, +produces a ``[num_train_docs, num_queries]`` score matrix, and evaluates it by +per-query LDS -- so ``--bank `` is all anyone needs to run one. + +The LDS itself reuses ``bergson``'s ``evaluate_retrained``: it computes each +subset's query-loss diff from the banked models once, caches it under the bank, +and every later baseline on the same bank/query reuses that cache. +""" + +import subprocess +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import yaml +from datasets import load_dataset + +from bergson.config.config import DataConfig +from bergson.magic.config import MagicConfig +from bergson.validate import evaluate_retrained + +REPO = Path(__file__).resolve().parents[2] + +# The default bank built on demand when ``--bank`` is omitted. +DEFAULT_BANK = REPO / "runs" / "gpt2_wikitext_bank" +DEFAULT_BANK_CONFIG = REPO / "examples" / "magic" / "gpt2_wikitext_bank.yaml" +# LDS query set (50 held-out docs); overridable per run. +DEFAULT_QUERY_SPLIT = "test[1:51]" + + +@dataclass +class BankSpec: + """The model + data a bank was built on, read from its ``config.yaml``.""" + + model: str + base_model: str # retrained/base if present, else the base model id + dataset: str + train_split: str + prompt_column: str + chunk_length: int + batch_size: int + + +def ensure_bank(bank: str | None) -> Path: + """Return a ready re-train bank, building the default one if none is given. + + A bank is ready when it has ``subsets.json`` and ``retrained/base``. When + ``bank`` is ``None`` and the default bank is absent, it is built by running + its config through ``bergson`` (a leave-k-out re-training job). + """ + if bank is not None: + path = Path(bank) + if not (path / "subsets.json").exists(): + raise FileNotFoundError( + f"{path} is not a re-train bank (no subsets.json); pass a dir " + "written with save_retrained_models=true" + ) + return path + + if not (DEFAULT_BANK / "subsets.json").exists(): + cmd = ["bergson", str(DEFAULT_BANK_CONFIG)] + print( + f"No bank passed and {DEFAULT_BANK} missing; building it:\n" + f" {' '.join(cmd)}" + ) + subprocess.run(cmd, cwd=REPO, check=True) + return DEFAULT_BANK + + +def read_bank_spec(bank: Path) -> BankSpec: + """Read the model + dataset a bank was built on from its ``config.yaml``.""" + with open(bank / "config.yaml") as f: + config = yaml.safe_load(f) + # The bank-building step (magic/train) is the one carrying model + data. + step = next( + v + for s in config["steps"] + for v in s.values() + if isinstance(v, dict) and "model" in v and "data" in v + ) + data = step["data"] + base = bank / "retrained" / "base" + return BankSpec( + model=step["model"], + base_model=str(base) if base.exists() else step["model"], + dataset=data["dataset"], + train_split=data.get("split", "train"), + prompt_column=data.get("prompt_column", "text"), + chunk_length=data.get("chunk_length", 0), + batch_size=step.get("batch_size", 64), + ) + + +def load_texts(spec: BankSpec, query_split: str) -> tuple[list[str], list[str]]: + """Return (train_texts, query_texts) in original doc order. + + Row i of ``train_texts`` is training doc id i, matching the ids in the + bank's ``subsets.json`` and the rows of every score matrix. + """ + train = load_dataset(spec.dataset, split=spec.train_split) + query = load_dataset(spec.dataset, split=query_split) + return list(train[spec.prompt_column]), list(query[spec.prompt_column]) + + +def evaluate_lds( + bank: Path, + score_path: Path, + out_dir: Path, + spec: BankSpec, + query_split: str, +) -> np.ndarray: + """Per-query LDS Spearman of a score matrix against the bank. + + Runs ``evaluate_retrained`` (no re-training; reuses the bank's cached + per-subset query losses) and returns the per-query Spearman correlations it + writes to ``summary.csv``. + """ + run_cfg = MagicConfig( + run_path=str(out_dir), + model=spec.model, + precision="fp32", + batch_size=spec.batch_size, + query=DataConfig( + dataset=spec.dataset, + split=query_split, + prompt_column=spec.prompt_column, + chunk_length=0, + ), + ) + evaluate_retrained(run_cfg, str(bank), score_path=str(score_path)) + + summary = np.genfromtxt(out_dir / "summary.csv", delimiter=",", names=True) + return np.atleast_1d(summary["spearman_corr"]).astype(float) + + +def save_scores(scores: np.ndarray, out_dir: Path, name: str) -> Path: + """Save a ``[num_train_docs, num_queries]`` score matrix as ``.npy``.""" + out_dir.mkdir(parents=True, exist_ok=True) + path = out_dir / f"{name}.npy" + np.save(path, scores.astype(np.float32)) + print(f"Saved scores {scores.shape} -> {path}") + return path + + +def report(name: str, rhos: np.ndarray) -> None: + """Print a per-query LDS summary consistent with the repo's convention.""" + print(f"\n=== {name}: per-query LDS Spearman (n={len(rhos)} queries) ===") + print(f" mean {np.mean(rhos):+.4f}") + print(f" median {np.median(rhos):+.4f}") + print(f" min {np.min(rhos):+.4f} max {np.max(rhos):+.4f}") + print(f" frac>0 {np.mean(rhos > 0):.2f}") diff --git a/examples/bank_baselines/gradient_baseline.py b/examples/bank_baselines/gradient_baseline.py new file mode 100644 index 00000000..ccf9a540 --- /dev/null +++ b/examples/bank_baselines/gradient_baseline.py @@ -0,0 +1,94 @@ +"""Baseline: gradient cosine similarity (TracIn-style, full parameters). + +Scores a train doc against a query by the cosine similarity of their full +per-example loss gradients on the bank's model -- the raw, unpreconditioned +influence signal that TrackStar/EK-FAC refine. A train doc whose gradient +aligns with the query's is predicted influential, so the loss-diff-convention +score is ``-cosine``. + +Full gradients are never stored: the 50 query gradients stay resident on the +GPU (GPT-2 is 124M params, ~25 GB in fp32) and each train-doc gradient is +computed, dotted against all queries, and discarded. + +Run with (builds the default bank if --bank is omitted): + python -m examples.bank_baselines.gradient_baseline --bank runs/retrain_bank_path +""" + +import argparse +from pathlib import Path + +import numpy as np +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + +from . import common + + +def per_example_grad( + text: str, model, tokenizer, device: str, max_len: int +) -> torch.Tensor: + """Flattened full-parameter gradient of the mean-CE loss on one document.""" + enc = tokenizer(text, return_tensors="pt", truncation=True, max_length=max_len).to( + device + ) + model.zero_grad(set_to_none=True) + model(input_ids=enc["input_ids"], labels=enc["input_ids"]).loss.backward() + return torch.cat( + [ + (p.grad if p.grad is not None else torch.zeros_like(p)).flatten() + for p in model.parameters() + ] + ) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--bank", default=None, help="re-train bank dir; built if omitted") + ap.add_argument("--query_split", default=common.DEFAULT_QUERY_SPLIT) + ap.add_argument("--out", default=str(common.REPO / "runs" / "bank_baselines")) + ap.add_argument("--device", default="cuda:0") + ap.add_argument("--max_len", type=int, default=1024) + args = ap.parse_args() + + bank = common.ensure_bank(args.bank) + spec = common.read_bank_spec(bank) + out_dir = Path(args.out) + + tokenizer = AutoTokenizer.from_pretrained(spec.model) + model = AutoModelForCausalLM.from_pretrained( + spec.base_model, dtype=torch.float32, attn_implementation="eager" + ).to(args.device) + model.eval() + + train_texts, query_texts = common.load_texts(spec, args.query_split) + + # Keep the query gradients resident; stream train gradients against them. + # Pre-allocate and fill row by row so we never hold a second copy (each full + # gradient is ~0.5 GB, the [Q, P] block ~25 GB in fp32). + print(f"Computing {len(query_texts)} query gradients ...") + n_params = sum(p.numel() for p in model.parameters()) + query_grads = torch.empty(len(query_texts), n_params, device=args.device) + for i, t in enumerate(query_texts): + query_grads[i] = per_example_grad( + t, model, tokenizer, args.device, args.max_len + ) + query_norms = query_grads.norm(dim=1).clamp_min(1e-12) + + print(f"Scoring {len(train_texts)} train gradients against queries ...") + scores = np.empty((len(train_texts), len(query_texts)), dtype=np.float32) + for i, text in enumerate(train_texts): + g = per_example_grad(text, model, tokenizer, args.device, args.max_len) + cos = (query_grads @ g) / (query_norms * g.norm().clamp_min(1e-12)) + scores[i] = cos.detach().cpu().numpy() + + scores = -scores # loss-diff convention (aligned gradient => influential) + score_path = common.save_scores(scores, out_dir, "gradient_scores") + + rhos = common.evaluate_lds( + bank, score_path, out_dir / "gradient_validate", spec, args.query_split + ) + common.report("gradient cosine similarity", rhos) + + +if __name__ == "__main__": + main() diff --git a/examples/bank_baselines/qwen3_baseline.py b/examples/bank_baselines/qwen3_baseline.py new file mode 100644 index 00000000..42576ace --- /dev/null +++ b/examples/bank_baselines/qwen3_baseline.py @@ -0,0 +1,76 @@ +"""Baseline: semantic search with Qwen3-Embedding-8B. + +A SOTA decoder-based embedding model (causal-attention Qwen3 backbone, +last-token pooling), top of the open MTEB leaderboard -- a much stronger, +attention-based embedder than the jina-v3 encoder. Training docs are embedded +as passages and queries with the retrieval prompt; a train doc scores against a +query by cosine similarity, so the loss-diff-convention score is ``-cosine``. + +(NVIDIA's NV-Embed-v2 is a comparable model but its custom modeling code is +incompatible with transformers 5.x; Qwen3-Embedding uses the native Qwen3 +architecture and loads cleanly. Pass --model Qwen/Qwen3-Embedding-4B for the +smaller, faster variant.) + +Run with (builds the default bank if --bank is omitted): + python -m examples.bank_baselines.qwen3_baseline --bank runs/retrain_bank_path +""" + +import argparse +from pathlib import Path + +from sentence_transformers import SentenceTransformer + +from . import common + +MODEL = "Qwen/Qwen3-Embedding-8B" + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--bank", default=None, help="re-train bank dir; built if omitted") + ap.add_argument("--query_split", default=common.DEFAULT_QUERY_SPLIT) + ap.add_argument("--out", default=str(common.REPO / "runs" / "bank_baselines")) + ap.add_argument("--model", default=MODEL) + ap.add_argument("--device", default="cuda:0") + ap.add_argument("--max_length", type=int, default=512) + ap.add_argument("--batch_size", type=int, default=8) + args = ap.parse_args() + + bank = common.ensure_bank(args.bank) + spec = common.read_bank_spec(bank) + out_dir = Path(args.out) + + model = SentenceTransformer( + args.model, model_kwargs={"torch_dtype": "bfloat16"}, device=args.device + ) + model.max_seq_length = args.max_length + + train_texts, query_texts = common.load_texts(spec, args.query_split) + print(f"Embedding {len(train_texts)} train docs (passages) ...") + train_emb = model.encode( + train_texts, + batch_size=args.batch_size, + normalize_embeddings=True, + show_progress_bar=True, + ) + print(f"Embedding {len(query_texts)} query docs (query prompt) ...") + query_emb = model.encode( + query_texts, + prompt_name="query", + batch_size=args.batch_size, + normalize_embeddings=True, + show_progress_bar=True, + ) + + cosine = train_emb @ query_emb.T # rows unit-norm => dot == cosine + scores = -cosine # loss-diff convention (similar => influential) + score_path = common.save_scores(scores, out_dir, "qwen3_scores") + + rhos = common.evaluate_lds( + bank, score_path, out_dir / "qwen3_validate", spec, args.query_split + ) + common.report(f"{args.model} semantic similarity", rhos) + + +if __name__ == "__main__": + main() diff --git a/examples/bank_baselines/semantic_baseline.py b/examples/bank_baselines/semantic_baseline.py new file mode 100644 index 00000000..2103251d --- /dev/null +++ b/examples/bank_baselines/semantic_baseline.py @@ -0,0 +1,102 @@ +"""Baseline: semantic-search similarity with a Jina AI embedding model. + +Embeds each training document as a retrieval passage and each query document +as a retrieval query with ``jinaai/jina-embeddings-v3`` (a strong 570M-param, +1024-dim, 8192-context text embedder), then scores a train doc against a query +by their cosine similarity -- ordinary dense semantic search. This ignores the +attributed model entirely; it is a pure content-similarity baseline. + +A training doc semantically similar to a query is predicted to be influential +(removing it should raise query loss), so the loss-diff-convention score is +``-cosine``. + +jina-embeddings-v3 ships custom modeling code that predates transformers 5.x, +so ``load_model`` patches two load-time incompatibilities (see there) to run it +for inference. + +Run with (builds the default bank if --bank is omitted): + python -m examples.bank_baselines.semantic_baseline --bank runs/retrain_bank_path +""" + +import argparse +from pathlib import Path + +import numpy as np +import torch +from transformers.modeling_utils import PreTrainedModel + +from . import common + +MODEL = "jinaai/jina-embeddings-v3" + + +def load_model(device: str): + # jina-embeddings-v3's custom code predates transformers 5.x; two fixes: + # (1) from_pretrained reads all_tied_weights_keys, which the custom class + # never defines -- give a benign empty default so load doesn't crash. + if not isinstance( + getattr(PreTrainedModel, "all_tied_weights_keys", None), property + ): + PreTrainedModel.all_tied_weights_keys = {} + from transformers import AutoModel + + model = AutoModel.from_pretrained( + MODEL, trust_remote_code=True, dtype=torch.float32 + ) + + # (2) its LoRA task adapters leave the per-forward lora_dropout_mask buffers + # uninitialized (NaN), so every embedding comes out NaN. In eval there + # is no dropout, so reset them to ones. + for name, buf in model.named_buffers(): + if "lora_dropout_mask" in name: + buf.data = torch.ones_like(buf) + return model.to(device).eval() + + +@torch.no_grad() +def encode(model, texts: list[str], task: str, batch_size: int) -> np.ndarray: + """L2-normalized embeddings via jina's task-specific encode API.""" + out = [] + for start in range(0, len(texts), batch_size): + emb = model.encode( + texts[start : start + batch_size], + task=task, + convert_to_numpy=True, + normalize_embeddings=True, + ) + out.append(np.asarray(emb, dtype=np.float32)) + return np.concatenate(out, axis=0) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--bank", default=None, help="re-train bank dir; built if omitted") + ap.add_argument("--query_split", default=common.DEFAULT_QUERY_SPLIT) + ap.add_argument("--out", default=str(common.REPO / "runs" / "bank_baselines")) + ap.add_argument("--device", default="cuda:0") + ap.add_argument("--batch_size", type=int, default=16) + args = ap.parse_args() + + bank = common.ensure_bank(args.bank) + spec = common.read_bank_spec(bank) + out_dir = Path(args.out) + model = load_model(args.device) + + train_texts, query_texts = common.load_texts(spec, args.query_split) + print(f"Embedding {len(train_texts)} train docs (retrieval.passage) ...") + train_emb = encode(model, train_texts, "retrieval.passage", args.batch_size) + print(f"Embedding {len(query_texts)} query docs (retrieval.query) ...") + query_emb = encode(model, query_texts, "retrieval.query", args.batch_size) + + cosine = train_emb @ query_emb.T # rows unit-norm => dot == cosine + scores = -cosine # loss-diff convention (similar => influential) + score_path = common.save_scores(scores, out_dir, "semantic_scores") + + rhos = common.evaluate_lds( + bank, score_path, out_dir / "semantic_validate", spec, args.query_split + ) + common.report("Jina v3 semantic similarity", rhos) + + +if __name__ == "__main__": + main() diff --git a/examples/magic/gpt2_wikitext_bank.yaml b/examples/magic/gpt2_wikitext_bank.yaml new file mode 100644 index 00000000..f47758b4 --- /dev/null +++ b/examples/magic/gpt2_wikitext_bank.yaml @@ -0,0 +1,43 @@ +# Build a GPT-2 / WikiText leave-k-out re-train bank (single node): MAGIC +# attribution plus 400 retrained leave-out models saved for later evaluation. + +# Run with `bergson examples/magic/gpt2_wikitext_bank.yaml` + +steps: + - magic: + run_path: runs/gpt2_wikitext_bank + model: gpt2 + overwrite: true + cleanup_ckpts: false + + data: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "train" + chunk_length: 0 + + query: + dataset: EleutherAI/bergson-wikitext-512-chunks + split: "test[0:1]" + chunk_length: 0 + + distributed: + nproc_per_node: 8 + nnode: 1 + + eps_root: 1e-8 + + batch_size: 64 + num_epochs: 4 + lr_schedule: + lr_scheduler_type: polynomial + lr: 0.0008 + lr_start: 1e-6 + lr_end: 0.00008 + warmup_steps: 0.25 + + subset_strategy: random + wandb_project: magic + num_subsets: 100 + subset_fraction: 0.01 + save_optimizer_state: True + save_retrained_models: True