add topk #6 (#8)

Author: gfdb

tests/test_gpu_augmentations.py

@@ -138,3 +138,33 @@ def _noop_add_noise(waveforms, sample_rate, **kwargs):
     out_wave, out_lengths = aug(waveforms, lengths=lengths)
     assert out_wave.shape[0] == waveforms.shape[0]
     assert out_lengths.data_ptr() == lengths.data_ptr()
+
+
+def test_wav2aug_top_k_limits_ops(monkeypatch):
+    """Test that top_k limits the number of augmentations used."""
+
+    def _noop_add_noise(waveforms, loader, **kwargs):
+        return waveforms
+
+    monkeypatch.setattr("wav2aug.gpu.wav2aug.add_noise", _noop_add_noise)
+
+    # top_k=3 should only include: noise, freq_drop, time_dropout
+    aug = Wav2Aug(sample_rate=16_000, top_k=3)
+    assert len(aug._base_ops) == 3
+
+    # top_k=6 should include: noise, freq_drop, time_dropout, speed_perturb, amp_clip, chunk_swap
+    aug = Wav2Aug(sample_rate=16_000, top_k=6)
+    assert len(aug._base_ops) == 6
+
+    # no pass should include all 9
+    aug = Wav2Aug(sample_rate=16_000)
+    assert len(aug._base_ops) == 9
+
+
+def test_wav2aug_top_k_invalid_raises():
+    """Test that invalid top_k values raise ValueError."""
+    with pytest.raises(ValueError, match="top_k must be between 1 and 9"):
+        Wav2Aug(sample_rate=16_000, top_k=0)
+
+    with pytest.raises(ValueError, match="top_k must be between 1 and 9"):
+        Wav2Aug(sample_rate=16_000, top_k=10)

wav2aug/gpu/wav2aug.py

@@ -22,6 +22,7 @@ def __init__(
         sample_rate: int,
         noise_dir: str | None = None,
         noise_preload: bool = True,
+        top_k: int = 9,
     ) -> None:
         """Initialize Wav2Aug.
 
@@ -31,6 +32,11 @@ def __init__(
                 default cached noise pack (auto-downloaded if needed).
             noise_preload: If True (default), preload all noise files into CPU RAM
                 at initialization for fast sampling. If False, load files on-demand.
+            top_k: Number of top augmentations to use, ordered by effectiveness.
+                Default is 9 (all augmentations). Common values: 3, 6, or 9.
+                Order (best to worst): Noise Addition, Freq Drop, Time Drop,
+                Speed Perturb, Amp Clip, Chunk Swap, Babble Noise, Amp Scale,
+                Polarity Inversion.
         """
         self.sample_rate = int(sample_rate)
 
@@ -41,24 +47,31 @@ def __init__(
             noise_dir = ensure_pack("pointsource_noises")
         self._noise_loader = NoiseLoader(noise_dir, sample_rate, preload=noise_preload)
 
-        self._base_ops: List[
-            Callable[[torch.Tensor, torch.Tensor | None], torch.Tensor]
-        ] = [
+        # All ops ordered by effectiveness (best first)
+        all_ops: List[Callable[[torch.Tensor, torch.Tensor | None], torch.Tensor]] = [
+            # top 3
             lambda x, lengths: add_noise(
                 x, self._noise_loader, snr_low=0.0, snr_high=10.0
             ),
-            lambda x, lengths: add_babble_noise(x),
-            lambda x, lengths: chunk_swap(x),
             lambda x, lengths: freq_drop(x),
-            lambda x, lengths: invert_polarity(x),
-            lambda x, lengths: rand_amp_clip(x),
-            lambda x, lengths: rand_amp_scale(x),
-            lambda x, lengths: speed_perturb(x, sample_rate=self.sample_rate),
             lambda x, lengths: time_dropout(
                 x, sample_rate=self.sample_rate, lengths=lengths
             ),
+            # top 6
+            lambda x, lengths: speed_perturb(x, sample_rate=self.sample_rate),
+            lambda x, lengths: rand_amp_clip(x),
+            lambda x, lengths: chunk_swap(x),
+            # all 9
+            lambda x, lengths: add_babble_noise(x),
+            lambda x, lengths: rand_amp_scale(x),
+            lambda x, lengths: invert_polarity(x),
         ]
 
+        if top_k < 1 or top_k > len(all_ops):
+            raise ValueError(f"top_k must be between 1 and {len(all_ops)}, got {top_k}")
+
+        self._base_ops = all_ops[:top_k]
+
     @torch.no_grad()
     def __call__(
         self,