✏️ Add attention_mask in SpeechRecognitionDataset

hezarai · May 6, 2024 · b8a6479 · b8a6479
1 parent 4b5f446
commit b8a6479
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/hezar/data/datasets/speech_recognition_dataset.py b/hezar/data/datasets/speech_recognition_dataset.py
@@ -73,10 +73,11 @@ def __getitem__(self, index):
         labels = self.tokenizer(
             transcript,
             max_length=self.config.labels_max_length,
-            return_tensors="pt"
-        )["token_ids"]
+            return_tensors="pt",
+        )
 
         return {
             "input_features": input_features,
-            "labels": labels
+            "labels": labels["token_ids"],
+            "attention_mask": labels["attention_mask"],
         }