[flyingdoubleG]

Summary

SurprisalScorer.compute_loss_with_kv_cache has two problems:

1. The explanation's KV cache is shared across batches. The cache is built once and passed as past_key_values to every batch. On current transformers this is a DynamicCache, which the forward pass extends in place with the batch's tokens. Every batch after the first then attends to the sentences of the earlier batches (and has shifted positions), and its losses are wrong. There is no error.
2. A partial last batch fails. explanation_tokens has already been repeated to batch_size rows, and the partial-batch branch calls repeat_interleave(current_batch_size) on it again, giving batch_size * current_batch_size rows. With len(samples) % batch_size != 0 the forward pass raises RuntimeError: Sizes of tensors must match except in dimension 2.

Repro: a randomly initialised 2-layer GPT-2, batch_size=2, four sentences. Each per-token loss is compared with a full forward pass of explanation + sentence:

sentence	batch	max |loss − reference| (main)	this PR
"The cat sat on the mat"	1	0.0000	0.0000
"Dogs bark loudly"	1	0.0000	0.0000
"A bird sings"	2	0.2677	0.0000
"Rain falls"	2	0.2605	0.0000

Change

Build the explanation cache inside the batch loop, sized to the current batch (explanation_tokens.expand(current_batch_size, -1)). This fixes both problems. The cost is one extra forward pass over the short explanation prompt per batch.

Tests

tests/test_scorers/test_surprisal.py builds a random 2-layer GPT-2 with the tokenizer fixture's vocabulary and uses 5 sentences with batch_size=2 (two full batches and a partial one). It checks every per-token loss against the full forward pass. It fails on main (the RuntimeError above) and passes with this change. ruff check and black --check pass. Locally I ran it with a stand-in tokenizer fixture instead of tests/conftest.py, since vLLM does not install on Windows.