From 225ef83ad1e3e65023a67cfcbb89dd8cceb000e4 Mon Sep 17 00:00:00 2001 From: Lukas Scheucher Date: Sat, 6 Jun 2026 20:24:04 +0200 Subject: [PATCH] feat(contextual): accumulator, streaming, finalize, and orchestrator Part 2/3 of the contextual word-drift pipeline (issue #34). The data path: re-stream FineWeb -> frozen-encoder forward -> on-the-fly per-word contextual centroids -> per-year finalized embeddings. Never stores per-token vectors. - contextual_aggregate.py: ContextualAccumulator holds GPU sum/sumsq/count (fp32, [V,768]). add_batch() finds first-subword positions via word_ids, resolves each word to its shared-vocab id (skipping stopwords/OOV/short/ digits, never folding into ), and folds the batch in with vectorized index_add_. Atomic npz save/load for resume. - contextual_stream.py: text_to_windows + iter_snapshot_windows re-stream FineWeb (language_score>=0.65, raw text, non-overlapping 128-word windows); run_year batches windows through the encoder into the accumulator, checkpointing after each completed snapshot (config-hash guarded). - contextual_finalize.py: centroid=sum/count, dispersion=Sum(sumsq/count - centroid^2), per-year mean-centering (anisotropy fix) + optional all-but-top-k PCA removal. Writes {year}.npy/_centered/_count/_dispersion. - scripts/contextual_drift.py: orchestrator with --stage {stream,finalize,all}, --smoke/--years/--model/--target-tokens. Resumable; on year completion drops the checkpoint but keeps the partial so finalize can be re-tuned cheaply. Verified via --smoke (2 snapshots/year): - artifacts shape (119466, 768) fp32; stopwords skipped (the=0); OOV skipped. - anisotropy fix works: raw random-pair cosine 0.47 -> centered ~0.000. - resume works: killed after snapshot 1, relaunch resumed at 300k and continued to 600k words; checkpoint deleted + partial kept on completion; --stage finalize re-derives from the partial without re-streaming. The drift stage + local eval land in Part 3 (analysis/contextual_drift.py). Co-Authored-By: Claude Opus 4.8 (1M context) --- pipeline/contextual_aggregate.py | 126 +++++++++++++++++++++++ pipeline/contextual_finalize.py | 103 +++++++++++++++++++ pipeline/contextual_stream.py | 167 +++++++++++++++++++++++++++++++ scripts/contextual_drift.py | 148 +++++++++++++++++++++++++++ 4 files changed, 544 insertions(+) create mode 100644 pipeline/contextual_aggregate.py create mode 100644 pipeline/contextual_finalize.py create mode 100644 pipeline/contextual_stream.py create mode 100644 scripts/contextual_drift.py diff --git a/pipeline/contextual_aggregate.py b/pipeline/contextual_aggregate.py new file mode 100644 index 000000000..20a00f12c --- /dev/null +++ b/pipeline/contextual_aggregate.py @@ -0,0 +1,126 @@ +"""On-the-fly contextual centroid accumulator. + +Never stores per-token vectors. Holds three preallocated GPU tensors keyed by +the shared vocab id -- `sum[V,H]`, `sumsq[V,H]`, `count[V]` -- and folds each +forward-pass batch in with vectorized `index_add_`. Accumulation is fp32 even +though the forward runs bf16. + +Pooling: **first-subword** (v1). For each word we take the hidden state of its +first WordPiece. Words that are stopwords / OOV (not in the shared vocab) / +shorter than MIN_WORD_LENGTH / all-digits are ignored (never accumulated, and +never folded into =0). +""" +from __future__ import annotations + +import os +import tempfile +from pathlib import Path + +import numpy as np +import torch + +from config import CONTEXTUAL_USE_STOPWORDS, CONTEXTUAL_STOPWORDS, MIN_WORD_LENGTH + + +class ContextualAccumulator: + def __init__( + self, + vocab: dict[str, int], + hidden_dim: int, + device: str = "cuda", + use_stopwords: bool = CONTEXTUAL_USE_STOPWORDS, + ) -> None: + self.vocab = vocab + self.V = len(vocab) + self.H = hidden_dim + self.device = device + self.use_stopwords = use_stopwords + self.stopwords = CONTEXTUAL_STOPWORDS if use_stopwords else frozenset() + + self.sum = torch.zeros((self.V, self.H), dtype=torch.float32, device=device) + self.sumsq = torch.zeros((self.V, self.H), dtype=torch.float32, device=device) + self.count = torch.zeros((self.V,), dtype=torch.float32, device=device) + + # Cache of word -> target id (or -1 if filtered), to avoid re-filtering + # the same string every batch. + self._target_cache: dict[str, int] = {} + + def _word_target(self, word: str) -> int: + cached = self._target_cache.get(word, -2) + if cached != -2: + return cached + if ( + len(word) < MIN_WORD_LENGTH + or word.isdigit() + or (self.use_stopwords and word in self.stopwords) + ): + tid = -1 + else: + tid = self.vocab.get(word, -1) # OOV -> -1 (never folded into =0) + self._target_cache[word] = tid + return tid + + @torch.inference_mode() + def add_batch( + self, + hidden: torch.Tensor, # [B, L, H], model dtype + word_ids: torch.Tensor, # [B, L] long, -1 for special/pad + windows: list[list[str]], # the B source word-windows + ) -> int: + """Fold one forward-pass batch into the accumulators. Returns #observations added.""" + B, L = word_ids.shape + + # First subword of each word = position where word_ids changes (and is real). + shifted = torch.full_like(word_ids, -2) + shifted[:, 1:] = word_ids[:, :-1] + is_first = (word_ids != shifted) & (word_ids >= 0) + + # Per-word target ids, padded to the batch's longest window. + w_max = max(len(w) for w in windows) + word_targets = torch.full((B, w_max), -1, dtype=torch.long, device=self.device) + for b, win in enumerate(windows): + row = [self._word_target(w) for w in win] + if row: + word_targets[b, : len(row)] = torch.tensor(row, device=self.device) + + # Gather each subword position's word target, keep only first subwords. + wid_clamped = word_ids.clamp(min=0) + gathered = torch.gather(word_targets, 1, wid_clamped) # [B, L] + target = torch.where(is_first, gathered, torch.full_like(gathered, -1)) + + flat_target = target.reshape(-1) + valid = flat_target >= 0 + if not bool(valid.any()): + return 0 + + idx = flat_target[valid] + h = hidden.reshape(-1, self.H).float()[valid] # fp32 accumulation + + self.sum.index_add_(0, idx, h) + self.sumsq.index_add_(0, idx, h * h) + self.count += torch.bincount(idx, minlength=self.V).float() + return int(idx.shape[0]) + + def to_numpy(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + return ( + self.sum.cpu().numpy(), + self.sumsq.cpu().numpy(), + self.count.cpu().numpy(), + ) + + def save_state(self, path: Path) -> None: + """Atomic npz write of sum/sumsq/count (resumable partial state).""" + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + s, sq, c = self.to_numpy() + fd, tmp = tempfile.mkstemp(dir=path.parent, suffix=".tmp.npz") + os.close(fd) + np.savez(tmp, sum=s, sumsq=sq, count=c) + # np.savez appends .npz to a path without the suffix; mkstemp already has it. + os.replace(tmp, path) + + def load_state(self, path: Path) -> None: + z = np.load(path) + self.sum = torch.from_numpy(z["sum"]).to(self.device) + self.sumsq = torch.from_numpy(z["sumsq"]).to(self.device) + self.count = torch.from_numpy(z["count"]).to(self.device) diff --git a/pipeline/contextual_finalize.py b/pipeline/contextual_finalize.py new file mode 100644 index 000000000..90c82b7a7 --- /dev/null +++ b/pipeline/contextual_finalize.py @@ -0,0 +1,103 @@ +"""Turn raw accumulators into per-year contextual embeddings. + +From `sum[V,H]`, `sumsq[V,H]`, `count[V]` (one year) we derive: + - centroid = sum / count (mean contextual vector) + - dispersion = sum_d(sumsq/count - centroid^2) (E||h - c||^2; polysemy signal) + - centered = centroid - per-year mean (BERT anisotropy fix) + [- optional all-but-top-k PCA removal] + +BERT's last layer is strongly anisotropic (random words sit at cos ~0.3-0.6), +which would swamp the drift signal, so cosine drift downstream runs on the +*centered* vectors. Dispersion is computed on raw centroids and is invariant to +centering (subtracting a constant per dim doesn't change variance). + +Writes `models/contextual/{year}.npy` (raw centroid), `_centered.npy`, +`_count.npy`, `_dispersion.npy`. +""" +from __future__ import annotations + +from pathlib import Path + +import numpy as np + +from config import ( + CONTEXTUAL_CENTERING, + CONTEXTUAL_MIN_COUNT, + CONTEXTUAL_PCA_REMOVE_K, +) + + +def load_partial(state_dir: Path, year: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + z = np.load(Path(state_dir) / f"{year}.partial.npz") + return z["sum"], z["sumsq"], z["count"] + + +def _remove_top_k(centered: np.ndarray, mask: np.ndarray, k: int) -> np.ndarray: + """All-but-the-top: project out the top-k principal components (fit on trusted rows).""" + trusted = centered[mask] + if trusted.shape[0] <= k: + return centered + # centered already has ~zero mean; SVD gives principal directions in Vt. + _, _, vt = np.linalg.svd(trusted - trusted.mean(axis=0), full_matrices=False) + comps = vt[:k] # [k, H] + return centered - (centered @ comps.T) @ comps + + +def finalize_year( + year: int, + sum_: np.ndarray, + sumsq: np.ndarray, + count: np.ndarray, + out_dir: Path, + *, + min_count: int = CONTEXTUAL_MIN_COUNT, + centering: bool = CONTEXTUAL_CENTERING, + pca_remove_k: int = CONTEXTUAL_PCA_REMOVE_K, +) -> dict: + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + sum_ = sum_.astype(np.float64) + sumsq = sumsq.astype(np.float64) + count = count.astype(np.float64) + + has = count > 0 + safe_count = np.where(has, count, 1.0)[:, None] + centroid = (sum_ / safe_count).astype(np.float32) + centroid[~has] = 0.0 + + var_per_dim = sumsq / safe_count - (sum_ / safe_count) ** 2 + dispersion = np.clip(var_per_dim.sum(axis=1), 0.0, None).astype(np.float32) + dispersion[~has] = 0.0 + + # Per-year mean over trusted words only (stable anisotropy estimate). + trusted = count >= min_count + if centering and trusted.any(): + mean_vec = centroid[trusted].mean(axis=0, keepdims=True) + centered = (centroid - mean_vec).astype(np.float32) + centered[~has] = 0.0 + if pca_remove_k > 0: + centered = _remove_top_k(centered, trusted, pca_remove_k).astype(np.float32) + centered[~has] = 0.0 + else: + centered = centroid.copy() + + np.save(out_dir / f"{year}.npy", centroid) + np.save(out_dir / f"{year}_centered.npy", centered) + np.save(out_dir / f"{year}_count.npy", count.astype(np.float32)) + np.save(out_dir / f"{year}_dispersion.npy", dispersion) + + stats = { + "year": year, + "words_with_obs": int(has.sum()), + "words_trusted": int(trusted.sum()), + "total_obs": int(count.sum()), + "mean_dispersion": float(dispersion[trusted].mean()) if trusted.any() else 0.0, + } + print( + f" [finalize {year}] {stats['words_with_obs']:,} words w/ obs, " + f"{stats['words_trusted']:,} trusted (>= {min_count}), " + f"{stats['total_obs']:,} total obs -> {out_dir}", + flush=True, + ) + return stats diff --git a/pipeline/contextual_stream.py b/pipeline/contextual_stream.py new file mode 100644 index 000000000..d923af4e5 --- /dev/null +++ b/pipeline/contextual_stream.py @@ -0,0 +1,167 @@ +"""Re-stream FineWeb and drive the contextual accumulator, per year. + +Reuses the streaming + per-snapshot checkpoint contract of +`pipeline/data_pipeline.py`: each year is a set of `CC-MAIN-YYYY-WW` snapshots, +filtered to `language_score >= LANGUAGE_SCORE_THRESHOLD`, with the per-year +token target split evenly across snapshots. Unlike the Word2Vec path we use the +**raw** `row["text"]` (the encoder's own tokenizer handles it), slice it into +non-overlapping word windows, and aggregate on the fly -- nothing is written to +`data/tokens`. + +Resumable: after each completed snapshot we atomically save the accumulator npz +and a checkpoint json. A relaunch loads the partial state and skips completed +snapshots, but only if the model/config hash matches (never mix configs). +""" +from __future__ import annotations + +import json +import time +from pathlib import Path + +from datasets import load_dataset +from tqdm import tqdm + +from config import ( + CONTEXTUAL_BATCH_SIZE, + CONTEXTUAL_MAX_LEN, + CONTEXTUAL_TARGET_TOKENS_PER_YEAR, + LANGUAGE_SCORE_THRESHOLD, +) +from pipeline.contextual_aggregate import ContextualAccumulator +from pipeline.contextual_model import ContextualEncoder +from pipeline.snapshot_registry import get_snapshots + + +def text_to_windows(text: str, max_len: int = CONTEXTUAL_MAX_LEN) -> list[list[str]]: + """Lowercase, whitespace-split, slice into non-overlapping word windows.""" + words = text.lower().split() + return [words[i : i + max_len] for i in range(0, len(words), max_len) if words[i : i + max_len]] + + +def iter_snapshot_windows(snapshot: str, max_words: int, max_len: int): + """Yield word-windows from one snapshot, stopping after ~max_words words.""" + ds = load_dataset("HuggingFaceFW/fineweb", name=snapshot, streaming=True, split="train") + ds = ds.filter(lambda x: x["language_score"] >= LANGUAGE_SCORE_THRESHOLD) + emitted = 0 + for row in ds: + for win in text_to_windows(row["text"], max_len): + yield win + emitted += len(win) + if emitted >= max_words: + return + + +def run_year( + year: int, + encoder: ContextualEncoder, + accumulator: ContextualAccumulator, + *, + state_dir: Path, + config_hash: str, + target_tokens: int = CONTEXTUAL_TARGET_TOKENS_PER_YEAR, + batch_size: int = CONTEXTUAL_BATCH_SIZE, + max_len: int = CONTEXTUAL_MAX_LEN, + snapshots: list[str] | None = None, + log_every_batches: int = 200, +) -> tuple[int, int]: + """Stream a year into the accumulator. Returns (total_words, total_windows). + + Resumes from any partial state in `state_dir`. Checkpoints after each + completed snapshot. + """ + state_dir = Path(state_dir) + state_dir.mkdir(parents=True, exist_ok=True) + snapshots = snapshots or get_snapshots(year) + tokens_per_snapshot = target_tokens // len(snapshots) + + partial_path = state_dir / f"{year}.partial.npz" + ckpt_path = state_dir / f"{year}_checkpoint.json" + + completed: list[str] = [] + total_words = 0 + total_windows = 0 + + if ckpt_path.exists(): + with open(ckpt_path) as f: + ckpt = json.load(f) + if ckpt.get("config_hash") != config_hash: + raise RuntimeError( + f"Refusing to resume year {year}: checkpoint config_hash " + f"{ckpt.get('config_hash')!r} != current {config_hash!r}. " + f"Delete {ckpt_path} and {partial_path} to restart this year." + ) + completed = ckpt["completed_snapshots"] + total_words = ckpt["total_tokens"] + total_windows = ckpt["total_windows"] + if partial_path.exists(): + accumulator.load_state(partial_path) + print( + f" [year {year}] resuming: {len(completed)}/{len(snapshots)} snapshots, " + f"{total_words:,} words", + flush=True, + ) + + def write_checkpoint() -> None: + tmp = ckpt_path.with_suffix(".json.tmp") + with open(tmp, "w") as f: + json.dump( + { + "completed_snapshots": completed, + "total_tokens": total_words, + "total_windows": total_windows, + "model": encoder.model_name, + "hidden_dim": encoder.hidden_dim, + "pooling": "first", + "max_len": max_len, + "config_hash": config_hash, + "updated_at": time.strftime("%Y-%m-%dT%H:%M:%S"), + }, + f, + ) + tmp.replace(ckpt_path) + + for snapshot in snapshots: + if snapshot in completed: + continue + print(f" [year {year}] {snapshot} (target {tokens_per_snapshot:,} words)...", flush=True) + + snapshot_words = 0 + n_batches = 0 + t0 = time.time() + batch: list[list[str]] = [] + + def flush() -> int: + hidden, _attn, word_ids = encoder.encode_windows(batch) + accumulator.add_batch(hidden, word_ids, batch) + return len(batch) + + for win in iter_snapshot_windows(snapshot, tokens_per_snapshot, max_len): + batch.append(win) + snapshot_words += len(win) + if len(batch) >= batch_size: + total_windows += flush() + n_batches += 1 + batch = [] + if n_batches % log_every_batches == 0: + wps = snapshot_words / max(1e-9, time.time() - t0) + print( + f" {snapshot_words:,}/{tokens_per_snapshot:,} words " + f"({wps:,.0f} words/s)", + flush=True, + ) + if snapshot_words >= tokens_per_snapshot: + break + if batch: + total_windows += flush() + + total_words += snapshot_words + completed.append(snapshot) + accumulator.save_state(partial_path) + write_checkpoint() + print( + f" [year {year}] {snapshot} done: {snapshot_words:,} words, " + f"cumulative {total_words:,} ({time.time()-t0:.0f}s)", + flush=True, + ) + + return total_words, total_windows diff --git a/scripts/contextual_drift.py b/scripts/contextual_drift.py new file mode 100644 index 000000000..995a83075 --- /dev/null +++ b/scripts/contextual_drift.py @@ -0,0 +1,148 @@ +"""Contextual word-drift pipeline entry point. + +Stages (mirrors the TWEC entry point's structure): + stream re-stream FineWeb per year -> accumulate contextual centroids -> + finalize that year (writes models/contextual/{year}*.npy). Resumable. + finalize re-derive {year}*.npy from existing partial accumulators (cheap way + to re-tune centering / PCA-removal without re-streaming). + all stream (which finalizes each year on completion). + +The drift stage (parquet + summary + eval) lands in a follow-up PR alongside +analysis/contextual_drift.py. + + uv run python scripts/contextual_drift.py --smoke # ~2 min code-path + resume check + uv run python scripts/contextual_drift.py --stage stream # full run (background, ~1.7 days) + uv run python scripts/contextual_drift.py --stage finalize # re-finalize from partials +""" +import argparse +import hashlib +import json +import sys +from pathlib import Path + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from config import ( + CONTEXTUAL_DIR, + CONTEXTUAL_LAYER, + CONTEXTUAL_MODEL, + CONTEXTUAL_POOLING, + CONTEXTUAL_MAX_LEN, + CONTEXTUAL_STATE_DIR, + CONTEXTUAL_TARGET_TOKENS_PER_YEAR, + CONTEXTUAL_USE_STOPWORDS, + MIN_WORD_LENGTH, + VOCAB_DIR, + YEARS, +) +from pipeline.contextual_aggregate import ContextualAccumulator +from pipeline.contextual_finalize import finalize_year, load_partial +from pipeline.contextual_model import ContextualEncoder +from pipeline.contextual_stream import run_year +from pipeline.snapshot_registry import get_snapshots +from pipeline.vocab import load_vocab + + +def compute_config_hash(model_name: str) -> str: + """Hash the knobs that affect accumulation; resume refuses on mismatch.""" + payload = json.dumps( + { + "model": model_name, + "layer": CONTEXTUAL_LAYER, + "pooling": CONTEXTUAL_POOLING, + "max_len": CONTEXTUAL_MAX_LEN, + "min_word_length": MIN_WORD_LENGTH, + "use_stopwords": CONTEXTUAL_USE_STOPWORDS, + }, + sort_keys=True, + ) + return hashlib.sha256(payload.encode()).hexdigest()[:16] + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--smoke", action="store_true") + ap.add_argument("--device", default="cuda") + ap.add_argument("--model", default=CONTEXTUAL_MODEL) + ap.add_argument("--years", type=int, nargs="+", default=None) + ap.add_argument("--stage", choices=["stream", "finalize", "all"], default="all") + ap.add_argument("--target-tokens", type=int, default=None) + a = ap.parse_args() + + if a.smoke: + years = a.years or [2014, 2018] + out_dir = Path("models/contextual_smoke") + state_dir = Path("data/contextual_state_smoke") + target_tokens = a.target_tokens or 600_000 # ~300k words/snapshot, fast code-path check + + def snapshots_for(y: int) -> list[str]: + return get_snapshots(y)[:2] # 2 snapshots/year (exercises checkpoint/resume) + else: + years = a.years or list(YEARS) + out_dir = CONTEXTUAL_DIR + state_dir = CONTEXTUAL_STATE_DIR + target_tokens = a.target_tokens or CONTEXTUAL_TARGET_TOKENS_PER_YEAR + + def snapshots_for(y: int) -> list[str] | None: + return None # all registered snapshots + + out_dir.mkdir(parents=True, exist_ok=True) + state_dir.mkdir(parents=True, exist_ok=True) + + vocab = load_vocab(VOCAB_DIR / "vocab.json") + config_hash = compute_config_hash(a.model) + print( + f"contextual-drift | model={a.model} stage={a.stage} years={years} " + f"vocab={len(vocab):,} target={target_tokens:,}/yr -> {out_dir}", + flush=True, + ) + + encoder = None + if a.stage in ("stream", "all"): + encoder = ContextualEncoder(model_name=a.model, device=a.device) + print(f" encoder hidden_dim={encoder.hidden_dim}", flush=True) + + for year in years: + final_path = out_dir / f"{year}.npy" + + if a.stage == "finalize": + s, sq, c = load_partial(state_dir, year) + finalize_year(year, s, sq, c, out_dir) + continue + + # stream / all + if final_path.exists(): + print(f"[year {year}] already finalized ({final_path}), skip", flush=True) + continue + + acc = ContextualAccumulator(vocab, encoder.hidden_dim, device=a.device) + total_words, total_windows = run_year( + year, + encoder, + acc, + state_dir=state_dir, + config_hash=config_hash, + target_tokens=target_tokens, + snapshots=snapshots_for(year), + ) + s, sq, c = acc.to_numpy() + finalize_year(year, s, sq, c, out_dir) + + # Year complete: drop the resume checkpoint, keep the partial npz so + # `--stage finalize` can re-tune centering/PCA without re-streaming. + ckpt = state_dir / f"{year}_checkpoint.json" + if ckpt.exists(): + ckpt.unlink() + print(f"[year {year}] done: {total_words:,} words, {total_windows:,} windows", flush=True) + + del acc + if a.device == "cuda": + torch.cuda.empty_cache() + + print("DONE", flush=True) + + +if __name__ == "__main__": + main()