Update README.md

Update of the example inference code

README.md

@@ -26,6 +26,8 @@ It achieves the following results on the evaluation set:
 
 ```python
 import os
+import random
+from glob import glob
 from typing import List, Optional, Union, Dict
 
 import tqdm
@@ -42,7 +44,6 @@ from transformers import (
     Wav2Vec2Processor
 )
 
-
 class CustomDataset(torch.utils.data.Dataset):
     def __init__(
         self,
@@ -68,19 +69,19 @@ class CustomDataset(torch.utils.data.Dataset):
             filepath = self.dataset[index]
         else:
             filepath = os.path.join(self.basedir, self.dataset[index])
-    
+
         speech_array, sr = torchaudio.load(filepath)
-    
+
         if speech_array.shape[0] > 1:
             speech_array = torch.mean(speech_array, dim=0, keepdim=True)
-    
+
         if sr != self.sampling_rate:
             transform = torchaudio.transforms.Resample(sr, self.sampling_rate)
             speech_array = transform(speech_array)
             sr = self.sampling_rate
-    
+
         len_audio = speech_array.shape[1]
-    
+
         # Pad or truncate the audio to match the desired length
         if len_audio < self.max_audio_len * self.sampling_rate:
             # Pad the audio if it's shorter than the desired length
@@ -89,9 +90,9 @@ class CustomDataset(torch.utils.data.Dataset):
         else:
             # Truncate the audio if it's longer than the desired length
             speech_array = speech_array[:, :self.max_audio_len * self.sampling_rate]
-    
+
         speech_array = speech_array.squeeze().numpy()
-    
+
         return {"input_values": speech_array, "attention_mask": None}
 
 
@@ -99,34 +100,37 @@ class CollateFunc:
     def __init__(
         self,
         processor: Wav2Vec2Processor,
-        max_length: Optional[int] = None,
         padding: Union[bool, str] = True,
         pad_to_multiple_of: Optional[int] = None,
+        return_attention_mask: bool = True,
         sampling_rate: int = 16000,
+        max_length: Optional[int] = None,
     ):
-        self.padding = padding
-        self.processor = processor
-        self.max_length = max_length
         self.sampling_rate = sampling_rate
+        self.processor = processor
+        self.padding = padding
         self.pad_to_multiple_of = pad_to_multiple_of
+        self.return_attention_mask = return_attention_mask
+        self.max_length = max_length
 
-    def __call__(self, batch: List):
-        input_features = []
-
-        for audio in batch:
-            input_tensor = self.processor(audio, sampling_rate=self.sampling_rate).input_values
-            input_tensor = np.squeeze(input_tensor)
-            input_features.append({"input_values": input_tensor})
+    def __call__(self, batch: List[Dict[str, np.ndarray]]):
+        # Extract input_values from the batch
+        input_values = [item["input_values"] for item in batch]
 
-        batch = self.processor.pad(
-            input_features,
+        batch = self.processor(
+            input_values,
+            sampling_rate=self.sampling_rate,
+            return_tensors="pt",
             padding=self.padding,
             max_length=self.max_length,
             pad_to_multiple_of=self.pad_to_multiple_of,
-            return_tensors="pt",
+            return_attention_mask=self.return_attention_mask
         )
 
-        return batch
+        return {
+            "input_values": batch.input_values,
+            "attention_mask": batch.attention_mask if self.return_attention_mask else None
+        }
 
 
 def predict(test_dataloader, model, device: torch.device):
@@ -175,15 +179,15 @@ def get_gender(model_name_or_path: str, audio_paths: List[str], label2id: Dict,
         batch_size=16,
         collate_fn=data_collator,
         shuffle=False,
-        num_workers=10
+        num_workers=2
     )
 
     preds = predict(test_dataloader=test_dataloader, model=model, device=device)
 
     return preds
 
-
 model_name_or_path = "alefiury/wav2vec2-large-xlsr-53-gender-recognition-librispeech"
+
 audio_paths = [] # Must be a list with absolute paths of the audios that will be used in inference
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")