feat(contextual): pipeline part 1 — config, frozen encoder, benchmark

config.py

@@ -28,3 +28,40 @@
 ALIGNED_DIR = MODELS_DIR / "aligned"
 DRIFT_DIR = MODELS_DIR / "drift"
 TENSORBOARD_DIR = Path("runs")
+
+# --- Contextual drift pipeline ---
+# Upgrade from static Word2Vec/TWEC drift to *contextual* drift: run the
+# historical corpus through a frozen transformer encoder and aggregate per-word
+# contextual centroids per year. A frozen shared encoder puts every year in one
+# coordinate system, so no Procrustes/TWEC compass is needed.
+CONTEXTUAL_MODEL = "bert-base-uncased"  # fast WordPiece tokenizer, lowercase, SDPA
+CONTEXTUAL_LAYER = -1  # which hidden state to pool (-1 = last_hidden_state)
+CONTEXTUAL_POOLING = "first"  # "first" subword | "mean" over subwords (mean deferred)
+CONTEXTUAL_MAX_LEN = 128  # non-overlapping window length (words, pre-tokenization)
+CONTEXTUAL_BATCH_SIZE = 128  # windows per forward; benchmark optimum on GB10 (~80k words/s)
+CONTEXTUAL_TARGET_TOKENS_PER_YEAR = TARGET_TOKENS_PER_YEAR  # 1B (full coverage)
+CONTEXTUAL_MIN_COUNT = 50  # min observations for a trusted (word, year) centroid
+CONTEXTUAL_CENTERING = True  # per-year mean-centering (BERT anisotropy fix)
+CONTEXTUAL_PCA_REMOVE_K = 0  # all-but-top-k PCA component removal; 0 = off
+CONTEXTUAL_USE_STOPWORDS = True  # skip function words when accumulating
+CONTEXTUAL_DIR = MODELS_DIR / "contextual"
+CONTEXTUAL_STATE_DIR = DATA_DIR / "contextual_state"
+
+# Small English function-word list. These are extremely high frequency, carry
+# little drift signal, and would dominate the accumulator; we skip them when
+# building the target-id tensor (gated by CONTEXTUAL_USE_STOPWORDS).
+CONTEXTUAL_STOPWORDS = frozenset({
+    "a", "an", "the", "and", "or", "but", "if", "then", "else", "when",
+    "of", "to", "in", "on", "at", "by", "for", "with", "about", "into",
+    "from", "up", "down", "out", "off", "over", "under", "as", "is", "am",
+    "are", "was", "were", "be", "been", "being", "have", "has", "had",
+    "do", "does", "did", "doing", "will", "would", "shall", "should",
+    "can", "could", "may", "might", "must", "not", "no", "nor", "so",
+    "than", "too", "very", "just", "this", "that", "these", "those",
+    "i", "you", "he", "she", "it", "we", "they", "me", "him", "her",
+    "us", "them", "my", "your", "his", "its", "our", "their", "mine",
+    "yours", "hers", "ours", "theirs", "who", "whom", "whose", "which",
+    "what", "where", "why", "how", "all", "any", "both", "each", "few",
+    "more", "most", "other", "some", "such", "only", "own", "same",
+    "here", "there", "again", "once", "also", "now", "ever", "never",
+})

pipeline/contextual_model.py

@@ -0,0 +1,100 @@
+"""Frozen transformer encoder for the contextual-drift pipeline.
+
+Loads a frozen `bert-base-uncased` (or any `AutoModel`) with its *fast*
+tokenizer and exposes `encode_windows`, which runs a bf16 `inference_mode`
+forward over a batch of pre-tokenized word windows and returns the chosen
+hidden-state layer together with the attention mask and a `word_ids` tensor
+that maps every subword position back to its source word index.
+
+Why a frozen encoder: it puts every year in one coordinate system directly, so
+no Procrustes/TWEC alignment is needed downstream. Why the fast tokenizer's
+`word_ids()`: it maps subwords -> words correctly for WordPiece/BPE/SP, unlike
+the brittle `"##"`-stripping heuristic.
+
+The hidden width is read from `model.config.hidden_size` -- never hardcode 768.
+"""
+from __future__ import annotations
+
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+from config import CONTEXTUAL_LAYER, CONTEXTUAL_MODEL
+
+
+class ContextualEncoder:
+    def __init__(
+        self,
+        model_name: str = CONTEXTUAL_MODEL,
+        device: str = "cuda",
+        layer: int = CONTEXTUAL_LAYER,
+        dtype: torch.dtype = torch.bfloat16,
+    ) -> None:
+        self.model_name = model_name
+        self.device = device
+        self.layer = layer
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+        if not self.tokenizer.is_fast:
+            raise RuntimeError(
+                f"{model_name} has no fast tokenizer; word_ids() alignment requires one."
+            )
+
+        self.model = AutoModel.from_pretrained(
+            model_name,
+            dtype=dtype,
+            attn_implementation="sdpa",
+        )
+        self.model.eval().requires_grad_(False).to(device)
+
+        self.hidden_dim: int = self.model.config.hidden_size
+        # Truncate to the model's positional limit; 128-word windows rarely exceed it.
+        self.max_subwords: int = min(
+            getattr(self.model.config, "max_position_embeddings", 512), 512
+        )
+
+    @torch.inference_mode()
+    def encode_windows(
+        self, windows: list[list[str]]
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Encode a batch of word windows.
+
+        Args:
+            windows: list of windows, each a list of lowercase word strings.
+
+        Returns:
+            hidden:     [B, L, H] hidden states of the configured layer (model dtype).
+            attn_mask:  [B, L] attention mask (1 = real token).
+            word_ids:   [B, L] long tensor; entry j is the source word index of
+                        subword j, or -1 for special/padding tokens. On `device`.
+        """
+        enc = self.tokenizer(
+            windows,
+            is_split_into_words=True,
+            padding=True,
+            truncation=True,
+            max_length=self.max_subwords,
+            return_tensors="pt",
+        )
+
+        # word_ids must be read from the BatchEncoding before moving tensors.
+        word_id_rows = [
+            [-1 if w is None else w for w in enc.word_ids(i)]
+            for i in range(len(windows))
+        ]
+        word_ids = torch.tensor(word_id_rows, dtype=torch.long, device=self.device)
+
+        input_ids = enc["input_ids"].to(self.device)
+        attn_mask = enc["attention_mask"].to(self.device)
+        token_type = enc.get("token_type_ids")
+        kwargs = {"input_ids": input_ids, "attention_mask": attn_mask}
+        if token_type is not None:
+            kwargs["token_type_ids"] = token_type.to(self.device)
+
+        if self.layer == -1:
+            out = self.model(**kwargs)
+            hidden = out.last_hidden_state
+        else:
+            out = self.model(**kwargs, output_hidden_states=True)
+            hidden = out.hidden_states[self.layer]
+
+        return hidden, attn_mask, word_ids

pyproject.toml

@@ -16,5 +16,6 @@ dependencies = [
     "tensorboard>=2.20.0",
     "torch>=2.12.0",
     "tqdm>=4.67.3",
+    "transformers>=5.10.0",
     "umap-learn>=0.5.12",
 ]

scripts/contextual_benchmark.py

@@ -0,0 +1,125 @@
+"""Throughput probe for the contextual-drift pipeline. RUN THIS FIRST.
+
+Pulls a few thousand real word-windows from one 2018 FineWeb snapshot (writes
+nothing to disk), warms up the GPU, then sweeps the forward-pass batch size and
+prints windows/sec, words/sec, subwords/sec and peak GPU memory, plus an
+hours-per-year and total-wall-clock table for the full 12 B-token run. Use the
+result to set CONTEXTUAL_BATCH_SIZE in config.py and confirm the ETA before
+launching the real run.
+
+    uv run python scripts/contextual_benchmark.py
+    uv run python scripts/contextual_benchmark.py --snapshot CC-MAIN-2018-30 --windows 3000
+"""
+import argparse
+import sys
+import time
+from pathlib import Path
+
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from config import (
+    CONTEXTUAL_MAX_LEN,
+    CONTEXTUAL_TARGET_TOKENS_PER_YEAR,
+    LANGUAGE_SCORE_THRESHOLD,
+    YEARS,
+)
+from pipeline.contextual_model import ContextualEncoder
+from pipeline.snapshot_registry import get_snapshots
+
+
+def collect_windows(snapshot: str, n_windows: int, max_len: int) -> list[list[str]]:
+    """Stream a snapshot and slice raw text into non-overlapping word windows.
+
+    Mirrors the windowing the real pipeline uses (lowercase + whitespace split,
+    fixed-length non-overlapping windows); kept inline so the probe has no
+    dependency on the not-yet-built streaming module.
+    """
+    from datasets import load_dataset
+
+    ds = load_dataset("HuggingFaceFW/fineweb", name=snapshot, streaming=True, split="train")
+    ds = ds.filter(lambda x: x["language_score"] >= LANGUAGE_SCORE_THRESHOLD)
+
+    windows: list[list[str]] = []
+    for row in ds:
+        words = row["text"].lower().split()
+        for i in range(0, len(words), max_len):
+            win = words[i : i + max_len]
+            if win:
+                windows.append(win)
+                if len(windows) >= n_windows:
+                    return windows
+    return windows
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--device", default="cuda")
+    ap.add_argument("--snapshot", default=None, help="default: a 2018 snapshot")
+    ap.add_argument("--windows", type=int, default=2000)
+    ap.add_argument("--batch-sizes", type=int, nargs="+", default=[128, 256, 512])
+    a = ap.parse_args()
+
+    snapshot = a.snapshot or get_snapshots(2018)[6]  # CC-MAIN-2018-30, mid-year
+    print(f"Loading encoder on {a.device}...", flush=True)
+    enc = ContextualEncoder(device=a.device)
+    print(f"  model={enc.model_name} hidden_dim={enc.hidden_dim} max_subwords={enc.max_subwords}", flush=True)
+
+    print(f"Pulling {a.windows} windows from {snapshot} (max_len={CONTEXTUAL_MAX_LEN})...", flush=True)
+    t0 = time.time()
+    windows = collect_windows(snapshot, a.windows, CONTEXTUAL_MAX_LEN)
+    avg_words = sum(len(w) for w in windows) / max(1, len(windows))
+    print(f"  got {len(windows)} windows, avg {avg_words:.1f} words/window ({time.time()-t0:.0f}s)", flush=True)
+
+    # Warm up (kernels, autotune, allocator).
+    print("Warming up...", flush=True)
+    for _ in range(3):
+        enc.encode_windows(windows[:a.batch_sizes[0]])
+    torch.cuda.synchronize()
+
+    print(f"\n{'batch':>6} {'win/s':>9} {'words/s':>11} {'subwd/s':>11} {'peakGB':>8}", flush=True)
+    print("-" * 50, flush=True)
+
+    results = []
+    for bs in a.batch_sizes:
+        torch.cuda.reset_peak_memory_stats()
+        torch.cuda.synchronize()
+        t0 = time.time()
+        total_windows = 0
+        total_subwords = 0
+        for i in range(0, len(windows), bs):
+            batch = windows[i : i + bs]
+            _, attn_mask, _ = enc.encode_windows(batch)
+            total_windows += len(batch)
+            total_subwords += int(attn_mask.sum().item())
+        torch.cuda.synchronize()
+        dt = time.time() - t0
+        win_s = total_windows / dt
+        words_s = win_s * avg_words
+        subwd_s = total_subwords / dt
+        peak_gb = torch.cuda.max_memory_allocated() / 1e9
+        results.append((bs, words_s))
+        print(f"{bs:>6} {win_s:>9.1f} {words_s:>11.0f} {subwd_s:>11.0f} {peak_gb:>8.2f}", flush=True)
+
+    # ETA table from the best (highest words/sec) config.
+    best_bs, best_words_s = max(results, key=lambda r: r[1])
+    hrs_year = CONTEXTUAL_TARGET_TOKENS_PER_YEAR / best_words_s / 3600
+    print(
+        f"\nBest: batch={best_bs} @ {best_words_s:,.0f} words/s",
+        flush=True,
+    )
+    print(
+        f"At {CONTEXTUAL_TARGET_TOKENS_PER_YEAR/1e9:.0f}B words/year: "
+        f"{hrs_year:.1f} h/year, {hrs_year*len(YEARS):.1f} h total "
+        f"({hrs_year*len(YEARS)/24:.1f} days) over {len(YEARS)} years.",
+        flush=True,
+    )
+    print(
+        "Note: words/s drives ETA (corpus target is in words); subwd/s reflects raw GPU load.",
+        flush=True,
+    )
+
+
+if __name__ == "__main__":
+    main()