Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,143 @@
+import torch
+from torch import nn
+from dataclasses import dataclass
+from transformers import AutoConfig, WhisperConfig, LlamaConfig, WhisperModel, AutoModelForCausalLM
+from typing import Callable, Optional, Union
+
+from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.cache_utils import Cache
+from transformers.processing_utils import Unpack
+
+
+@dataclass
+class ManaraConfig:
+    audio_config: WhisperConfig
+    text_config: LlamaConfig
+
+
+class MultiModalProjector(nn.Module):
+
+    def __init__(self, config: ManaraConfig):
+        super().__init__()
+        self.linear_1 = nn.Linear(config.audio_config.d_model,
+                                  config.text_config.hidden_size, bias=False)
+        self.act = nn.SiLU()
+        self.pool = nn.Conv1d(
+            config.text_config.hidden_size,
+            config.text_config.hidden_size,
+            kernel_size=2,
+            stride=2,
+            bias=False
+        )
+        self.linear_2 = nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size,
+                                  bias=False)
+
+    def forward(self, audio_features):
+        hidden_states = self.linear_1(audio_features)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.pool(hidden_states.transpose(1, 2)).transpose(1, 2)
+        hidden_states = self.linear_2(hidden_states)
+        return hidden_states
+
+
+class ManaraForConditionalGeneration(nn.Module):
+    def __init__(self):
+        super().__init__()
+        audio = 'openai/whisper-large-v3'
+        text = 'babs/llama-multi-lm'
+        self.text_config = LlamaConfig.from_pretrained(text)
+        self.audio_config = WhisperConfig.from_pretrained(audio)
+
+        config = ManaraConfig(self.audio_config, self.text_config)
+        self.vocab_size = config.text_config.vocab_size
+
+        audio_tower = WhisperModel.from_pretrained(audio,
+                                                   attn_implementation='flash_attention_2')
+        self.audio_tower = audio_tower.encoder
+        self.audio_tower.requires_grad_(False)
+        del audio_tower
+
+        self.language_model = AutoModelForCausalLM.from_pretrained(text,
+                                                                   attn_implementation='flash_attention_2')
+        self.multi_modal_projector = MultiModalProjector(config)
+        self.audio_token_id = 128002
+
+    def get_input_embeddings(self):
+        return self.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.language_model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.language_model.get_output_embeddings()
+
+    def set_output_embeddings(self, new_embeddings):
+        self.language_model.set_output_embeddings(new_embeddings)
+
+    def set_decoder(self, decoder):
+        self.language_model.set_decoder(decoder)
+
+    def get_decoder(self):
+        return self.language_model.get_decoder()
+
+    def get_audio_embeds(self, input_features: torch.FloatTensor):
+        audio_outputs = self.audio_tower(input_features)
+        audio_hidden_states = audio_outputs.last_hidden_state
+        # audio_hidden_states = audio_hidden_states.reshape(-1, self.config.audio_config.intermediate_size)
+        audio_embeds = self.multi_modal_projector(audio_hidden_states)
+        return audio_embeds
+
+    def forward(
+            self,
+            input_ids: Optional[torch.LongTensor] = None,
+            input_features: Optional[torch.FloatTensor] = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[Cache] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            cache_position: Optional[torch.LongTensor] = None,
+            logits_to_keep: Union[int, torch.Tensor] = 0,
+            **kwargs,
+    ) -> CausalLMOutputWithPast:
+
+        if inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings()(input_ids)
+
+        if input_features is not None:
+            audio_embeds = self.get_audio_embeds(input_features)
+
+            audio_token_mask = (input_ids == self.audio_token_id)
+
+            if audio_token_mask.sum() != audio_embeds.shape[0] * audio_embeds.shape[1]:
+                raise ValueError(
+                    f"The number of audio tokens in input_ids ({audio_token_mask.sum()}) does not match "
+                    f"the number of audio features ({audio_embeds.shape[0] * audio_embeds.shape[1]}). "
+                    "Check your data preparation.")
+
+            inputs_embeds[audio_token_mask] = audio_embeds.flatten(0, 1).to(inputs_embeds.dtype)
+
+        outputs = self.language_model(
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            labels=labels,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            logits_to_keep=logits_to_keep,
+            return_dict=True,
+            **kwargs,
+        )
+        return outputs
+
+
+if __name__ == "__main__":
+    model = ManaraForConditionalGeneration()
+    print(model)
+    print("Audio Token ID:", model.audio_token_id)
+    print("Vocab Size:", model.vocab_size)
+    print("Text Config:", model.text_config)
+    print("Audio Config:", model.audio_config)
+    print("Model Parameters:", sum(p.numel() for p in model.parameters() if p.requires_grad))