Merge pull request #4 from alex-jw-brooks/conformer_pr_updates

feature attention mask bugfix

Author: alex-jw-brooks

src/transformers/models/granite_speech/feature_extraction_granite_speech.py

@@ -83,7 +83,7 @@ def __call__(
         audio_embed_sizes = self._get_num_audio_features(audio_lengths)
         speech_inputs["audio_embed_sizes"] = audio_embed_sizes
         # todo: input_features_mask is not a great name, because input_features and input_features mask have different shapes (before/after the projector)
-        speech_inputs["input_features_mask"] = torch.arange(max(audio_embed_sizes)).view(1, -1) <= torch.tensor(
+        speech_inputs["input_features_mask"] = torch.arange(max(audio_embed_sizes)).view(1, -1) < torch.tensor(
             audio_embed_sizes
         ).view(-1, 1)
         return BatchFeature(data=speech_inputs)

src/transformers/models/granite_speech/modeling_granite_speech.py

@@ -1253,6 +1253,8 @@ def get_merged_audio_embeddings(self, input_ids, audio_features, input_features_
         and potentially labels.
         """
         is_audio_index = input_ids == self.config.audio_token_index
+        assert torch.all(is_audio_index.int().sum(dim=1) == input_features_mask.int().sum(dim=1)).item(), \
+            "number of features should align"
         llm_input_ids = torch.where(is_audio_index, 0, input_ids)
         inputs_embeds = self.language_model.get_input_embeddings()(llm_input_ids)  # [bsz, # features, hidden size]