diff --git a/src/lighteval/models/transformers/transformers_model.py b/src/lighteval/models/transformers/transformers_model.py index 64e790a2f..94c117b67 100644 --- a/src/lighteval/models/transformers/transformers_model.py +++ b/src/lighteval/models/transformers/transformers_model.py @@ -713,6 +713,18 @@ def _padded_greedy_until( add_special_tokens=self.add_special_tokens, ).to(self.device) + # `truncation="longest_first"` above drops the overflow of any context longer than + # `max_context_continuation_size_allowed`, so `tokenized` alone no longer tells us how many + # tokens each context holds. `contexts` holds strings, so `len(context)` is a character + # count and cannot be compared against a token count: re-tokenize without a cap to get the + # real token counts (`verbose=False` silences the tokenizer's over-long-input warning). + context_token_counts = [ + len(input_ids) + for input_ids in self.tokenizer( + contexts, add_special_tokens=self.add_special_tokens, verbose=False + )["input_ids"] + ] + # The main question for this step is the following: # Would we rather truncate the prompt to allow generation to go to max_new_tokens, at the risk # of losing some meaning, or have some generations that are exceedingly short? @@ -739,7 +751,10 @@ def _padded_greedy_until( input_ids=tokenized["input_ids"], input_lengths=[len(item == 1) for item in tokenized["attention_mask"]], input_mask=tokenized["attention_mask"], - truncated=[max(len(c) - tokenized["input_ids"].shape[1], 0) for c in contexts], + truncated=[ + max(context_token_count - tokenized["input_ids"].shape[1], 0) + for context_token_count in context_token_counts + ], padded=[sum(mask == 0) for mask in tokenized["attention_mask"]], ) diff --git a/tests/unit/models/test_transformers_model.py b/tests/unit/models/test_transformers_model.py index da7c925ae..690e0804e 100644 --- a/tests/unit/models/test_transformers_model.py +++ b/tests/unit/models/test_transformers_model.py @@ -426,5 +426,49 @@ def test_transformers_model_use_chat_template_with_different_model_names( self.assertEqual(model.use_chat_template, model._tokenizer.chat_template is not None) +class TestTransformersModelTruncatedTokensCount(unittest.TestCase): + """`truncated_tokens_count` must count tokens, not `len(context) - token_count`.""" + + @patch("lighteval.models.transformers.transformers_model.Accelerator") + def setUp(self, mock_accelerator): + """Set up shared model instance for all tests.""" + mock_accelerator_instance = Mock() + mock_accelerator_instance.device = torch.device("cpu") + mock_accelerator_instance.is_main_process = True + mock_accelerator_instance.prepare = lambda x: x + mock_accelerator_instance.gather = lambda x: x + mock_accelerator_instance.gather_for_metrics = lambda x: x + mock_accelerator.return_value = mock_accelerator_instance + # batch_size is set explicitly so the model is not run to detect the largest batch size + self.model = TransformersModel(TransformersModelConfig(model_name="gpt2", batch_size=1)) + + def _truncated_for(self, query: str) -> list[int]: + """Runs _padded_greedy_until with generation stubbed out and returns Batch.truncated.""" + doc = Doc(query=query, choices=[], gold_index=0, generation_size=4, stop_sequences=["Question:"]) + captured = {} + + def fake_generate(**kwargs): + captured["truncated"] = list(kwargs["batch"].truncated) + return [ModelResponse(text=["stub"], output_tokens=[[0]], input_tokens=[0])] + + with patch("lighteval.models.transformers.transformers_model.DataLoader") as mock_dataloader: + mock_dataloader.return_value = [[doc]] + with patch.object(self.model, "_generate", side_effect=fake_generate): + self.model._padded_greedy_until([doc]) + + return captured["truncated"] + + def test_context_fitting_the_context_window_is_not_reported_as_truncated(self): + query = "The capital of France is" + self.assertLess(len(self.model.tokenizer(query)["input_ids"]), self.model.max_length) + self.assertEqual(self._truncated_for(query), [0]) + + def test_truncated_tokens_count_matches_the_dropped_token_count(self): + query = "hello world " * 700 + token_count = len(self.model.tokenizer(query)["input_ids"]) + self.assertGreater(token_count, self.model.max_length) + self.assertEqual(self._truncated_for(query), [token_count - self.model.max_length]) + + if __name__ == "__main__": unittest.main()