Update mllama_audio_model.py

mllama_audio_model.py

@@ -29,18 +29,16 @@ class MllamaAudioModel(MllamaPreTrainedModel):
         input_ids: torch.LongTensor = None,
         return_dict: Optional[bool] = None,
     ) -> Union[BaseModelOutput, Tuple[torch.Tensor, ...]]:
+        bs, max_num_img, l, d = audio_features.shape
         input_embeddings = self.text_embedding(torch.clamp(input_ids, min=0))
-        audio_embeddings = self.audio_embedding(input_features=audio_features[0])['last_hidden_state']
-        print(f"{audio_embeddings.shape=}")
-        bs, max_num_img, _, _ = audio_features.shape
+        audio_embeddings = self.audio_embedding(input_features=audio_features.view((bs*max_num_img, l, d)))['last_hidden_state']
+        audio_embeddings = audio_embeddings.view((bs, max_num_img, -1, start_of_audio.shape[-1]))        
 
         for i in range(bs):
             for j in range(max_num_img):
                 audio_id = -1 - j
                 if torch.any(input_ids[i] == audio_id):
                     idx = input_ids[i] == audio_id
-                    print(f"{audio_features[i, j].shape=}")
-                    print(f"{self.start_of_audio.shape=}")
                     input_embeddings[i][idx] = torch.concat([self.start_of_audio, audio_embeddings[i, j], self.end_of_audio])
         
         if return_dict: