Update dataset.py

dataset.py

@@ -6,21 +6,17 @@ import librosa
 import numpy as np
 
 class DemoDataset(Dataset):
-    def __init__(self, demonstration_paths, query_path, sample_rate=16000):
+    def __init__(self, demonstration_paths, demonstration_labels, query_path, sample_rate=16000):
         self.sample_rate = sample_rate
         self.query_path = query_path
         
         # Convert to list if single path
-        if isinstance(demonstration_paths, str):
-            self.demonstration_paths = [demonstration_paths]
-        else:
-            self.demonstration_paths = demonstration_paths
-            
+        self.demonstration_paths = demonstration_paths
+        self.demonstration_labels = [0 if label == 'bonafide' else 1 for label in demonstration_labels]
+
         # Load feature extractor
-        self.feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base")
+        self.feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/w2v-bert-2.0")
         
-        print(f'Number of demonstration audios: {len(self.demonstration_paths)}')
-        print(f'Query audio: {self.query_path}')
 
     def load_pad(self, path, max_length=64000):
         """Load and pad audio file"""
@@ -74,9 +70,12 @@ class DemoDataset(Dataset):
             )
             prompt_features.append(prompt_feature)
         
+        prompt_labels = torch.tensor(self.demonstration_labels, dtype=torch.long)
+        
         return {
             'main_features': main_features,
             'prompt_features': prompt_features,
+            'prompt_labels': prompt_labels,
             'file_name': os.path.basename(self.query_path),
             'file_path': self.query_path
         }
@@ -114,9 +113,12 @@ def collate_fn(batch):
     file_names = [item['file_name'] for item in batch]
     file_paths = [item['file_path'] for item in batch]
 
+    prompt_labels = torch.tensor([item['prompt_labels'] for item in batch], dtype=torch.long)
+
     return {
         'main_features': main_features,
         'prompt_features': prompt_features,
+        'prompt_labels': prompt_labels,
         'file_names': file_names,
         'file_paths': file_paths
     }