From fd70464fa74c101ab3bc60e0c3db7d5e3b75fe90 Mon Sep 17 00:00:00 2001 From: Raushan Turganbay Date: Fri, 11 Oct 2024 15:41:46 +0200 Subject: [PATCH] Fix flaky tests (#34069) * fix mllama only * allow image token index --- src/transformers/models/mllama/modeling_mllama.py | 2 +- utils/check_config_attributes.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/src/transformers/models/mllama/modeling_mllama.py b/src/transformers/models/mllama/modeling_mllama.py index 0bc77eaeec..e486e149e3 100644 --- a/src/transformers/models/mllama/modeling_mllama.py +++ b/src/transformers/models/mllama/modeling_mllama.py @@ -2214,7 +2214,7 @@ class MllamaForConditionalGeneration(MllamaPreTrainedModel, GenerationMixin): # If we're in pre-fill or cacheless decoding step, then we need pixel_values and aspect ratios # to compute image hidden states, otherwise they are cached within each cross attn layer - if (input_ids == self.config.image_token_index).any(): + if cache_position[0] == 0: model_inputs["pixel_values"] = pixel_values model_inputs["aspect_ratio_ids"] = aspect_ratio_ids model_inputs["aspect_ratio_mask"] = aspect_ratio_mask diff --git a/utils/check_config_attributes.py b/utils/check_config_attributes.py index 7bd9379636..83fe07fef2 100644 --- a/utils/check_config_attributes.py +++ b/utils/check_config_attributes.py @@ -243,6 +243,7 @@ def check_attribute_being_used(config_class, attributes, default_value, source_s "pad_index", "unk_index", "mask_index", + "image_token_index", # for VLMs "image_size", "use_cache", "out_features",