Upload modeling_apex.py with huggingface_hub

modeling_apex.py

@@ -1,4 +1,3 @@
-import os
 import json
 import torch
 import torch.nn as nn
@@ -9,7 +8,9 @@ from transformers import PreTrainedModel, AutoProcessor, AutoModel
 from .configuration_apex import APEXConfig
 
 
-# Building blocks
+# -------------------------------
+# BUILDING BLOCKS
+# -------------------------------
 class SharedBlock(nn.Module):
     def __init__(self, in_dim, out_dim, dropout):
         super().__init__()
@@ -54,18 +55,22 @@ class TaskBranch(nn.Module):
         return torch.sigmoid(self.branch(x)) * self.scale + self.shift
 
 
+# -------------------------------
 # APEX MODEL
+# -------------------------------
 class APEXModel(PreTrainedModel):
-    config_class = APEXConfig
+    config_class                    = APEXConfig
+    _keys_to_ignore_on_load_missing = [r"mert\..*", r"mert_processor\..*"]
 
     def __init__(self, config: APEXConfig):
         super().__init__(config)
 
-        # MERT encoder + processor
+        # Load MERT processor and encoder fresh from HuggingFace
         self.mert_processor = AutoProcessor.from_pretrained(
-            config.mert_model_name, trust_remote_code=True
+            config.mert_model_name,
+            trust_remote_code = True
         )
-        with torch.device('cpu'):
+        with torch.device("cpu"):
             self.mert = AutoModel.from_pretrained(
                 config.mert_model_name,
                 trust_remote_code = True,
@@ -121,13 +126,11 @@ class APEXModel(PreTrainedModel):
         waveform, sr = sf.read(audio_path, dtype="float32")
         waveform     = torch.from_numpy(waveform)
 
-        # Stereo to mono
         if len(waveform.shape) > 1 and waveform.shape[1] > 1:
             waveform = waveform.mean(dim=1)
 
         waveform = waveform.to(self.device)
 
-        # Resample if needed
         if sr != self.target_sr:
             waveform = TAF.resample(waveform, sr, self.target_sr)
 
@@ -142,12 +145,10 @@ class APEXModel(PreTrainedModel):
             if segment.numel() == 0:
                 break
 
-            # Zero-pad if needed
             if segment.shape[0] < segment_len:
                 pad_len = segment_len - segment.shape[0]
                 segment = torch.nn.functional.pad(segment, (0, pad_len))
 
-            # MERT forward
             inputs = self.mert_processor(
                 segment.cpu().numpy(),
                 sampling_rate  = self.target_sr,
@@ -158,23 +159,20 @@ class APEXModel(PreTrainedModel):
             with torch.no_grad():
                 outputs = self.mert(**inputs, output_hidden_states=True)
 
-            # Extract layers and aggregate
             all_hidden = torch.stack([
                 outputs.hidden_states[i].mean(dim=1)
                 for i in self.config.layer_indices
-            ])  # [4, 1, 768]
-            all_hidden = all_hidden.squeeze(1)  # [4, 768]
+            ])
+            all_hidden = all_hidden.squeeze(1)
 
-            # Conv1d aggregation
             pooled = self.aggregator(
                 all_hidden.unsqueeze(0)
-            ).squeeze()  # [768]
+            ).squeeze()
 
             segment_embeddings.append(pooled)
 
             del segment, inputs, outputs, all_hidden, pooled
 
-        # Average across segments to song-level embedding
         song_embedding = torch.stack(segment_embeddings).mean(dim=0)
         return song_embedding