Fix audio decoding: use dataset.decode() to force decode all audio before expansion

training/augmentation.py

@@ -307,45 +307,67 @@ def expand_dataset_with_speed_augmentation(
     print(f"  Original size: {len(dataset):,} samples")
     print(f"  Speed factors: {speed_factors}")
     
-    # Use indexed access to ensure audio decoding happens
-    for idx in range(len(dataset)):
-        example = dataset[idx]
+    # Force decode all audio files before iterating
+    # This ensures AudioDecoder objects are decoded to dict format with 'array' and 'sampling_rate'
+    print("  Decoding audio files...")
+    try:
+        # Use decode() to force decoding of all audio files
+        # This converts AudioDecoder objects to dict format
+        dataset_decoded = dataset.decode()
+        print(f"  ✓ Audio decoding complete")
+    except Exception as e:
+        print(f"  ⚠ Warning: Failed to decode dataset: {e}")
+        print(f"  Continuing with undecoded dataset (may be slower)...")
+        dataset_decoded = dataset
+    
+    # Use indexed access to iterate over decoded dataset
+    for idx in range(len(dataset_decoded)):
+        example = dataset_decoded[idx]
         original_id = example.get(id_column, f"sample_{idx}")
+        
+        # After decode(), audio should be a dict with 'array' and 'sampling_rate'
         audio_data = example.get(audio_column)
         transcription = example.get(transcription_column, "")
         
-        # Handle audio data format - decode if needed
+        # Handle audio data format
+        # After decode(), audio_data should be a dict with 'array' and 'sampling_rate'
         audio_array = None
         audio_sr = target_sr
         
         if audio_data is not None:
-            # Handle dict format (most common - audio is decoded when accessed via dataset[idx])
             if isinstance(audio_data, dict):
                 audio_array = audio_data.get("array")
                 audio_sr = audio_data.get("sampling_rate", target_sr)
-            # Handle Audio object with .array attribute
+            # Fallback: handle Audio object with .array attribute (in case decode() didn't work)
             elif hasattr(audio_data, "array"):
                 audio_array = audio_data.array
                 audio_sr = getattr(audio_data, "sampling_rate", target_sr)
-            # Handle if it's already a numpy array
+            # Fallback: handle numpy array directly
             elif isinstance(audio_data, np.ndarray):
                 audio_array = audio_data
-            # If it's still an AudioDecoder, try to access it properly
+            # If still an AudioDecoder, try to decode it manually
             elif hasattr(audio_data, "__class__") and "Decoder" in str(type(audio_data)):
-                # Try to get the decoded value - sometimes accessing via dict key helps
                 try:
-                    # Force decoding by accessing the audio column again with proper indexing
-                    decoded_example = dataset[idx]
-                    audio_data = decoded_example.get(audio_column)
-                    if isinstance(audio_data, dict):
-                        audio_array = audio_data.get("array")
-                        audio_sr = audio_data.get("sampling_rate", target_sr)
-                    elif hasattr(audio_data, "array"):
-                        audio_array = audio_data.array
-                        audio_sr = getattr(audio_data, "sampling_rate", target_sr)
+                    if hasattr(audio_data, "decode"):
+                        decoded = audio_data.decode()
+                        if isinstance(decoded, dict):
+                            audio_array = decoded.get("array")
+                            audio_sr = decoded.get("sampling_rate", target_sr)
+                        elif hasattr(decoded, "array"):
+                            audio_array = decoded.array
+                            audio_sr = getattr(decoded, "sampling_rate", target_sr)
                 except Exception as e:
-                    print(f"⚠ Warning: Failed to decode audio for sample {original_id}: {e}, skipping...")
+                    print(f"⚠ Warning: Failed to decode AudioDecoder for sample {original_id}: {e}, skipping...")
+                    skipped_count += 1
+                    if skipped_count <= 5:
+                        print(f"  Audio data type: {type(audio_data)}")
                     continue
+            else:
+                print(f"⚠ Warning: Unexpected audio format for sample {original_id}: {type(audio_data)}, skipping...")
+                skipped_count += 1
+                if skipped_count <= 5:
+                    print(f"  Audio data type: {type(audio_data)}")
+                continue
         
         if audio_array is None:
             skipped_count += 1
@@ -407,7 +429,7 @@ def expand_dataset_with_speed_augmentation(
         
         # Progress update every 1000 samples
         if (idx + 1) % 1000 == 0:
-            print(f"  Processed {idx + 1:,}/{len(dataset):,} samples...")
+            print(f"  Processed {idx + 1:,}/{len(dataset_decoded):,} samples...")
     
     if len(expanded_examples) == 0:
         raise ValueError(
@@ -423,7 +445,7 @@ def expand_dataset_with_speed_augmentation(
     from datasets import Audio
     expanded_dataset = expanded_dataset.cast_column(audio_column, Audio(sampling_rate=target_sr))
     
-    print(f"✓ Expanded dataset: {len(expanded_dataset):,} samples ({len(expanded_dataset) / len(dataset):.1f}x)")
+    print(f"✓ Expanded dataset: {len(expanded_dataset):,} samples ({len(expanded_dataset) / len(dataset_decoded):.1f}x)")
     if skipped_count > 0:
         print(f"⚠ Skipped {skipped_count} samples during expansion")