Update tasks/audio.py

tasks/audio.py

@@ -5,6 +5,9 @@ from sklearn.metrics import accuracy_score
 import random
 import os
 import torch
+from torch.utils.data import DataLoader
+from Model_Loader import load_model  
+
 
 from .utils.evaluation import AudioEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
@@ -14,7 +17,7 @@ load_dotenv()
 
 router = APIRouter()
 
-DESCRIPTION = "Random Baseline"
+DESCRIPTION = "Quantized M5"
 ROUTE = "/audio"
 
 
@@ -55,30 +58,46 @@ async def evaluate_audio(request: AudioEvaluationRequest):
     #--------------------------------------------------------------------------------------------   
     
     # Make random predictions (placeholder for actual model inference)
-    def preprocess_audio(example):
-    """Convert dataset into tensors."""
-        waveform = torch.tensor(example["audio"]["array"], dtype=torch.float32).unsqueeze(0)  # Add batch dim
-        label = torch.tensor(example["label"], dtype=torch.long)  # Ensure labels are `int64`
-    return waveform, label
-
     model_path = "quantized_teacher_m5_static.pth"
     model, device = load_model(model_path)
 
+    def preprocess_audio(example, target_length=32000):
+            """
+        Convert dataset into tensors:
+        - Convert to tensor
+        - Normalize waveform
+        - Pad/truncate to `target_length`
+        """
+        waveform = torch.tensor(example["audio"]["array"], dtype=torch.float32).unsqueeze(0)  # Add batch dim
+    
+        # Normalize waveform
+        waveform = (waveform - waveform.mean()) / (waveform.std() + 1e-6)
+    
+        # Pad or truncate to fixed length
+        if waveform.shape[1] < target_length:
+            pad = torch.zeros(1, target_length - waveform.shape[1])
+            waveform = torch.cat((waveform, pad), dim=1)  # Pad
+        else:
+            waveform = waveform[:, :target_length]  # Truncate
+    
+        label = torch.tensor(example["label"], dtype=torch.long)  # Ensure int64
+    return {"waveform": waveform, "label": label}
+
+
 
-    train_test = train_test.map(preprocess_audio)
+    train_test = train_test.map(preprocess_audio, batched=True)
     test_dataset = train_test.map(preprocess_audio)
 
-    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=True)
+    train_loader = DataLoader(train_test, batch_size=32, shuffle=True)
 
     
-    true_labels = test_dataset["label"]
+    true_labels = train_dataset["label"]
     predictions = []
 
     with torch.no_grad():
-    for waveforms, labels in test_loader:
+    for waveforms, labels in train_loader:
         waveforms, labels = waveforms.to(device), labels.to(device)
         
-        # Run Model
         outputs = model(waveforms)
         predicted_label = torch.argmax(F.softmax(outputs, dim=1), dim=1)
         true_labels.extend(labels.cpu().numpy())