diff --git a/src/lighteval/models/abstract_model.py b/src/lighteval/models/abstract_model.py index ccddb404c..a23069d32 100644 --- a/src/lighteval/models/abstract_model.py +++ b/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) diff --git a/src/lighteval/models/transformers/transformers_model.py b/src/lighteval/models/transformers/transformers_model.py index 64e790a2f..1cd31f887 100644 --- a/src/lighteval/models/transformers/transformers_model.py +++ b/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, diff --git a/tests/unit/models/test_abstract_model.py b/tests/unit/models/test_abstract_model.py index d062fb556..d91290564 100644 --- a/tests/unit/models/test_abstract_model.py +++ b/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) == [] diff --git a/tests/unit/models/test_transformers_model.py b/tests/unit/models/test_transformers_model.py index da7c925ae..c9ff50299 100644 --- a/tests/unit/models/test_transformers_model.py +++ b/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")