feat(contextual): pipeline part 2 — accumulator, streaming, finalize, orchestrator

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 <UNK>=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 <UNK>=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)

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

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