Spaces:
Sleeping
Sleeping
| from typing import Any, Self | |
| import torch | |
| from torch import nn | |
| from transformers import PreTrainedModel | |
| from transformers.modeling_outputs import CausalLMOutputWithPast | |
| from src.modeling.adapter import AudioAdapter | |
| from src.modeling.cross_attention import GatedCrossAttention | |
| from src.modeling.lfm2_decoder import AudioLfm2ForCausalLM, add_cross_attention | |
| class LambdaAudioModel(nn.Module): | |
| def __init__( | |
| self, | |
| encoder: PreTrainedModel, | |
| decoder: AudioLfm2ForCausalLM, | |
| encoder_hidden_size: int = 768, | |
| ) -> None: | |
| super().__init__() | |
| self.backbones_trainable = True | |
| self.encoder = encoder | |
| self.decoder = decoder | |
| self.adapter = AudioAdapter(encoder_hidden_size, decoder.config.hidden_size) | |
| add_cross_attention(self.decoder) | |
| def encode( | |
| self, | |
| input_values: torch.Tensor, | |
| audio_padding_mask: torch.BoolTensor, | |
| ) -> tuple[torch.Tensor, torch.BoolTensor]: | |
| # --------------------------------------------------------- | |
| # Encode audio and create the mask for encoder output frames. | |
| # --------------------------------------------------------- | |
| encoder_output = self.encoder( | |
| input_values=input_values, | |
| padding_mask=audio_padding_mask, | |
| ) | |
| encoder_padding_mask = self.encoder.forward_padding_mask( | |
| encoder_output.last_hidden_state, | |
| audio_padding_mask, | |
| ) | |
| return self.adapter( | |
| encoder_output.last_hidden_state, | |
| encoder_padding_mask, | |
| ) | |
| def forward( | |
| self, | |
| input_values: torch.Tensor, | |
| audio_padding_mask: torch.BoolTensor, | |
| input_ids: torch.LongTensor, | |
| attention_mask: torch.Tensor, | |
| labels: torch.LongTensor, | |
| ) -> CausalLMOutputWithPast: | |
| audio_features, encoded_padding_mask = self.encode( | |
| input_values, | |
| audio_padding_mask, | |
| ) | |
| audio_key_values = self.decoder.project_audio_key_values(audio_features) | |
| return self.decoder( | |
| input_ids=input_ids, | |
| attention_mask=attention_mask, | |
| labels=labels, | |
| audio_key_values=audio_key_values, | |
| audio_padding_mask=encoded_padding_mask, | |
| ) | |
| def generate( | |
| self, | |
| input_values: torch.Tensor, | |
| audio_padding_mask: torch.BoolTensor, | |
| input_ids: torch.LongTensor, | |
| attention_mask: torch.Tensor, | |
| **kwargs: Any, | |
| ) -> torch.LongTensor: | |
| audio_features, encoded_padding_mask = self.encode( | |
| input_values, | |
| audio_padding_mask, | |
| ) | |
| audio_key_values = self.decoder.project_audio_key_values(audio_features) | |
| return self.decoder.generate( | |
| input_ids=input_ids, | |
| attention_mask=attention_mask, | |
| audio_key_values=audio_key_values, | |
| audio_padding_mask=encoded_padding_mask, | |
| **kwargs, | |
| ) | |
| def set_backbones_trainable(self, trainable: bool) -> None: | |
| # --------------------------------------------------------- | |
| # Keep audio conditioning trainable while freezing pretrained weights. | |
| # --------------------------------------------------------- | |
| self.backbones_trainable = trainable | |
| self.requires_grad_(trainable) | |
| self.adapter.requires_grad_(True) | |
| for module in self.decoder.modules(): | |
| if isinstance(module, GatedCrossAttention): | |
| module.requires_grad_(True) | |
| self.train(self.training) | |
| def train(self, mode: bool = True) -> Self: | |
| super().train(mode) | |
| if mode and not self.backbones_trainable: | |
| self.encoder.eval() | |
| self.decoder.eval() | |
| for module in self.decoder.modules(): | |
| if isinstance(module, GatedCrossAttention): | |
| module.train() | |
| return self | |