feat(magic): support per-token per-query MAGIC

[luciaquirke]

Opus slop, ignore:

Stacked on #412 (base is fix/magic-doc-ids-fresh-runs). That PR makes doc_ids.pt conditional on per_token alone, which is what lets it cover the 3-D case here.

Note: this branch is still named fix/magic-reject-per-token-per-query from an earlier revision of this PR that rejected the combination rather than supporting it. GitHub can't rename a PR's head branch — say the word and I'll re-open under feat/per-token-per-query.

Problem

attribute_tokens=True with query_method="none" produced an unusable score tensor.

Per-token weights are [rows, seq_len] and the per-query stack used dim=1, giving [rows, num_queries, seq_len] — query axis in the middle, which nothing downstream reads. validate_scores takes shape[-1] as the query count, so it compared seq_len against the query document count and died naming the wrong dimension:

ValueError: scores has 8 query columns but the query dataset has 2 documents;
multi-query validation requires one score column per query document

on a run with 2 queries and seq_len 8.

Fix

Stack on dim=-1. The query axis then comes last in both modes — [rows, num_queries] per-doc, [rows, seq_len, num_queries] per-token — matching the layout Scores.to_grid already produces for multi-query token score directories, which load_attribution_scores already flags multi_query.

validate_scores needed no change: shape[-1] is the query count and reshape(-1, num_queries) flattens the leading axes into leave-out units — documents in 2-D, token positions in 3-D. tests/test_per_token_lds.py already asserts that expression against a [docs, seq_len, queries] grid. dim=-1 is identical to dim=1 for 1-D inputs, so per-doc per-query scores are unchanged.

Fix the padding trim in the same path, which applied weight_pad_count regardless of rank while the main scoring path picks by rank. The two differ once doc_ids are present (pad rows route to one synthetic doc id):

5 docs @ batch_size 4 -> rows=8 pad_count=3 weight_pad_count=1

so the run kept 7 of its 8 padded rows instead of trimming to 5, leaving pad rows in the saved scores.

Teach both .pt classifiers about 3-D — scores_are_per_token so a reloaded run sizes its weights per-token, and _pt_scores_are_per_query so it is recognised as multi-query. 3-D needs no config lookup to disambiguate, unlike 2-D.

Second commit: format from the config, not the shape

9527dbc is a separate commit, revertable on its own.

Layout was inferred from tensor rank in several places, which is guesswork — a 2-D .pt is [docs, seq_len] or [docs, queries] depending only on how the run was configured. #407 established the fix for one call site (read the config.yaml beside scores.pt), but scores_are_per_token was left sniffing shapes and the parsing lived inline rather than beside the other config readers.

Adds read_first_step_config to config_io, next to read_config and load_subconfig, and routes both classifiers through it. The flags say everything the rank could:

query_method	attribute_tokens	layout
none	yes	[docs, seq_len, queries]
none	no	[docs, queries]
mean/sum	yes	[docs, seq_len]
mean/sum	no	[docs]

So a run is per-query iff query_method is none, whatever rank results, and per-token iff attribute_tokens is set. Score directories keep reading info.json, and load_attribution_scores now takes the query count from the store's num_scores instead of re-deriving it from the grid it just built.

Shape survives in exactly one place: a .pt with no config beside it, where nothing else is knowable and only rank 3 is unambiguous.

cfg_attributes_tokens reads attribute_tokens with the deprecated per_token as an alias, in one place, so a run written with the current field name is no longer missed.

Test cleanup

Dropped six tests that asserted torch's own view()/reshape() indexing semantics. They called no bergson code, so they could not fail unless PyTorch itself changed; the behaviour they stood in for is covered end to end by the per-query aggregation test.

One existing test changed

test_load_attribution_scores_pt_per_query asserted that a 2-D tensor whose config said query_method: none and per_token: true was single-query. No run produces that pair — attributing tokens per query yields rank 3 — so the case was describing an unreachable artifact and pinning the shape-derived answer for it. Repointed at the 3-D tensor such a run does produce.

Net change in bergson/: +19 / −17, four of them functional.

Third commit: one writer for the token score format

766a9b7 is also revertable on its own.

save_sequence_scores delegates to MemmapSequenceScoreWriter, but its token counterpart reimplemented MemmapTokenScoreWriter inline — the memmap creation, offsets.npy, and an info.json payload identical field for field. The on-disk token score format was written from two places that had to be kept in step by hand.

That matters more after the commit above: scores_are_per_token now reads info.json["attribute_tokens"] as authoritative, so drift between the two writers stops being cosmetic and becomes a misclassification.

Delegating needs the writer to accept what it actually uses. It only ever took a Dataset in order to call compute_num_token_grads on it, while save_token_scores already holds the offsets those counts came from — so __init__ now takes num_token_grads, and a from_dataset classmethod covers the four callers that hold a dataset.

It also gives the token writer the overwrite flag its sequence twin already had, which delegation requires: save_token_scores wrote with mode="w+" unconditionally, and without overwrite the writer would silently reuse a stale scores.bin rather than replacing it.

Net 19 lines out of score_writer.py, and one place left that knows the format.

Fourth commit: drop .npy score support

12f0d1f, also revertable on its own.

bergson never writes a .npy score file — every writer emits a score directory — so .npy was an ingest-only path for arrays produced outside the pipeline, and nothing in the repo feeds one: no config sets scores: to a .npy, and the examples that save scores.npy read it straight back with np.load rather than through load_attribution_scores.

Removes the branch from load_attribution_scores, scores_are_per_token and worker's score-path dispatch, and with it ArrayScores, which existed only to give a bare array the Scores interface.

It also carried its own rules. A .npy could not have a score_cfg, so it alone skipped the higher_is_better negation and had to arrive in the loss-diff convention already; and it was the one input whose multi-query flag came from a raw column count. Score directories record num_scores in info.json, so the surviving formats all describe themselves.

The bank-loss-cache tests used .npy to hand a score matrix to evaluate_retrained; they now write a score directory via save_sequence_scores, and the multi_query parametrization still passes both ways.

Net −49 lines.

Verification

End-to-end on CPU, 5 docs / 2 queries / seq_len 8, attribute_tokens=True, query_method="none":

scores.pt shape=(5, 8, 2)          # was (7, 2, 8)
doc_ids.pt written: True
doc_ids shape=(5, 8) indexes scores' leading axes: True
per-(doc,query) agg shape=(5, 2) sum preserved: True

and the leave-subset-out validation that previously raised now completes:

Query 0: Spearman -0.5000 (p=6.67e-01)  Pearson -0.2178 (p=8.60e-01)
Query 1: Spearman  1.0000 (p=0.00e+00)  Pearson  0.7076 (p=5.00e-01)
Mean Spearman across 2 queries: 0.2500

(Correlations are meaningless at 3 subsets on a randomly-initialised model — the point is the machinery runs.)

Testing

• test_per_query_per_token_aggregates_to_per_doc is the numerical gate: per-token per-query scores summed over each document's tokens reproduce the per-doc per-query run to 1e-5.
• Both new end-to-end tests fail on the parent commit with expected (5, 8, 2), got (7, 2, 8) — wrong axis order and untrimmed padding together.
• test_three_dim_scores_load_as_per_token_multi_query covers the reload path.
• 109 passed, 29 skipped across every test file touching magic / validate / scores, including the existing per-query and multi-query suites — so per-doc per-query is unaffected.
• pre-commit clean.

🤖 Generated with Claude Code