Update model.py

model.py

@@ -0,0 +1,116 @@
+
+import torch
+import torch.nn as nn
+import transformers
+from typing import Optional, Tuple, Union, List
+from config import ModelConfig
+
+class ModelProjector(nn.Module):
+    def __init__(self, config: ModelConfig, audio_hidden_size: int):
+        super().__init__()
+        self.stack_factor = config.stack_factor
+        input_dim = audio_hidden_size * self.stack_factor
+        
+        self.linear1 = nn.Linear(input_dim, config.hidden_size)
+        self.act = nn.GELU() if config.projector_act == 'gelu' else nn.ReLU()
+        self.linear2 = nn.Linear(config.hidden_size, config.hidden_size)
+        self.norm = nn.LayerNorm(config.hidden_size)
+    
+    def forward(self, audio_features: torch.Tensor) -> torch.Tensor:
+        if audio_features.dim() == 3 and audio_features.shape[1] < audio_features.shape[2]:
+             audio_features = audio_features.transpose(1, 2)
+        
+        B, T, C = audio_features.shape
+        
+        if T % self.stack_factor != 0:
+            pad_len = self.stack_factor - (T % self.stack_factor)
+            audio_features = torch.nn.functional.pad(audio_features, (0, 0, 0, pad_len))
+            T = T + pad_len
+            
+        audio_features = audio_features.view(B, T // self.stack_factor, C * self.stack_factor)
+        
+        x = self.linear1(audio_features)
+        x = self.act(x)
+        x = self.linear2(x)
+        x = self.norm(x)
+        return x
+
+class MultiModalModel(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+        
+        self.audio_encoder = transformers.AutoModel.from_pretrained(config.audio_model_id).encoder
+        for param in self.audio_encoder.parameters():
+            param.requires_grad = False
+            
+        audio_hidden_size = self.audio_encoder.config.hidden_size
+        
+        self.llm = transformers.AutoModelForCausalLM.from_pretrained(config.text_model_id, trust_remote_code=True)
+        self.llm_hidden_size = self.llm.config.hidden_size
+        
+        self.projector = ModelProjector(config, audio_hidden_size)
+        if config.hidden_size != self.llm_hidden_size:
+             self.projector.linear2 = nn.Linear(config.hidden_size, self.llm_hidden_size)
+             self.projector.norm = nn.LayerNorm(self.llm_hidden_size)
+
+        
+    def forward(
+        self, 
+        input_ids: torch.Tensor, 
+        audio_values: Optional[torch.Tensor] = None, 
+        labels: Optional[torch.Tensor] = None,
+        **kwargs
+    ):
+        inputs_embeds = self.llm.get_input_embeddings()(input_ids)
+        
+        if audio_values is not None:
+             audio_outputs = self.audio_encoder(audio_values)
+             audio_features = audio_outputs.last_hidden_state
+             
+             audio_projected = self.projector(audio_features)
+             
+             inputs_embeds = torch.cat([audio_projected, inputs_embeds], dim=1)
+             
+             if labels is not None:
+                  audio_labels = torch.full((audio_projected.shape[0], audio_projected.shape[1]), -100, device=labels.device, dtype=labels.dtype)
+                  labels = torch.cat([audio_labels, labels], dim=1)
+                  
+             if "attention_mask" in kwargs:
+                  audio_mask = torch.ones((audio_projected.shape[0], audio_projected.shape[1]), device=inputs_embeds.device, dtype=kwargs["attention_mask"].dtype)
+                  kwargs["attention_mask"] = torch.cat([audio_mask, kwargs["attention_mask"]], dim=1)
+
+        # Match LLM dtype (e.g. bfloat16) to avoid "float != bfloat16" in linear layers
+        llm_dtype = next(self.llm.parameters()).dtype
+        inputs_embeds = inputs_embeds.to(llm_dtype)
+        if labels is not None:
+            labels = labels.to(llm_dtype) if labels.dtype.is_floating_point else labels
+
+        # Drop non-tensor keys (e.g. continuation) so LLM forward doesn't receive them
+        kwargs = {k: v for k, v in kwargs.items() if isinstance(v, torch.Tensor)}
+        outputs = self.llm(
+            inputs_embeds=inputs_embeds,
+            labels=labels,
+            **kwargs
+        )
+        
+        return outputs
+    
+    def generate(self, input_ids, audio_values=None, **kwargs):
+        inputs_embeds = self.llm.get_input_embeddings()(input_ids)
+        
+        if audio_values is not None:
+             audio_outputs = self.audio_encoder(audio_values)
+             audio_features = audio_outputs.last_hidden_state
+             audio_projected = self.projector(audio_features)
+             inputs_embeds = torch.cat([audio_projected, inputs_embeds], dim=1)
+             
+             if "attention_mask" in kwargs:
+                  audio_mask = torch.ones((audio_projected.shape[0], audio_projected.shape[1]), device=inputs_embeds.device, dtype=kwargs["attention_mask"].dtype)
+                  kwargs["attention_mask"] = torch.cat([audio_mask, kwargs["attention_mask"]], dim=1)
+             inputs_embeds = inputs_embeds.to(next(self.llm.parameters()).dtype)
+
+        return self.llm.generate(inputs_embeds=inputs_embeds, **kwargs)
+
+    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
+        self.llm.gradient_checkpointing_enable(gradient_checkpointing_kwargs=gradient_checkpointing_kwargs)