feat(magic): per-token per-query MAGIC, and score-format cleanups

[luciaquirke]

Opus slop

Reopens #413 against main. #413 was based on fix/magic-doc-ids-fresh-runs, and when #412 was squash-merged into main the squash did not carry this work — #413 then merged into a base branch that main no longer tracks, so all five commits were stranded. Same commits, rebased onto main, no conflicts.

1. Per-token per-query MAGIC

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

on a run with 2 queries and seq_len 8.

Stacking on dim=-1 puts the query axis last in both modes — [rows, num_queries] per-doc, [rows, seq_len, num_queries] per-token — matching the layout Scores.to_grid already produces. validate_scores needed no change: shape[-1] is the query count and reshape(-1, num_queries) flattens the leading axes into leave-out units either way. dim=-1 is identical to dim=1 for 1-D inputs, so per-doc per-query scores are unchanged.

Also fixes the padding trim in that path, which applied weight_pad_count regardless of rank while the main scoring path picks by rank. With doc_ids present the two differ, so a 5-doc dataset at batch_size 4 kept 7 of its 8 padded rows instead of trimming to 5.

Verified end to end on CPU (5 docs / 2 queries / seq_len 8): scores.pt is (5, 8, 2), was (7, 2, 8); doc_ids (5, 8) indexes its leading axes; per-(doc, query) aggregation preserves the total; and the leave-subset-out validation that previously raised now completes with per-query correlations.

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 an independent per-doc per-query run to 1e-5. It fails on the parent with expected (5, 8, 2), got (7, 2, 8).

2. Format from the config, not the shape

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; 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 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]

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

One existing test changed: it 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 it was pinning a shape-derived answer for an unreachable artifact. Repointed at the 3-D tensor such a run does produce.

Also drops six tests that asserted torch's own view()/reshape() semantics; they called no bergson code and could not fail unless PyTorch changed.

3. One writer for the token score format

save_sequence_scores delegates to its writer class; save_token_scores reimplemented MemmapTokenScoreWriter inline — memmap creation, offsets.npy, and an info.json payload identical field for field. That matters more after commit 2, which makes info.json["attribute_tokens"] authoritative: drift between the writers becomes a misclassification.

The writer only ever took a Dataset to call compute_num_token_grads on it, so __init__ now takes num_token_grads and a from_dataset classmethod covers the callers holding a dataset. It also gains the overwrite flag its sequence twin already had — without it, delegation would silently reuse a stale scores.bin. Net 19 lines out of score_writer.py.

4. Drop .npy score support

bergson never writes a .npy score file, and nothing in the repo feeds one: no config sets scores: to a .npy, and the examples that save scores.npy read it back with np.load rather than through load_attribution_scores. Removes the branch from three call sites and with it ArrayScores.

It also carried its own rules — a .npy could not have a score_cfg, so it alone skipped the higher_is_better negation, and it was the one input whose multi-query flag came from a raw column count. Net −49 lines.

5. Inline the per-token config lookup

cfg_attributes_tokens had one caller left after the earlier inlining, and reaching from magic.cli into validate for a two-key dict lookup bought nothing.

Testing

107 passed, 29 skipped across every test file touching magic, validate or scores, on main. pre-commit clean. Each commit is revertable on its own.

🤖 Generated with Claude Code