update

audio_processing_mllama.py

@@ -20,7 +20,7 @@ def build_audio_tokens(text: List[str], audio_features: Union[Dict, List[List[np
     return text
             
 def get_num_embeddings(num_framses, adapter_kernel_size=7, adapter_stride=4) -> int:
-    return math.ceil((num_framses - adapter_kernel_size) / adapter_stride) + 1 + 2 # 2 = <|begin_of_audio|>, <|end_of_audio|>
+    return math.ceil((num_framses - adapter_kernel_size + adapter_stride) / adapter_stride) + 1 + 2 # 2 = <|begin_of_audio|>, <|end_of_audio|>
 
 
 class MllamaAudioFeatureExtractor(SeamlessM4TFeatureExtractor):

config.json

@@ -142,10 +142,10 @@
   },
   "audio_token_index": 128257,
   "auto_map": {
-    "AutoConfig": "AlexHung29629/test_mllama_11B_v3--configuration_llama3.Llama3Config",
-    "AutoModel": "AlexHung29629/test_mllama_11B_v3--modeling_llama3.Llama3ForConditionalGeneration",
-    "AutoProcessor": "AlexHung29629/test_mllama_11B_v3--processing_mllama.MllamaProcessor",
-    "AutoFeatureExtractor": "AlexHung29629/test_mllama_11B_v3--audio_processing_mllama.MllamaAudioFeatureExtractor"
+    "AutoConfig": "configuration_llama3.Llama3Config",
+    "AutoModel": "modeling_llama3.Llama3ForConditionalGeneration",
+    "AutoProcessor": "processing_mllama.MllamaProcessor",
+    "AutoFeatureExtractor": "audio_processing_mllama.MllamaAudioFeatureExtractor"
   },
   "image_token_index": 128256,
   "model_type": "llama3",

mllama_audio_model.py

@@ -1,22 +1,62 @@
 from typing import Optional, Tuple, Union
-import math
 import torch
 from torch import nn
 from transformers.modeling_outputs import BaseModelOutput
-from transformers import Wav2Vec2BertModel, Wav2Vec2BertConfig, Wav2Vec2BertPreTrainedModel, MllamaPreTrainedModel
-from transformers.models.mllama.configuration_mllama import MllamaTextConfig
+from transformers import Wav2Vec2BertModel, Wav2Vec2BertConfig, MllamaPreTrainedModel
+from transformers.models.wav2vec2_bert.modeling_wav2vec2_bert import Wav2Vec2BertAdapterLayer
 from .configuration_llama3 import Llama3Config
 
+class AudioAdapter(nn.Module):
+    def __init__(self, config: Wav2Vec2BertConfig):
+        super().__init__()
+        # feature dim might need to be down-projected
+        if config.output_hidden_size != config.hidden_size:
+            self.proj = nn.Linear(config.hidden_size, config.output_hidden_size)
+        else:
+            self.proj = None
+        self.layers = nn.ModuleList(Wav2Vec2BertAdapterLayer(config) for _ in range(config.num_adapter_layers))
+        self.layerdrop = config.layerdrop
 
+        self.kernel_size = config.adapter_kernel_size
+        self.stride = config.adapter_stride
+
+    def _compute_sub_sample_lengths_from_attention_mask(self, seq_lens):
+        if seq_lens is None:
+            return seq_lens
+        pad = self.kernel_size // 2
+        seq_lens = ((seq_lens + 2 * pad - self.kernel_size) / self.stride) + 1
+        return seq_lens.floor()
+
+    def forward(self, hidden_states, attention_mask=None):
+        # down project hidden_states if necessary
+        if self.proj is not None:
+            hidden_states = self.proj(hidden_states)
+
+        sub_sampled_lengths = None
+        if attention_mask is not None:
+            sub_sampled_lengths = (attention_mask.size(1) - (1 - attention_mask.int()).sum(1)).to(hidden_states.device)
+
+        for layer in self.layers:
+            layerdrop_prob = torch.rand([])
+            sub_sampled_lengths = self._compute_sub_sample_lengths_from_attention_mask(sub_sampled_lengths)
+            if not self.training or (layerdrop_prob > self.layerdrop):
+                hidden_states = layer(
+                    hidden_states, attention_mask=attention_mask, sub_sampled_lengths=sub_sampled_lengths
+                )
+
+        return hidden_states
+
+        
 class Llama3Embedding(MllamaPreTrainedModel):
     config_class = Llama3Config
     base_model_prefix = "audio_model"
     def __init__(self, config: Llama3Config):
         super().__init__(config)
-        assert config.audio_config.add_adapter is True, f'{type(self).__name__} requires add adapter to be true.'
         assert config.audio_config.output_hidden_size == config.text_config.hidden_size
         self.text_embeddings = nn.Embedding(config.text_config.vocab_size, config.text_config.hidden_size, config.text_config.pad_token_id)
-        self.audio_embedding = Wav2Vec2BertModel(config.audio_config)
+        config.audio_config.add_adapter = False
+        self.audio_encoder = Wav2Vec2BertModel(config.audio_config)
+        self.audio_adapter = AudioAdapter(config.audio_config)
         self.start_of_audio = nn.Parameter(data=torch.zeros((1, config.audio_config.output_hidden_size)), requires_grad=True)
         self.end_of_audio = nn.Parameter(data=torch.zeros((1, config.audio_config.output_hidden_size)), requires_grad=True)
         self.text_config = config.text_config
@@ -30,7 +70,8 @@ class Llama3Embedding(MllamaPreTrainedModel):
         if audio_features is None:
             return input_embeddings
         bs, max_num_img, l, d = audio_features.shape
-        audio_embeddings = self.audio_embedding(input_features=audio_features.view((bs*max_num_img, l, d)))['last_hidden_state']
+        audio_embeddings = self.audio_encoder(input_features=audio_features.view((bs*max_num_img, l, d)))['last_hidden_state']
+        audio_embeddings = self.audio_adapter(audio_embeddings)
         audio_embeddings = audio_embeddings.view((bs, max_num_img, -1, self.start_of_audio.shape[-1]))        
 
         for i in range(bs):