Upload folder using huggingface_hub

modeling_finelap.py

@@ -1,9 +1,8 @@
-# modeling_finelap.py
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import torchaudio
 from transformers import PreTrainedModel, RobertaModel, RobertaTokenizer
-
 from .configuration_finelap import FineLAPConfig
 from .modeling_eat import EATModel 
 
@@ -13,163 +12,149 @@ class FineLAPModel(PreTrainedModel):
     def __init__(self, config: FineLAPConfig):
         super().__init__(config)
         self.config = config
-
         self.audio_encoder = EATModel(config.audio_config)
         self.audio_width = getattr(config.audio_config, 'hidden_size', 768) 
-
-        self.text_encoder = RobertaModel.from_pretrained(
-            config.text_encoder_name, 
-            add_pooling_layer=False,
-        )
         
+        self.text_encoder = RobertaModel.from_pretrained(config.text_encoder_name, add_pooling_layer=False)
         self.text_width = self.text_encoder.config.hidden_size
+        self.tokenizer = RobertaTokenizer.from_pretrained(config.text_encoder_name)
+        
         self.embed_size = config.embed_size
 
-        if config.temp_global != 0:
-            self.temp_global = nn.Parameter(torch.ones([]) * config.temp_global)
-        if config.b_global != 0:
-            self.b_global = nn.Parameter(torch.ones([]) * config.b_global)
-        if config.temp_local != 0:
-            self.temp_local = nn.Parameter(torch.ones([]) * config.temp_local)
-        if config.b_local != 0:
-            self.b_local = nn.Parameter(torch.ones([]) * config.b_local)
-
-        self.global_audio_proj = nn.Sequential(
-            nn.Linear(self.audio_width, self.embed_size),
-            nn.ReLU(),
-            nn.Linear(self.embed_size, self.embed_size),
-        )
-        self.global_text_proj = nn.Sequential(
-            nn.Linear(self.text_width, self.embed_size),
-            nn.ReLU(),
-            nn.Linear(self.embed_size, self.embed_size),
-        )
+        for param in ['temp_global', 'b_global', 'temp_local', 'b_local']:
+            val = getattr(config, param, None)
+            if val is not None:
+                self.register_parameter(param, nn.Parameter(torch.ones([]) * val))
+
+        self.global_audio_proj = nn.Sequential(nn.Linear(self.audio_width, self.embed_size), nn.ReLU(), nn.Linear(self.embed_size, self.embed_size))
+        self.global_text_proj = nn.Sequential(nn.Linear(self.text_width, self.embed_size), nn.ReLU(), nn.Linear(self.embed_size, self.embed_size))
 
-        # 5. Local Audio Projection Layer
         self.local_audio_proj_type = config.local_audio_proj_type
         if self.local_audio_proj_type == "rnn":
-            self.local_audio_proj = nn.GRU(
-                input_size=self.audio_width,
-                hidden_size=int(self.embed_size / 2),
-                num_layers=2,
-                batch_first=True,
-                bidirectional=True
-            )
-        elif self.local_audio_proj_type == "linear":
-            self.local_audio_proj = nn.Sequential(
-                nn.Linear(self.audio_width, self.embed_size),
-                nn.ReLU(),
-                nn.Linear(self.embed_size, self.embed_size)
-            )
+            self.local_audio_proj = nn.GRU(input_size=self.audio_width, hidden_size=int(self.embed_size / 2), num_layers=2, batch_first=True, bidirectional=True)
         elif self.local_audio_proj_type == "transformer":
-            encoder_layer = nn.TransformerEncoderLayer(
-                d_model=self.embed_size,
-                nhead=8,
-                dim_feedforward=self.embed_size * 4,
-                dropout=0.1,
-                activation='relu',
-                batch_first=True
-            )
-            transformer_encoder = nn.TransformerEncoder(encoder_layer=encoder_layer, num_layers=2)
-            self.local_audio_proj = nn.Sequential(
-                nn.Linear(self.audio_width, self.embed_size),
-                transformer_encoder
-            )
-        elif self.local_audio_proj_type == "transformer_linearlast":
-            encoder_layer = nn.TransformerEncoderLayer(
-                d_model=self.audio_width,
-                nhead=8,
-                dim_feedforward=self.audio_width * 4,
-                dropout=0.1,
-                activation='relu',
-                batch_first=True
-            )
-            transformer_encoder = nn.TransformerEncoder(encoder_layer=encoder_layer, num_layers=2)
-            self.local_audio_proj = nn.Sequential(
-                transformer_encoder,
-                nn.Linear(self.audio_width, self.embed_size),
-            )
-        else:
-            raise ValueError(f"Invalid local audio proj type: {self.local_audio_proj_type}")
-
+            l = nn.TransformerEncoderLayer(d_model=self.embed_size, nhead=8, dim_feedforward=self.embed_size * 4, dropout=0.1, activation='relu', batch_first=True)
+            self.local_audio_proj = nn.Sequential(nn.Linear(self.audio_width, self.embed_size), nn.TransformerEncoder(l, num_layers=2))
+        elif self.local_audio_proj_type == "linear":
+            self.local_audio_proj = nn.Sequential(nn.Linear(self.audio_width, self.embed_size), nn.ReLU(), nn.Linear(self.embed_size, self.embed_size))
+        
         self.post_init()
 
-
-    def encode_audio(self, audio_mel):
-
-        outputs = self.audio_encoder.extract_features(audio_mel)
-        audio_encoded_raw = outputs['x'] if isinstance(outputs, dict) else outputs
-
-        audio_cls = audio_encoded_raw[:, 0:1, :]
-        audio_patches = audio_encoded_raw[:, 1:, :]
+    def load_audio(self, audio_path, device=None):
+        device = device or self.device
+        wav, sr = torchaudio.load(audio_path)
+        if wav.shape[0] > 1:
+            wav = wav.mean(dim=0, keepdim=True)
+        if sr != 16000:
+            wav = torchaudio.functional.resample(wav, sr, 16000)
+        wav = wav.squeeze(0)
+        wav = wav - wav.mean()
         
-        B, T, D = audio_patches.shape
-        ds_factor = 8
-        audio_patches_downsampled = audio_patches.reshape(
-            B, T // ds_factor, ds_factor, D
-        ).mean(dim=2)
-        
-        # [B, 1+T//8, D]
-        audio_encoded = torch.cat([audio_cls, audio_patches_downsampled], dim=1)
-        return audio_encoded
-
-
-    def encode_text(self, input_ids, attention_mask):
-        outputs = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask)
-        return outputs.last_hidden_state
-
-
-    def get_global_text_embeds(self, input_ids, attention_mask):
-        text_feats = self.encode_text(input_ids, attention_mask)
-        text_embeds = F.normalize(self.global_text_proj(text_feats[:, 0, :]), dim=-1)
-        return text_embeds
-    
+        mel = torchaudio.compliance.kaldi.fbank(
+            wav.unsqueeze(0), htk_compat=True, sample_frequency=16000, 
+            use_energy=False, window_type='hanning', num_mel_bins=128, 
+            dither=0.0, frame_shift=10
+        )
+        target_len = 1024
+        if mel.shape[0] < target_len:
+            mel = F.pad(mel, (0, 0, 0, target_len - mel.shape[0]))
+        else:
+            mel = mel[:target_len, :]
+        mel = ((mel - (-4.268)) / (4.569 * 2)).unsqueeze(0).unsqueeze(0).to(device)
+        return mel
 
-    def get_global_audio_embeds(self, audio_mel):
-        audio_feats = self.encode_audio(audio_mel)
-        
+    def encode_audio(self, audio_path):
+        audio_mel = self.load_audio(audio_path)
+        outputs = self.audio_encoder.extract_features(audio_mel)
+        raw = outputs['x'] if isinstance(outputs, dict) else outputs
+        B, T, D = raw[:, 1:, :].shape
+        ds = 8
+        patches = raw[:, 1:, :].reshape(B, T // ds, ds, D).mean(dim=2)
+        return torch.cat([raw[:, 0:1, :], patches], dim=1)
+
+    def get_global_text_embeds(self, text_labels, device=None):
+        device = device or self.device
+        t_in = self.tokenizer(text_labels, padding=True, truncation=True, return_tensors="pt").to(device)
+        feat = self.text_encoder(input_ids=t_in["input_ids"], attention_mask=t_in["attention_mask"]).last_hidden_state
+        return F.normalize(self.global_text_proj(feat[:, 0, :]), dim=-1)
+
+    def get_global_audio_embeds(self, audio_path):
+        audio_feats = self.encode_audio(audio_path)
         if self.config.unify_audio_proj:
             audio_embeds = self.local_audio_proj(audio_feats)
             if self.config.local_audio_proj_type == "rnn":
                 audio_embeds = audio_embeds[0]
-            global_audio_embeds = F.normalize(audio_embeds[:, 0, :], dim=-1)
-            return global_audio_embeds
+            return F.normalize(audio_embeds[:, 0, :], dim=-1)
         else:
             audio_cls_feat = audio_feats[:, 0, :]
-            audio_embeds = F.normalize(self.global_audio_proj(audio_cls_feat), dim=-1)
-            return audio_embeds
-
-
-    def get_dense_audio_embeds(self, audio_mel):
-        audio_feats = self.encode_audio(audio_mel)
-        audio_patches = audio_feats[:, 1:, :]
-        
-        audio_embeds = self.local_audio_proj(audio_patches)
-        if self.config.local_audio_proj_type == "rnn":
-            audio_embeds = audio_embeds[0]
-            
-        if self.config.normalize_dense_audio_embeds:
-            audio_embeds = F.normalize(audio_embeds, dim=-1)
-        return audio_embeds
-
-    
-    def forward(self, audio_mel=None, input_ids=None, attention_mask=None, return_dict=True):
-        global_audio_embeds = None
-        dense_audio_embeds = None
-        global_text_embeds = None
-
-        if audio_mel is not None:
-            global_audio_embeds = self.get_global_audio_embeds(audio_mel)
-            dense_audio_embeds = self.get_dense_audio_embeds(audio_mel)
-            
-        if input_ids is not None:
-            global_text_embeds = self.get_global_text_embeds(input_ids, attention_mask)
-
-        if not return_dict:
-            return (global_audio_embeds, dense_audio_embeds, global_text_embeds)
-            
-        return {
-            "global_audio_embeds": global_audio_embeds,
-            "dense_audio_embeds": dense_audio_embeds,
-            "global_text_embeds": global_text_embeds
-        }
+            return F.normalize(self.global_audio_proj(audio_cls_feat), dim=-1)
+
+    def get_dense_audio_embeds(self, audio_path):
+        patches = self.encode_audio(audio_path)[:, 1:, :]
+        out = self.local_audio_proj(patches)
+        embeds = out[0] if self.local_audio_proj_type == "rnn" else out
+        return F.normalize(embeds, dim=-1) if self.config.normalize_dense_audio_embeds else embeds
+
+    @torch.no_grad()
+    def get_frame_level_score(self, audio_path, text_labels, device=None):
+        device = device or self.device
+        self.to(device)
+        self.eval()
+
+        dense_audio = self.get_dense_audio_embeds(audio_path).squeeze(0)
+        text_embeds = self.get_global_text_embeds(text_labels, device)
+
+        sim = torch.matmul(text_embeds, dense_audio.transpose(-1, -2))
+        if hasattr(self, "temp_local"):
+            sim = sim / self.temp_local
+        if hasattr(self, "b_local"):
+            sim = sim + self.b_local
+        return F.sigmoid(sim)
+
+    @torch.no_grad()
+    def get_clip_level_score(self, audio_path, text_labels, device=None):
+        device = device or self.device
+        self.to(device)
+        self.eval()
+
+        global_audio = self.get_global_audio_embeds(audio_path)
+        global_text = self.get_global_text_embeds(text_labels, device)
+
+        logits = torch.matmul(global_text, global_audio.transpose(-1, -2))
+        if hasattr(self, "temp_global"): 
+            logits = logits / self.temp_global
+        if hasattr(self, "b_global"): 
+            logits = logits + self.b_global
+        return torch.sigmoid(logits).squeeze(-1)
+
+    @torch.no_grad()
+    def plot_frame_level_score(self, audio_path, text_labels, output_path="similarity_plot.png", device=None):
+        import matplotlib.pyplot as plt
+        import numpy as np
+
+        scores = self.get_frame_level_score(audio_path, text_labels, device)
+        sim_matrix_np = scores.cpu().numpy()
+
+        fig, ax = plt.subplots(figsize=(14, 8))
+        im = ax.imshow(sim_matrix_np, aspect='auto', cmap='viridis', interpolation='nearest')
+        ax.set_xlabel('Time Frames', fontsize=12)
+        ax.set_ylabel('Labels', fontsize=12)
+        ax.set_title('Frame-level Audio-Text Similarity', fontsize=14)
+        ax.set_yticks(range(len(text_labels)))
+        ax.set_yticklabels(text_labels)
+
+        cbar = plt.colorbar(im, ax=ax)
+        cbar.set_label('Similarity Score', rotation=270, labelpad=20)
+
+        plt.tight_layout()
+        plt.savefig(output_path, dpi=150, bbox_inches='tight')
+        plt.close()
+
+    def forward(self, audio_path=None, text_labels=None):
+        res = {}
+        if audio_path is not None:
+            res["global_audio_embeds"] = self.get_global_audio_embeds(audio_path) if not self.config.unify_audio_proj else None
+            res["dense_audio_embeds"] = self.get_dense_audio_embeds(audio_path)
+        if text_labels is not None:
+            res["global_text_embeds"] = self.get_global_text_embeds(text_labels)
+        return res