Batch tokenization in the transformers backend's scoring loop

src/lighteval/models/abstract_model.py

@@ -353,5 +353,80 @@ def tok_encode_pair(self, context, continuations: list[str], pairwise: bool = Fa
 
         return context_encs, continuations_encs
 
+    def _batch_tok_encode(self, strings: list[str], add_special_tokens: bool) -> list[list[int]]:
+        """Tokenize a list of strings in a single tokenizer call, without padding.
+
+        Equivalent to `[self.tok_encode(s, add_special_tokens) for s in strings]`,
+        but issues one call instead of one per string, which lets a fast
+        tokenizer batch and parallelize the work instead of paying per-call
+        Python overhead for every string.
+        """
+        if not strings:
+            return []
+        return self.tokenizer(strings, add_special_tokens=add_special_tokens, padding=False)["input_ids"]
+
+    def tok_encode_pair_batch(
+        self, contexts: list[str], continuations_list: list[list[str]]
+    ) -> tuple[list[list[list[int]]], list[list[list[int]]]]:
+        """Batched equivalent of `tok_encode_pair(context, continuations, pairwise=True)`
+        for a whole list of documents.
+
+        On a large benchmark, calling `tok_encode_pair` once per document tokenizes
+        every context and every continuation with its own tokenizer call: a batch
+        of 32 documents with 4 choices each makes 32 + 128 separate calls. This
+        does the same encoding (including the trailing-space handling) but makes
+        exactly two tokenizer calls for the whole batch: one for every context,
+        one for every continuation across every document.
+
+        Args:
+            contexts: One context string per document.
+            continuations_list: One list of continuation strings per document,
+                aligned with `contexts`.
+
+        Returns:
+            Tuple of (context token ids, continuation token ids), each a list
+            with one entry per document, matching what `tok_encode_pair(...,
+            pairwise=True)` returns for a single document.
+        """
+        if getattr(self, "move_trailing_context_space", True):
+            adjusted_contexts = []
+            adjusted_continuations_list = []
+            for context, continuations in zip(contexts, continuations_list):
+                n_spaces = len(context) - len(context.rstrip())
+                if n_spaces > 0:
+                    adjusted_continuations_list.append([context[-n_spaces:] + cont for cont in continuations])
+                    adjusted_contexts.append(context[:-n_spaces])
+                else:
+                    adjusted_continuations_list.append(continuations)
+                    adjusted_contexts.append(context)
+            contexts = adjusted_contexts
+            continuations_list = adjusted_continuations_list
+
+        # One call for every context in the batch.
+        context_encs = self._batch_tok_encode(contexts, add_special_tokens=self.add_special_tokens)
+
+        # One call for every continuation across every document, flattened so the
+        # tokenizer sees the whole batch at once, then split back up per document.
+        flat_continuations = [cont for continuations in continuations_list for cont in continuations]
+        flat_continuation_encs = self._batch_tok_encode(flat_continuations, add_special_tokens=False)
+
+        continuation_encs_list: list[list[list[int]]] = []
+        flat_index = 0
+        for continuations in continuations_list:
+            n = len(continuations)
+            continuation_encs_list.append(flat_continuation_encs[flat_index : flat_index + n])
+            flat_index += n
+
+        # Mirrors the pairwise branch of tok_encode_pair: strip a trailing eos
+        # token (it would otherwise make the model ignore the context) and repeat
+        # the context encoding once per continuation.
+        context_encs_list: list[list[list[int]]] = []
+        for context_enc, continuations in zip(context_encs, continuations_list):
+            if len(context_enc) > 0 and context_enc[-1] == self.tokenizer.eos_token_id:
+                context_enc = context_enc[:-1]
+            context_encs_list.append([context_enc] * len(continuations))
+
+        return context_encs_list, continuation_encs_list
+
     def tok_decode(self, tokens: torch.LongTensor) -> list[str]:
         return self.tokenizer.batch_decode(tokens, skip_special_tokens=True)

src/lighteval/models/transformers/transformers_model.py

@@ -958,13 +958,11 @@ def _loglikelihood_tokens(  # noqa: C901
 
             for batch in tqdm(dataloader, disable=self.disable_tqdm):
                 batch_contexts: list[str] = [self.prompt_manager.prepare_prompt(doc) for doc in batch]
-                batch_tokenized_contexts = []
-                batch_tokenized_continuations = []
-
-                for context, doc in zip(batch_contexts, batch):
-                    doc_contexts, doc_continuations = self.tok_encode_pair(context, doc.choices, pairwise=True)
-                    batch_tokenized_contexts.append(doc_contexts)
-                    batch_tokenized_continuations.append(doc_continuations)
+                # Two tokenizer calls for the whole batch instead of one pair of
+                # calls per document (see tok_encode_pair_batch).
+                batch_tokenized_contexts, batch_tokenized_continuations = self.tok_encode_pair_batch(
+                    batch_contexts, [doc.choices for doc in batch]
+                )
 
                 prepared_batch = self.prepare_batch_logprob(
                     tokenized_contexts=batch_tokenized_contexts,

tests/unit/models/test_abstract_model.py

@@ -57,3 +57,52 @@ def test_tok_encode_pair_move_trailing_context_space():
     model.move_trailing_context_space = False
     _, cont_kept = model.tok_encode_pair(context, continuation, pairwise=True)
     assert cont_kept == bare
+
+
+def test_tok_encode_pair_batch_matches_per_document_pairwise_encoding():
+    # tok_encode_pair_batch exists purely as a performance optimization: it must
+    # produce exactly what calling tok_encode_pair(..., pairwise=True) once per
+    # document would, just with fewer tokenizer calls.
+    model = DummyModel(config=DummyModelConfig(seed=42))
+    model._tokenizer = AutoTokenizer.from_pretrained("gpt2")
+
+    contexts = ["The capital of France is", "Question: 2+2= ", "No trailing space here"]
+    continuations_list = [
+        [" Paris", " London", " Berlin"],
+        ["4", "five"],
+        ["!", "?"],
+    ]
+
+    expected_contexts = []
+    expected_continuations = []
+    for context, continuations in zip(contexts, continuations_list):
+        context_enc, continuation_enc = model.tok_encode_pair(context, continuations, pairwise=True)
+        expected_contexts.append(context_enc)
+        expected_continuations.append(continuation_enc)
+
+    batch_contexts, batch_continuations = model.tok_encode_pair_batch(contexts, continuations_list)
+
+    assert batch_contexts == expected_contexts
+    assert batch_continuations == expected_continuations
+
+
+def test_tok_encode_pair_batch_respects_move_trailing_context_space():
+    model = DummyModel(config=DummyModelConfig(seed=42))
+    model._tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    contexts = ["Answer: "]
+    continuations_list = [["Paris"]]
+    bare = [model.tok_encode("Paris", add_special_tokens=False)]
+
+    model.move_trailing_context_space = True
+    _, cont_moved = model.tok_encode_pair_batch(contexts, continuations_list)
+    assert cont_moved[0] != bare
+
+    model.move_trailing_context_space = False
+    _, cont_kept = model.tok_encode_pair_batch(contexts, continuations_list)
+    assert cont_kept[0] == bare
+
+
+def test_batch_tok_encode_empty_list_returns_empty_list():
+    model = DummyModel(config=DummyModelConfig(seed=42))
+    model._tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    assert model._batch_tok_encode([], add_special_tokens=True) == []

tests/unit/models/test_transformers_model.py

@@ -394,6 +394,30 @@ def mock_gather(tensor):
         # Restore original gather function
         self.model.accelerator.gather_for_metrics = lambda x: x
 
+    @patch("lighteval.models.transformers.transformers_model.DataLoader")
+    def test_loglikelihood_batches_tokenization_per_minibatch(self, mock_dataloader):
+        """The scoring loop should tokenize each mini-batch with two calls to
+        _batch_tok_encode (one for every context, one for every continuation),
+        not one pair of calls per document. That's the whole point of
+        tok_encode_pair_batch: with 3 documents below, a per-document loop
+        would make 6 calls; batched, it makes 2 regardless of how many
+        documents or choices are in the mini-batch."""
+        docs = [
+            Doc(query="What is the capital of France?", choices=["London", "Berlin", "Paris", "Madrid"], gold_index=2),
+            Doc(query="What is 2+2?", choices=["3", "4", "5"], gold_index=1),
+            Doc(query="What color is the sky?", choices=["Blue", "Green"], gold_index=0),
+        ]
+        mock_dataloader.return_value = [docs]
+        if hasattr(self.model.accelerator, "prepare"):
+            self.model.accelerator.prepare = Mock(side_effect=lambda x: x)
+
+        with patch.object(
+            TransformersModel, "_batch_tok_encode", wraps=self.model._batch_tok_encode
+        ) as mock_batch_encode:
+            self.model._loglikelihood_tokens(docs)
+
+        self.assertEqual(mock_batch_encode.call_count, 2)
+
 
 class TestTransformersModelUseChatTemplate(unittest.TestCase):
     @patch("lighteval.models.transformers.transformers_model.Accelerator")