Mayank022/Audio-Language-Model

1 2import torch3import torch.nn as nn4import transformers5from typing import Optional, Tuple, Union, List6from config import ModelConfig7 8class ModelProjector(nn.Module):9    def __init__(self, config: ModelConfig, audio_hidden_size: int):10        super().__init__()11        self.stack_factor = config.stack_factor12        input_dim = audio_hidden_size * self.stack_factor13        14        self.linear1 = nn.Linear(input_dim, config.hidden_size)15        self.act = nn.GELU() if config.projector_act == 'gelu' else nn.ReLU()16        self.linear2 = nn.Linear(config.hidden_size, config.hidden_size)17        self.norm = nn.LayerNorm(config.hidden_size)18    19    def forward(self, audio_features: torch.Tensor) -> torch.Tensor:20        if audio_features.dim() == 3 and audio_features.shape[1] < audio_features.shape[2]:21             audio_features = audio_features.transpose(1, 2)22        23        B, T, C = audio_features.shape24        25        if T % self.stack_factor != 0:26            pad_len = self.stack_factor - (T % self.stack_factor)27            audio_features = torch.nn.functional.pad(audio_features, (0, 0, 0, pad_len))28            T = T + pad_len29            30        audio_features = audio_features.view(B, T // self.stack_factor, C * self.stack_factor)31        32        x = self.linear1(audio_features)33        x = self.act(x)34        x = self.linear2(x)35        x = self.norm(x)36        return x37 38class MultiModalModel(nn.Module):39    def __init__(self, config: ModelConfig):40        super().__init__()41        self.config = config42        43        self.audio_encoder = transformers.AutoModel.from_pretrained(config.audio_model_id).encoder44        for param in self.audio_encoder.parameters():45            param.requires_grad = False46            47        audio_hidden_size = self.audio_encoder.config.hidden_size48        49        self.llm = transformers.AutoModelForCausalLM.from_pretrained(config.text_model_id, trust_remote_code=True)50        self.llm_hidden_size = self.llm.config.hidden_size51        52        self.projector = ModelProjector(config, audio_hidden_size)53        if config.hidden_size != self.llm_hidden_size:54             self.projector.linear2 = nn.Linear(config.hidden_size, self.llm_hidden_size)55             self.projector.norm = nn.LayerNorm(self.llm_hidden_size)56 57        58    def forward(59        self, 60        input_ids: torch.Tensor, 61        audio_values: Optional[torch.Tensor] = None, 62        labels: Optional[torch.Tensor] = None,63        **kwargs64    ):65        inputs_embeds = self.llm.get_input_embeddings()(input_ids)66        67        if audio_values is not None:68             audio_outputs = self.audio_encoder(audio_values)69             audio_features = audio_outputs.last_hidden_state70             71             audio_projected = self.projector(audio_features)72             73             inputs_embeds = torch.cat([audio_projected, inputs_embeds], dim=1)74             75             if labels is not None:76                  audio_labels = torch.full((audio_projected.shape[0], audio_projected.shape[1]), -100, device=labels.device, dtype=labels.dtype)77                  labels = torch.cat([audio_labels, labels], dim=1)78                  79             if "attention_mask" in kwargs:80                  audio_mask = torch.ones((audio_projected.shape[0], audio_projected.shape[1]), device=inputs_embeds.device, dtype=kwargs["attention_mask"].dtype)81                  kwargs["attention_mask"] = torch.cat([audio_mask, kwargs["attention_mask"]], dim=1)82 83        # Match LLM dtype (e.g. bfloat16) to avoid "float != bfloat16" in linear layers84        llm_dtype = next(self.llm.parameters()).dtype85        inputs_embeds = inputs_embeds.to(llm_dtype)86        if labels is not None:87            labels = labels.to(llm_dtype) if labels.dtype.is_floating_point else labels88 89        # Drop non-tensor keys (e.g. continuation) so LLM forward doesn't receive them90        kwargs = {k: v for k, v in kwargs.items() if isinstance(v, torch.Tensor)}91        outputs = self.llm(92            inputs_embeds=inputs_embeds,93            labels=labels,94            **kwargs95        )96        97        return outputs98    99    def generate(self, input_ids, audio_values=None, **kwargs):100        inputs_embeds = self.llm.get_input_embeddings()(input_ids)101        102        if audio_values is not None:103             audio_outputs = self.audio_encoder(audio_values)104             audio_features = audio_outputs.last_hidden_state105             audio_projected = self.projector(audio_features)106             inputs_embeds = torch.cat([audio_projected, inputs_embeds], dim=1)107             108             if "attention_mask" in kwargs:109                  audio_mask = torch.ones((audio_projected.shape[0], audio_projected.shape[1]), device=inputs_embeds.device, dtype=kwargs["attention_mask"].dtype)110                  kwargs["attention_mask"] = torch.cat([audio_mask, kwargs["attention_mask"]], dim=1)111             inputs_embeds = inputs_embeds.to(next(self.llm.parameters()).dtype)112 113        return self.llm.generate(inputs_embeds=inputs_embeds, **kwargs)114 115    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):116        self.llm.gradient_checkpointing_enable(gradient_checkpointing_kwargs=gradient_checkpointing_kwargs)117