Fix AudioDecoder handling: use get_all_samples() method to extract audio data from AudioDecoder objects

training/augmentation.py

@@ -307,66 +307,89 @@ def expand_dataset_with_speed_augmentation(
     print(f"  Original size: {len(dataset):,} samples")
     print(f"  Speed factors: {speed_factors}")
     
-    # 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]
+    # Access audio column directly using bracket notation to trigger decoding
+    # HuggingFace datasets automatically decode AudioDecoder when accessed via bracket notation
+    print("  Processing audio files (decoding on-the-fly)...")
+    
+    # Use indexed access to iterate over dataset
+    for idx in range(len(dataset)):
+        example = dataset[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)
+        # Access audio column using bracket notation - this should trigger automatic decoding
+        # Try bracket access first (triggers decoding), fallback to .get() if needed
+        try:
+            audio_data = example[audio_column]  # Bracket access triggers decoding
+        except (KeyError, TypeError):
+            try:
+                audio_data = example.get(audio_column)
+            except:
+                audio_data = None
+        
         transcription = example.get(transcription_column, "")
         
         # Handle audio data format
-        # After decode(), audio_data should be a dict with 'array' and 'sampling_rate'
+        # AudioDecoder objects need to use get_all_samples() to extract audio data
         audio_array = None
         audio_sr = target_sr
         
         if audio_data is not None:
+            # Case 1: Already decoded dict format (ideal case)
             if isinstance(audio_data, dict):
                 audio_array = audio_data.get("array")
                 audio_sr = audio_data.get("sampling_rate", target_sr)
-            # 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)
-            # Fallback: handle numpy array directly
-            elif isinstance(audio_data, np.ndarray):
-                audio_array = audio_data
-            # If still an AudioDecoder, try to decode it manually
+            # Case 2: AudioDecoder object - use get_all_samples() method
             elif hasattr(audio_data, "__class__") and "Decoder" in str(type(audio_data)):
                 try:
-                    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)
+                    # AudioDecoder has get_all_samples() method that returns AudioSamples
+                    if hasattr(audio_data, "get_all_samples"):
+                        audio_samples = audio_data.get_all_samples()
+                        # AudioSamples has .data (PyTorch tensor) and .sample_rate
+                        if hasattr(audio_samples, "data"):
+                            # Convert PyTorch tensor to numpy array
+                            if hasattr(audio_samples.data, "numpy"):
+                                audio_array = audio_samples.data.numpy()
+                            elif hasattr(audio_samples.data, "cpu"):
+                                # If it's on GPU, move to CPU first
+                                audio_array = audio_samples.data.cpu().numpy()
+                            else:
+                                # Try to convert directly
+                                audio_array = np.array(audio_samples.data)
+                            
+                            # Handle multi-channel audio (take first channel if stereo)
+                            if audio_array.ndim > 1:
+                                audio_array = audio_array[0] if audio_array.shape[0] == 1 else audio_array.mean(axis=0)
+                            
+                            # Get sampling rate
+                            if hasattr(audio_samples, "sample_rate"):
+                                audio_sr = audio_samples.sample_rate
+                            elif hasattr(audio_samples, "sampling_rate"):
+                                audio_sr = audio_samples.sampling_rate
+                        else:
+                            raise ValueError("AudioSamples object doesn't have 'data' attribute")
+                    else:
+                        raise ValueError("AudioDecoder doesn't have 'get_all_samples' method")
                 except Exception as e:
                     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)}")
+                        import traceback
+                        print(f"  Error details: {traceback.format_exc()}")
                     continue
+            # Case 3: Audio object with .array attribute
+            elif hasattr(audio_data, "array"):
+                audio_array = audio_data.array
+                audio_sr = getattr(audio_data, "sampling_rate", target_sr)
+            # Case 4: Already a numpy array
+            elif isinstance(audio_data, np.ndarray):
+                audio_array = audio_data
             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)}")
+                    print(f"  Available attributes: {dir(audio_data)[:10]}")
                 continue
         
         if audio_array is None:
@@ -429,7 +452,7 @@ def expand_dataset_with_speed_augmentation(
         
         # Progress update every 1000 samples
         if (idx + 1) % 1000 == 0:
-            print(f"  Processed {idx + 1:,}/{len(dataset_decoded):,} samples...")
+            print(f"  Processed {idx + 1:,}/{len(dataset):,} samples...")
     
     if len(expanded_examples) == 0:
         raise ValueError(
@@ -445,7 +468,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_decoded):.1f}x)")
+    print(f"✓ Expanded dataset: {len(expanded_dataset):,} samples ({len(expanded_dataset) / len(dataset):.1f}x)")
     if skipped_count > 0:
         print(f"⚠ Skipped {skipped_count} samples during expansion")