fix(magic): score per-query MAGIC by document, not by row

[luciaquirke]

Opus slop

The bug

Per-query MAGIC crashes on a chunked query set (--query.chunk_length 32, i.e. any raw dataset you let bergson chunk):

IndexError: index 2 is out of bounds for dimension 0 with size 2
  bergson/magic/data_stream.py:132   "example_weight": self.weights[indices]

A chunked query set carries a per-token doc_ids column: a row can pack several documents and a document can span several rows, so query i is not row i. compute_per_query_magic_scores selected row qi and sized the per-query stream's weights by row, while DataStream indexes weights by document id whenever the batch has doc_ids — hence the overrun. Unchunked query sets (every shipped example uses chunk_length: 0) have no doc_ids column and were unaffected.

Fix

Score per document, which is the unit everything downstream already uses (validate_scores requires one column per query document; per_doc_query_losses aggregates losses by doc_ids):

• pick the rows holding document qi's tokens,
• mask every other document's tokens out of the labels, so the loss is that document's own mean cross-entropy,
• drop doc_ids from the per-query stream, so its row-sized weights are indexed by row. compute_query_gradients deletes example_weight anyway — the document restriction now lives in the labels.

Padding cycles through the document's own rows instead of appending dead ones: a pad row is a real term in the query loss (the weights that would silence it are discarded), and an all-pad batch on some rank would silently scale the query gradient by 1/world_size. With one row per document this repeats that row — exactly what pad_dataset_to_batch_size did — so unchunked runs are bit-for-bit unchanged.

Chunking drops the tail that doesn't fill a chunk, so a short document (e.g. a blank line in wikitext) can leave no tokens behind. Those score zeros, matching the zero baseline loss per_doc_query_losses gives them; their correlation is undefined, so the reported mean Spearman now averages the queries that have one rather than coming out nan for the whole run: