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

[scheuclu]

Contextual word-drift pipeline — Part 2/3: accumulation + streaming + finalize + orchestrator

Second slice of the contextual word-drift pipeline (#34). Builds the actual data path: re-stream FineWeb → frozen-encoder forward → on-the-fly per-word contextual centroids → per-year finalized embeddings. Never stores per-token vectors (storage collapses to ~9 GB of partials + ~9 GB of artifacts).

Stacked on #35 (Part 1). Review/merge that first; this PR's base is the Part 1 branch.

What's here

File	What
pipeline/contextual_aggregate.py	ContextualAccumulator: preallocated GPU sum[V,H], sumsq[V,H], count[V] (fp32). add_batch() finds first-subword positions, resolves each word → shared vocab id (skipping stopwords / OOV / short / digits), and folds the batch in with three vectorized index_add_s. Atomic save_state/load_state (npz).
pipeline/contextual_stream.py	text_to_windows + iter_snapshot_windows (re-stream FineWeb, language_score >= 0.65, raw text, non-overlapping 128-word windows) and run_year() which batches windows → encoder → accumulator, checkpointing after each completed snapshot (atomic npz + json, config-hash guarded).
pipeline/contextual_finalize.py	finalize_year(): centroid = sum/count, dispersion = Σ(sumsq/count − centroid²) (polysemy signal), per-year mean-centering (BERT anisotropy fix) + optional all-but-top-k PCA removal. Writes {year}.npy, _centered.npy, _count.npy, _dispersion.npy.
scripts/contextual_drift.py	Orchestrator: --stage {stream,finalize,all}, --smoke, --years, --model, --target-tokens. Resumable (skips finalized {year}.npy, resumes partials). On year completion finalizes + drops the checkpoint but keeps the partial so --stage finalize can re-tune centering/PCA without re-streaming.

The drift stage (parquet + summary) + local eval land in Part 3 alongside analysis/contextual_drift.py.

Design notes

• First-subword pooling via word_ids() (vectorized; "mean" is a deferred knob).
• fp32 accumulation even though the forward runs bf16.
• No alignment — a frozen shared encoder already puts every year in one coordinate system.
• Resume safety: a config_hash over {model, layer, pooling, max_len, min_word_length, use_stopwords} is stored in the checkpoint; resume refuses on mismatch rather than mixing incompatible accumulations.

Verification (smoke)

uv run python scripts/contextual_drift.py --smoke --years 2018 (2 snapshots, 600k words):

  [year 2018] CC-MAIN-2018-05 done: 300,031 words, cumulative 300,031 (41s)
  [year 2018] CC-MAIN-2018-09 done: 300,039 words, cumulative 600,070 (42s)
  [finalize 2018] 26,820 words w/ obs, 992 trusted (>= 50), 251,022 total obs
[year 2018] done: 600,070 words, 5,272 windows
DONE

Artifact checks (models/contextual_smoke/2018*):

shapes: (119466, 768) (119466, 768) (119466,) (119466,)   dtype: float32
words with obs: 26820   words >=50 obs: 992
  house    count= 111  disp=101.6   the   count=0  (stopword skipped)
  bank     count=  40  disp= 88.6   covid OOV (not in shared vocab)

anisotropy (mean random-pair cosine over 992 trusted words):
  raw      : 0.4703   (BERT anisotropy, as expected)
  centered : -0.0004  (~0 after mean-centering -> fix works)

Resume check (kill after snapshot 1, relaunch):

# killed after snapshot 1, then relaunched:
  [year 2014] resuming: 1/2 snapshots, 300,002 words
  [year 2014] CC-MAIN-2014-15 (target 300,000 words)...        # snapshot 1 skipped
  [year 2014] CC-MAIN-2014-15 done: 300,030 words, cumulative 600,032 (51s)
  [finalize 2014] 28,102 words w/ obs, 940 trusted, 247,467 total obs
DONE
# post-completion: 2014_checkpoint.json deleted, 2014.partial.npz (0.73 GB) kept
# --stage finalize re-derived identical numbers from the partial (no re-stream)

🤖 Generated with Claude Code

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
language-drift	Ready	Preview, Comment	Jun 6, 2026 6:24pm