Update modeling_llama3.py

modeling_llama3.py

@@ -6,7 +6,7 @@ import torch.utils.checkpoint
 from torch import nn
 
 import transformers
-from transformers import MllamaPreTrainedModel, MllamaVisionModel, MllamaForCausalLM, AutoModel, AutoModelForCausalLM
+from transformers import MllamaPreTrainedModel, MllamaVisionModel, MllamaForCausalLM, Wav2Vec2BertConfig, AutoModel, AutoModelForCausalLM
 from transformers.cache_utils import Cache, StaticCache
 from transformers.generation import GenerationMixin
 from transformers.modeling_attn_mask_utils import AttentionMaskConverter
@@ -14,7 +14,6 @@ from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutpu
 from transformers.utils import logging
 from transformers.models.mllama.modeling_mllama import _prepare_cross_attention_mask, MllamaCrossAttentionDecoderLayer, MllamaSelfAttentionDecoderLayer, MllamaTextRMSNorm, MllamaRotaryEmbedding
 from transformers.models.mllama.configuration_mllama import MllamaTextConfig
-
 from .configuration_llama3 import Llama3Config
 from .mllama_audio_model import Llama3Embedding
 
@@ -25,27 +24,27 @@ class Llama3PreTrainedModel(MllamaPreTrainedModel):
     config_class = Llama3Config
     base_model_prefix = "model"
 
-class Llama3TextModel(Llama3PreTrainedModel):
-    config_class = Llama3Config
+class Llama3TextModel(MllamaPreTrainedModel):
+    config_class = MllamaTextConfig
     base_model_prefix = "language_model.model"
 
-    def __init__(self, config: Llama3Config):
+    def __init__(self, config: MllamaTextConfig, audio_config: Wav2Vec2BertConfig):
         super().__init__(config)
-        self.padding_idx = config.text_config.pad_token_id
-        self.vocab_size = config.text_config.vocab_size
-        self.embed_tokens = Llama3Embedding(config.audio_config, config.text_config)
-        self.cross_attention_layers = config.text_config.cross_attention_layers
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        #self.embed_tokens = Llama3Embedding(audio_config, config)
+        self.cross_attention_layers = config.cross_attention_layers
 
         layers = []
-        for layer_idx in range(config.text_config.num_hidden_layers):
+        for layer_idx in range(config.num_hidden_layers):
             if layer_idx in self.cross_attention_layers:
-                layers.append(MllamaCrossAttentionDecoderLayer(config.text_config, layer_idx))
+                layers.append(MllamaCrossAttentionDecoderLayer(config, layer_idx))
             else:
-                layers.append(MllamaSelfAttentionDecoderLayer(config.text_config, layer_idx))
+                layers.append(MllamaSelfAttentionDecoderLayer(config, layer_idx))
 
         self.layers = nn.ModuleList(layers)
-        self.norm = MllamaTextRMSNorm(config.text_config.hidden_size, eps=config.text_config.rms_norm_eps)
-        self.rotary_emb = MllamaRotaryEmbedding(config=config.text_config)
+        self.norm = MllamaTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = MllamaRotaryEmbedding(config=config)
         self.gradient_checkpointing = False
         self.post_init()
 
@@ -57,8 +56,8 @@ class Llama3TextModel(Llama3PreTrainedModel):
 
     def forward(
         self,
-        input_ids: Optional[torch.LongTensor] = None,
-        audio_features: Optional[torch.Tensor] = None,
+        #input_ids: Optional[torch.LongTensor] = None,
+        #audio_features: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.LongTensor] = None,
         cross_attention_states: Optional[torch.FloatTensor] = None,
@@ -94,15 +93,15 @@ class Llama3TextModel(Llama3PreTrainedModel):
         torch.Size([1, 13, 4096])
         ```
         """
-        output_attentions = output_attentions if output_attentions is not None else self.config.text_config.output_attentions
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
-        use_cache = use_cache if use_cache is not None else self.config.text_config.use_cache
-        return_dict = return_dict if return_dict is not None else self.config.text_config.use_return_dict
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        if (input_ids is None) ^ (inputs_embeds is not None):
-            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+        #if (input_ids is None) ^ (inputs_embeds is not None):
+        #    raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
 
         if self.gradient_checkpointing and self.training and use_cache:
             logger.warning_once(
@@ -110,8 +109,9 @@ class Llama3TextModel(Llama3PreTrainedModel):
             )
             use_cache = False
 
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids=input_ids, audio_features=audio_features)
+        #if inputs_embeds is None:
+        #    inputs_embeds = self.embed_tokens(input_ids=input_ids, audio_features=audio_features)
+        
 
         hidden_states = inputs_embeds
 
@@ -214,7 +214,7 @@ class Llama3TextModel(Llama3PreTrainedModel):
         past_key_values: Cache,
         output_attentions: bool,
     ):
-        if self.config.text_config._attn_implementation == "flash_attention_2":
+        if self.config._attn_implementation == "flash_attention_2":
             if attention_mask is not None and 0.0 in attention_mask:
                 return attention_mask
             return None
@@ -258,7 +258,7 @@ class Llama3TextModel(Llama3PreTrainedModel):
         )
 
         if (
-            self.config.text_config._attn_implementation == "sdpa"
+            self.config._attn_implementation == "sdpa"
             and attention_mask is not None
             and attention_mask.device.type == "cuda"
             and not output_attentions
@@ -306,25 +306,26 @@ class Llama3TextModel(Llama3PreTrainedModel):
 
         return causal_mask
 
-class Llama3ForCausalLM(Llama3PreTrainedModel, GenerationMixin):
-    config_class = Llama3Config
+class Llama3ForCausalLM(MllamaPreTrainedModel, GenerationMixin):
+    config_class = MllamaTextConfig
     base_model_prefix = "model"
-    _tied_weights_keys = ["lm_head.weight"]
+    #_tied_weights_keys = ["lm_head.weight"]
 
-    def __init__(self, config):
+    def __init__(self, config: MllamaTextConfig):
         super().__init__(config)
-        self.text_config = config.get_text_config()
-        self.vocab_size = self.text_config.vocab_size
+        self.vocab_size = config.vocab_size
         self.model = Llama3TextModel._from_config(config, attn_implementation=config._attn_implementation)
-        self.lm_head = nn.Linear(self.text_config.hidden_size, self.vocab_size, bias=False)
+        self.lm_head = nn.Linear(config.hidden_size, self.vocab_size, bias=False)
 
         self.post_init()
 
     def get_input_embeddings(self):
-        return self.model.embed_tokens.text_embeddings
+        #return self.model.embed_tokens.text_embeddings
+        return None
 
     def set_input_embeddings(self, value):
-        self.model.embed_tokens.text_embeddings = value
+        #self.model.embed_tokens.text_embeddings = value
+        pass
 
     def get_output_embeddings(self):
         return self.lm_head
@@ -340,7 +341,7 @@ class Llama3ForCausalLM(Llama3PreTrainedModel, GenerationMixin):
 
     def forward(
         self,
-        input_ids: torch.LongTensor = None,
+        #input_ids: torch.LongTensor = None,
         attention_mask: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.LongTensor] = None,
         cross_attention_states: Optional[torch.LongTensor] = None,
@@ -357,15 +358,15 @@ class Llama3ForCausalLM(Llama3PreTrainedModel, GenerationMixin):
         num_logits_to_keep: int = 0,
         **loss_kwargs,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
-        output_attentions = output_attentions if output_attentions is not None else self.text_config.output_attentions
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.text_config.output_hidden_states
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
-        return_dict = return_dict if return_dict is not None else self.text_config.use_return_dict
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
         outputs = self.model(
-            input_ids=input_ids,
+            #input_ids=input_ids,
             cross_attention_states=cross_attention_states,
             attention_mask=attention_mask,
             position_ids=position_ids,
@@ -402,7 +403,9 @@ class Llama3ForCausalLM(Llama3PreTrainedModel, GenerationMixin):
 AutoModelForCausalLM.register(Llama3Config, Llama3ForCausalLM)
 transformers.Llama3ForCausalLM = Llama3ForCausalLM
 
-class Llama3ForConditionalGeneration(Llama3PreTrainedModel, GenerationMixin):
+class Llama3ForConditionalGeneration(MllamaPreTrainedModel, GenerationMixin):
+    config_class = Llama3Config
+    base_model_prefix = "model"
     _supports_quantized_cache = False  # quant cache not supported in encoder-decoder setting
 
     def __init__(self, config: Llama3Config):
@@ -415,6 +418,7 @@ class Llama3ForConditionalGeneration(Llama3PreTrainedModel, GenerationMixin):
 
         self.vision_model = MllamaVisionModel._from_config(config.vision_config)
         self.language_model = Llama3ForCausalLM._from_config(config)
+        self.embeddings = Llama3Embedding(config.audio_config, config.text_config)
         self.multi_modal_projector = nn.Linear(
             config.vision_config.vision_output_dim,
             config.text_config.hidden_size,