Update mllama_audio_model.py

mllama_audio_model.py

@@ -4,6 +4,7 @@ import torch
 from torch import nn
 from transformers.modeling_outputs import BaseModelOutput
 from transformers import Wav2Vec2BertModel, Wav2Vec2BertConfig, MllamaPreTrainedModel
+from transformers.models.llama.modeling_llama import LlamaRMSNorm
 from transformers.models.wav2vec2_bert.modeling_wav2vec2_bert import _compute_new_attention_mask, _prepare_4d_attention_mask, Wav2Vec2BertFeedForward, Wav2Vec2BertSelfAttention, Wav2Vec2BertFeatureProjection
 from .configuration_llama3 import Llama3Config
 
@@ -50,6 +51,8 @@ class AudioAdapter(nn.Module):
         else:
             self.proj = None
         self.layers = nn.ModuleList(Wav2Vec2BertAdapterLayer(config) for _ in range(config.num_adapter_layers))
+        self.final_adapter_norm = LlamaRMSNorm(config.output_hidden_size)
+        self.final_adapter_norm.weight.data.fill_(0.4)
 
     def forward(self, hidden_states, attention_mask=None):
         # down project hidden_states if necessary
@@ -61,6 +64,8 @@ class AudioAdapter(nn.Module):
                 hidden_states
             )
 
+        hidden_states = self.final_adapter_norm(hidden_states)
+
         return hidden_states