Token maps: map every late-interaction token as its own point, with token-level SAE features

[enjalot]

What

For ColBERT-style late-interaction embeddings (which already store per-token vectors), this adds a token granularity through the whole product: instead of one point per document (mean-pooled), the map shows one point per token, while the table keeps documents readable by showing the parent text with the selected token highlighted in context. Token-level SAEs plug in as first-class feature surfaces.

Pipeline

ls-embed  <ds> text colbert-jinaai___jina-colbert-v2      # unchanged
ls-tokenize <ds> embedding-001                            # NEW: token strings + char offsets
ls-umap   <ds> embedding-001 25 0.1 --granularity tokens  # one 2D point per token
ls-cluster <ds> umap-001 1000 25 0.0 --method hdbscan     # token-frequency default labels
ls-sae    <ds> embedding-001 <model> 32_512 --granularity tokens --checkpoint <dir>
ls-scope  <ds> embedding-001 umap-001 cluster-001 default "Token map" "..." --sae_id sae-001

Full docs in docs/token-maps.md.

Key design points

• Alignment is exact, or nothing is written. ls-tokenize replays pylate's document pipeline (tokenize to document_length-1, insert the [DocumentMarker] at position 1, drop the 32-symbol punctuation skiplist) with return_offsets_mapping, and validates every document's kept-token count against the stored num_tokens. Verified against real models: jina-colbert-v2 and answerai-colbert-small-v1, including truncation, empty strings, and punctuation-only edge cases.
• Nothing materializes the full token set. Token vectors stream out of LanceDB in bounded batches everywhere; the UMAP fits on ≤ --fit_sample uniformly sampled tokens and batch-transforms the rest (cuML or CPU).
• The frontend invariant survives. Token scope parquets keep ls_index == positional index (it is the global token index), so ScatterGL/selection plumbing is untouched; parent_index/token_pos/token_str columns carry the document linkage, and POST /api/tokens/indexed resolves token indices → parent rows + char spans for the highlight cell.
• SAE surface generalizes. ls-sae gains --granularity tokens and --checkpoint <local dir>; the Explore feature modal/filter/coloring now enable from the scope's own SAE metadata (not just the CDN-labeled allowlist).

Verified end to end (RTX 5090)

4,000 fineweb-edu chunks → jina-colbert-v2 → 1,009,572 token points with the 64K BatchTopK token SAE (enjalot/sae-jina-colbert-v2-tokens-64K):

step	result
ls-embed	25 s (shared GPU)
ls-tokenize	1,009,572/1,009,572 tokens aligned, zero mismatches
token UMAP (cuML)	~90 s (1M-sample fit + batch transform)
hdbscan (cuML)	61 clusters + noise, token-frequency labels
token SAE encode	[1009572, 32] top-k, 52/65,536 dead features
Explore UI	1M points render in ~15 s; snippet highlighting correct incl. mid-word subwords; per-token feature modal works

Known limitations (documented)

• NN text search + column filters are hidden for token scopes (they return document indices; mapping MaxSim's per-token matches onto the map is follow-up).
• Scoped LanceDB export (per-scope vector search) is skipped at token granularity.
• The scope-rows JSON payload at 1M tokens is ~227 MB (browser parses it in ~15 s); a binary transport is future work.
• Token SAEs need a latentsae release with BatchTopK + load_from_disk (currently satisfied via the local latent-sae checkout).

Tests

• uv run pytest tests/ -q → 298 passed (8 new: tokenize alignment/mismatch-abort/streaming, token umap both fit paths, granularity guards, token SAE, /api/tokens/indexed)
• cd web && npm run test → 121 passed (18 new tokenSnippet cases); lint at baseline; production build clean

🤖 Generated with Claude Code

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 793418a620

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Reject mismatched token SAEs

For token scopes this only checks granularity, so ls-scope ... --sae_id sae-001 still succeeds when that token SAE was produced for a different embedding, or for a stale token run with fewer rows. The UI and /api/tokens/indexed then index the SAE HDF5 by the current token indices, which can silently show activations for the wrong token vectors or 500 on out-of-range indices. Please also validate the SAE metadata's embedding_id and rows against the current embedding/token count before accepting it.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Guard URL column filters for token scopes

Hiding column filters from the dropdown is not enough because FilterContext still hydrates ?column=...&value=... and runs the COLUMN filter path, whose result is a list of document row indices. In a token scope those values are then treated as token indices and fetched through /api/tokens/indexed, so a shared or stale column-filter URL selects arbitrary tokens instead of matching parent documents. Add the same isTokenScope guard to URL hydration or the COLUMN switch that the SEARCH path now has.

Useful? React with 👍 / 👎.

[enjalot]

Round 2 (pushed): merged latest main + UX fixes from testing

Merged origin/main (Amber Console restyle, 3D explore views, FilterDataTable body/connected/standalone refactor) — token mode re-grafted onto the new structure: token-scope awareness now flows as props into the presentational table body, and the token hover lives in the new floating hover-card Panel. All 303 pytest + 121 vitest pass; lint at baseline.

Binary scope-rows transport — token scopes now fetch the scope parquet bytes directly (existing /api/files route) and parse client-side with hyparquet, selecting only the columns Explore reads. 19.8 MB on the wire instead of 227 MB of JSON; map renders in ~2.7 s instead of ~15 s on the 1M-token demo. Row scopes keep the JSON path.

Token stays visible in the table snippet — three-segment flex cell: before-context ellipsizes from the left (dir="rtl" + <bdi dir="ltr">), after-context from the right, the highlighted token never clips out of view.

Passage context everywhere — the hover card hydrates the parent passage via /api/tokens/indexed (debounced) and highlights the token in place; the point-detail drawer highlights the token's span inside the full document text.

SAE feature labels — new saeLabels map keyed by the SAE's own model repo (scope-declared SAEs aren't tied to an embedding-model entry); ScopeContext loads the latent-taxonomy label parquet, merges per-dataset stats, and falls back to generic labels if the parquet is unpublished. Verified end-to-end with the real COLBERT_JINA_64K features.parquet (65,487 span labels) — labels go live for everyone once the pending latent-taxonomy site package deploys.

🤖 Generated with Claude Code

[enjalot]

SAE labels now working end-to-end from the deployed latent-taxonomy site (a18b647).

After the COLBERT_JINA_64K package deployed, labels still failed in the browser — root cause was infrastructure, not the wiring: GitHub Pages returns 416 for range requests when the client negotiates gzip (byte offsets from the uncompressed length don't exist in the compressed representation, and browsers always send Accept-Encoding: gzip). hyparquet's asyncBufferFromUrl is range-based, so the label parquet read failed. Reproducible with curl: a plain ranged GET returns the PAR1 footer; the identical request with Accept-Encoding: gzip returns 416.

getFeatures now fetches the whole parquet (a few MB, gzip-friendly single GET) and parses from the buffer via parquetReadObjects on named columns — which also removes the old positional f[0]/f[6] column-order coupling. This likely explains any past flakiness in NOMIC label loading too.

Verified headless against the live CDN: feature modal shows the full span labels + per-dataset stats for the 64K token SAE.

🤖 Generated with Claude Code

[enjalot]

Filter single-select fix (d0a6f76): selecting a feature while a cluster filter is open now replaces the cluster — previously both params stacked in the URL, the restore effect (which reads only the first param) kept showing the cluster chip, and clear wiped both. All filter writers now share applyFilterToUrlParams + clear the other filter hooks' state; also the filter effect re-runs when SAE feature labels finish loading (a feature filter applied early used to silently return empty). Verified live: ?cluster=12 → funnel click in the feature modal → ?feature=19, chip shows the feature, map recolors by activation. 5 new vitest cases.

Also registered the MiniLM Stage-J 49K label parquet in saeLabels.

[enjalot]

Follow-up fix: the feature modal's "Top 15" sliced the raw h5 top-k arrays, but latentsae stores them unsorted (topk(..., sorted=False)) — so the list showed an arbitrary subset and dropped genuinely strong features (a rank-6 activation sat at array position 31, invisible). Now sorted by activation before slicing, zero-act padding dropped, and the latent-taxonomy link resolves the scope's own SAE model instead of the hardcoded NOMIC one. 🤖