new batching

model.py

@@ -5,7 +5,6 @@ from transformers import AutoModel, AutoConfig, AutoFeatureExtractor
 import torchaudio
 from safetensors import safe_open
 from typing import List, Dict
-import time
 
 torch.backends.cuda.matmul.allow_tf32 = True 
 torch.backends.cuda.enable_flash_sdp(True)
@@ -66,11 +65,9 @@ class WavLMForMusicDetection(nn.Module):
     ) -> torch.Tensor:
         """
         Apply attention-based pooling over time dimension.
-
         Args:
             hidden_states (torch.Tensor): [batch_size, seq_len, hidden_size]
             attention_mask (torch.Tensor): [batch_size, seq_len] — mask to ignore padding
-
         Returns:
             torch.Tensor: [batch_size, hidden_size] — context vector
         """
@@ -94,21 +91,22 @@ class WavLMForMusicDetection(nn.Module):
     ) -> torch.Tensor:
         """
         Forward pass for inference.
-
         Args:
             input_values (torch.Tensor): [batch_size, audio_seq_len] — raw audio waveform
             attention_mask (torch.Tensor): [batch_size, audio_seq_len] — input mask (1 = real, 0 = pad)
-
         Returns:
             torch.Tensor: [batch_size, 1] — probability that audio contains music
         """
         assert isinstance(input_values, torch.Tensor), f"Expected torch.Tensor, got {type(input_values)}"
         assert isinstance(attention_mask, torch.Tensor), f"Expected torch.Tensor, got {type(attention_mask)}"
 
-        outputs = self.wavlm(input_values.to(self.device), attention_mask=attention_mask.to(self.device))
+
+        input_values = input_values.to(dtype=self.dtype, device=self.device)
+        attention_mask = attention_mask.to(device=self.device, dtype=self.dtype)
+
+        outputs = self.wavlm(input_values, attention_mask=attention_mask)
         hidden_states = outputs.last_hidden_state  # [B, T', D]
 
-        # Align attention mask with downsampled hidden states
         input_length = attention_mask.size(1)
         hidden_length = hidden_states.size(1)
         ratio = input_length / hidden_length
@@ -125,10 +123,8 @@ class WavLMForMusicDetection(nn.Module):
     def _prepare_batches(self, audio_paths: List[str]) -> List[List[str]]:
         """
         Split list of audio paths into batches of size `self.batch_size`.
-
         Args:
             audio_paths (List[str]): List of paths to audio files.
-
         Returns:
             List[List[str]]: List of batches, each batch is a list of paths.
         """
@@ -151,10 +147,8 @@ class WavLMForMusicDetection(nn.Module):
     def _preprocess_audio_batch(self, audio_paths: List[str]) -> Dict[str, torch.Tensor]:
         """
         Load and preprocess a batch of audio files.
-
         Args:
             audio_paths (List[str]): List of file paths.
-
         Returns:
             Dict with keys:
                 "input_values": tensor [B, T]
@@ -191,10 +185,8 @@ class WavLMForMusicDetection(nn.Module):
     def predict_proba(self, audio_paths: List[str]) -> torch.Tensor:
         """
         Predict music probability for a list of audio files.
-
         Args:
             audio_paths (List[str]): List of audio file paths.
-
         Returns:
             torch.Tensor: [N] — probabilities for each audio file.
         """
@@ -212,21 +204,68 @@ class WavLMForMusicDetection(nn.Module):
             all_probs.append(probs)
 
         return torch.cat(all_probs, dim=0)
+
+    def convert_to_bf16(self):
+        self.wavlm = self.wavlm.to(torch.bfloat16)
+        self.pool_attention = self.pool_attention.to(torch.bfloat16)
+        self.classifier = self.classifier.to(torch.bfloat16)
+        self.dtype = torch.bfloat16
+        return self
+    
+    def predict_proba_smart_batching(
+        self,
+        audio_paths: List[str],
+        audio_lengths: List[float]
+    ) -> torch.Tensor:
+      
+        assert len(audio_paths) == len(audio_lengths), \
+            f"Mismatch: {len(audio_paths)} paths vs {len(audio_lengths)} lengths"
+        
+        was_training = self.training
+        self.eval()
+        
+        try:
+            indexed_audios = [
+                (i, path, length)
+                for i, (path, length) in enumerate(zip(audio_paths, audio_lengths))
+            ]
+            
+            sorted_audios = sorted(indexed_audios, key=lambda x: x[2])
+            batches = []
+            for i in range(0, len(sorted_audios), self.batch_size):
+                batch = sorted_audios[i:i + self.batch_size]
+                batches.append(batch)
+            
+            results = {}
+            
+            for batch in batches:
+                batch_paths = [item[1] for item in batch]
+                batch_indices = [item[0] for item in batch]
+                
+                inputs = self._preprocess_audio_batch(batch_paths)
+                inputs = {k: v.to(self.device) for k, v in inputs.items()}
+                
+                with torch.no_grad():
+                    probs = self.forward(**inputs).squeeze(-1)
+                
+                if probs.dim() == 0:
+                    probs = probs.unsqueeze(0)
+                
+                for idx, prob in zip(batch_indices, probs):
+                    results[idx] = prob.cpu()
+            
+            all_probs = [results[i] for i in range(len(audio_paths))]
+            return torch.stack(all_probs)
+        finally:
+            if was_training:
+                self.train()
     
 if __name__ == "__main__":
     device = 'cuda:0'
     checkpoint_path = './music_detection.safetensors'
-    model = WavLMForMusicDetection('microsoft/wavlm-base-plus', batch_size=32, device=device)
-    
+    model = WavLMForMusicDetection('microsoft/wavlm-base-plus', batch_size=8, device=device)
+    model.convert_to_bf16()
+    model.eval()
     with safe_open(checkpoint_path, framework="pt", device=device) as f:
         state_dict = {key: f.get_tensor(key) for key in f.keys()}
     model.load_state_dict(state_dict)
-    global_start = time.time()
-    paths = [
-        '/92.mp3',
-        '133.mp3',
-        '113.mp3',
-        '30.mp3'
-    ]
-    print(model.predict_proba(paths))
-