Repository: huggingface/transformers
File: tests/generation/test_utils.py (Python)

Code before:
        # 1. Sanity check: the model's `prepare_inputs_for_generation` comes from `GenerationMixin`
        self.assertTrue("GenerationMixin" in str(model.prepare_inputs_for_generation))

        # 2. If we pass input ids by themselves, we should get back the same input ids
        model_inputs = model.prepare_inputs_for_generation(input_ids)
        self.assertTrue(torch.all(model_inputs["input_ids"] == input_ids))

        # 3. `use_cache` (and other kwargs) are forwarded
        self.assertFalse("use_cache" in model_inputs)  # From the previous input, there is no `use_cache`
        model_inputs = model.prepare_inputs_for_generation(input_ids, use_cache=True, foo="bar")
        self.assertTrue(model_inputs["use_cache"] is True)
        self.assertTrue(model_inputs["foo"] == "bar")

        # 4. We never discard data from input ids and expect it to be already slice
        init_input_ids = input_ids[:, :2]
        init_dynamic_cache = model(init_input_ids, past_key_values=dynamic_cache).past_key_values
        init_cache_position = cache_position[dynamic_cache.get_seq_length() :]
        model_inputs = model.prepare_inputs_for_generation(
            input_ids,
            past_key_values=init_dynamic_cache,
            cache_position=init_cache_position,
            attention_mask=attention_mask,
            position_ids=position_ids,
        )
        self.assertTrue("past_key_values" in model_inputs)
        self.assertTrue(torch.all(model_inputs["cache_position"] == init_cache_position))
        self.assertTrue(model_inputs["input_ids"].shape[-1] == input_ids.shape[1])
        self.assertTrue(model_inputs["position_ids"].shape[-1] == input_ids.shape[1])
        self.assertTrue(model_inputs["attention_mask"].shape[-1] == input_ids.shape[1])

        # 4.1: We do not have to pass prepared `cache_positions` to slice the data!
        # Data is sliced based on input ids length
        model_inputs = model.prepare_inputs_for_generation(
            init_input_ids,
            past_key_values=init_dynamic_cache,
            attention_mask=attention_mask,
            position_ids=position_ids,
        )
        self.assertTrue("past_key_values" in model_inputs)
        self.assertTrue(model_inputs["cache_position"] is None)
        self.assertTrue(model_inputs["input_ids"].shape[-1] == init_input_ids.shape[1])
        self.assertTrue(model_inputs["position_ids"].shape[-1] == init_input_ids.shape[1])
        self.assertTrue(model_inputs["attention_mask"].shape[-1] == 3)  # attn mask is never sliced, we need PAD mask

        # 5. If we pass a `static_cache`, the attention mask will be prepared as a static shape 4D mask
        max_cache_len = 10
        static_cache = StaticCache(config=config, max_cache_len=max_cache_len)
        static_cache = model(init_input_ids, past_key_values=static_cache).past_key_values
        model_inputs = model.prepare_inputs_for_generation(
            init_input_ids, past_key_values=static_cache, attention_mask=attention_mask, cache_position=cache_position
        )

Review comment (bug):
for this block of assertions I would use `assertEqual` so it if it fails we get the two values explicitely displayed

Code after:
        # 1. Sanity check: the model's `prepare_inputs_for_generation` comes from `GenerationMixin`
        self.assertTrue("GenerationMixin" in str(model.prepare_inputs_for_generation))

        # 2. If we pass input ids by themselves, we should get back the same input ids
        model_inputs = model.prepare_inputs_for_generation(input_ids)
        self.assertListEqual(model_inputs["input_ids"].tolist(), input_ids.tolist())

        # 3. `use_cache` (and other kwargs) are forwarded
        self.assertFalse("use_cache" in model_inputs)  # From the previous input, there is no `use_cache`
        model_inputs = model.prepare_inputs_for_generation(input_ids, use_cache=True, foo="bar")
        self.assertTrue(model_inputs["use_cache"] is True)
        self.assertEqual(model_inputs["foo"], "bar")

        # 4. We never discard data from input ids and expect it to be already slice
        init_input_ids = input_ids[:, :2]
        init_dynamic_cache = model(init_input_ids, past_key_values=dynamic_cache).past_key_values
        init_cache_position = cache_position[dynamic_cache.get_seq_length() :]
        model_inputs = model.prepare_inputs_for_generation(
            input_ids,
            past_key_values=init_dynamic_cache,
            cache_position=init_cache_position,
            attention_mask=attention_mask,
            position_ids=position_ids,
        )
        self.assertTrue("past_key_values" in model_inputs)
        self.assertListEqual(model_inputs["cache_position"].tolist(), init_cache_position.tolist())
        self.assertEqual(model_inputs["input_ids"].shape[-1], input_ids.shape[1])
        self.assertEqual(model_inputs["position_ids"].shape[-1], input_ids.shape[1])
        self.assertEqual(model_inputs["attention_mask"].shape[-1], input_ids.shape[1])

        # 4.1: We do not have to pass prepared `cache_positions` to slice the data!
        # Data is sliced based on input ids length
        model_inputs = model.prepare_inputs_for_generation(
            init_input_ids,
            past_key_values=init_dynamic_cache,
            attention_mask=attention_mask,
            position_ids=position_ids,
        )
        self.assertTrue("past_key_values" in model_inputs)
        self.assertEqual(model_inputs["cache_position"], None)
        self.assertEqual(model_inputs["input_ids"].shape[-1], init_input_ids.shape[1])
        self.assertEqual(model_inputs["position_ids"].shape[-1], init_input_ids.shape[1])
        self.assertEqual(model_inputs["attention_mask"].shape[-1], 3)  # attn mask is never sliced, we need PAD mask

        # 5. If we pass a `static_cache`, the attention mask will be prepared as a static shape 4D mask
        max_cache_len = 10
        static_cache = StaticCache(config=config, max_cache_len=max_cache_len)
        static_cache = model(init_input_ids, past_key_values=static_cache).past_key_values
        model_inputs = model.prepare_inputs_for_generation(
            init_input_ids, past_key_values=static_cache, attention_mask=attention_mask, cache_position=cache_position
        )