Add VADSegmentationStage for voice activity detection - #1578
Conversation
Greptile SummaryThis PR adds Compared to earlier review rounds, many issues have been resolved: Two issues remain before merge:
Confidence Score: 4/5Near-safe to merge; two P1 issues — test suite skipped on CPU CI and nested-mode early-exit returning the wrong type — should be addressed before merging. The implementation has improved significantly across review rounds and the core VAD logic is correct. The two P1 findings are a test infrastructure problem that leaves regressions undetected, and a type contract violation in an error path of nested mode. Neither causes data corruption in normal operation but both affect correctness.
Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[AudioTask input] --> B{_vad_model loaded?}
B -- No --> ERR[raise RuntimeError]
B -- Yes --> D[_resolve_audio]
D -- waveform present --> E[ensure_waveform_2d]
D -- waveform missing, filepath valid --> F[load_audio_file]
D -- both missing --> G[return None]
F --> E
G --> BUG{nested mode?}
BUG -- nested=True --> BUG2[BUG: returns empty list instead of AudioTask]
BUG -- nested=False --> EL[return empty list]
E --> K[_get_vad_segments]
K --> L{sample_rate in SILERO_SUPPORTED_RATES?}
L -- Yes --> M[pass waveform as-is]
L -- No --> N[Resample to 16kHz via torchaudio]
M --> O[get_speech_timestamps]
N --> O
O --> P{segments empty?}
P -- Yes, nested=True --> R[return task with segments=[]]
P -- Yes, nested=False --> S[return empty list]
P -- No, nested=True --> U[build segment dicts stored in task.data.segments]
U --> V[delete waveform from task.data]
V --> W[return single AudioTask]
P -- No, nested=False --> X[build one AudioTask per segment]
X --> Y[return list of AudioTasks fan-out]
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| 'duration_sec': (end_ms - start_ms) / 1000.0, | ||
| 'original_file': item.get('audio_filepath', 'unknown'), | ||
| } | ||
|
|
||
| # Copy any metadata from original item |
There was a problem hiding this comment.
Lazy lock creation is not thread-safe
_get_lock() itself is not protected by a lock. Two threads can both observe _init_lock is None concurrently and each create a separate threading.Lock() instance. When _initialize_model then calls with self._get_lock(), the two threads could be holding different lock objects, so the double-checked locking pattern no longer prevents concurrent model initialization.
The idiomatic fix is to create the lock once, before it can be shared across threads. Because threading.Lock is not picklable, it must be created lazily — but the creation itself must be guarded. A common solution is to use a module-level or class-level "meta-lock" for the creation step, or to simply create the lock unconditionally in __setstate__ and in __post_init__ using a sentinel that is reset on unpickling:
# In __post_init__:
self._init_lock = threading.Lock() # safe, object not yet shared
# In __setstate__:
self.__dict__.update(state)
self._init_lock = threading.Lock() # recreate after unpicklingThis way _get_lock() is no longer needed and the double-checked locking in _initialize_model becomes correct.
|
|
||
|
|
||
| @dataclass | ||
| class VADConfig: |
There was a problem hiding this comment.
Discussed offline. We do not need a config class per stage.
| def __getstate__(self): | ||
| """Return state for pickling, excluding unpicklable objects.""" | ||
| state = self.__dict__.copy() | ||
| # Remove the lock and model - they'll be recreated | ||
| state['_init_lock'] = None | ||
| state['_vad_model'] = None | ||
| state['_vad_utils'] = None | ||
| state['_device'] = None | ||
| return state | ||
|
|
||
| def __setstate__(self, state): | ||
| """Restore state after unpickling.""" | ||
| self.__dict__.update(state) | ||
| self._init_lock = None | ||
| self._vad_model = None | ||
| self._vad_utils = None | ||
| self._device = None | ||
|
|
||
| def _get_lock(self): | ||
| """Get or create the initialization lock (lazy initialization).""" | ||
| if self._init_lock is None: | ||
| self._init_lock = threading.Lock() | ||
| return self._init_lock |
There was a problem hiding this comment.
IIUC we can remove these.
| return [], ['audio', 'waveform', 'sample_rate', 'start_ms', 'end_ms', 'segment_num', 'duration_sec'] | ||
|
|
||
| def ray_stage_spec(self) -> dict[str, Any]: | ||
| from nemo_curator.backends.experimental.utils import RayStageSpecKeys |
There was a problem hiding this comment.
This can be a top-level import.
| # Determine device based on resources and CUDA availability | ||
| # Use GPU if _resources.gpus > 0 and CUDA is available | ||
| use_gpu = self._resources.gpus > 0 and torch.cuda.is_available() |
There was a problem hiding this comment.
What would you recommend as the default here? GPU or no GPU?
| if use_gpu: | ||
| self._device = torch.device(f'cuda:{torch.cuda.current_device()}') | ||
| model = model.to(self._device) | ||
| logger.info(f"Silero VAD model loaded on GPU: {self._device}") |
There was a problem hiding this comment.
It should just work OOTB with self._device = torch.device("cuda").
| def test_defaults(self) -> None: | ||
| cfg = VADConfig() | ||
| assert cfg.min_interval_ms == 500 | ||
| assert cfg.min_duration_sec == 2.0 | ||
| assert cfg.max_duration_sec == 60.0 | ||
| assert cfg.threshold == 0.5 | ||
| assert cfg.speech_pad_ms == 300 | ||
|
|
||
| def test_from_dict(self) -> None: | ||
| cfg = VADConfig.from_dict({"min_duration_sec": 3.0, "threshold": 0.6}) | ||
| assert cfg.min_duration_sec == 3.0 | ||
| assert cfg.threshold == 0.6 | ||
| assert cfg.max_duration_sec == 60.0 | ||
|
|
||
| def test_from_dict_none(self) -> None: | ||
| cfg = VADConfig.from_dict(None) | ||
| assert cfg.threshold == 0.5 | ||
|
|
||
| def test_from_dict_ignores_unknown(self) -> None: | ||
| cfg = VADConfig.from_dict({"unknown": 99, "threshold": 0.3}) | ||
| assert cfg.threshold == 0.3 |
| assert restored.threshold == original.threshold | ||
| assert restored.speech_pad_ms == original.speech_pad_ms | ||
|
|
||
| def test_get(self) -> None: |
| def test_stage_properties(self) -> None: | ||
| stage = VADSegmentationStage() | ||
| assert stage.name == "VADSegmentation" | ||
| assert stage.inputs() == (["data"], []) | ||
| _, output_keys = stage.outputs() | ||
| for key in ["audio", "waveform", "sample_rate", "start_ms", "end_ms", | ||
| "segment_num", "duration_sec"]: | ||
| assert key in output_keys | ||
|
|
||
| def test_config_overrides_params(self) -> None: | ||
| cfg = VADConfig(min_duration_sec=3.0, max_duration_sec=20.0, threshold=0.7) | ||
| stage = VADSegmentationStage(config=cfg) | ||
| assert stage.min_duration_sec == 3.0 | ||
| assert stage.max_duration_sec == 20.0 | ||
| assert stage.threshold == 0.7 | ||
|
|
||
| def test_ray_stage_spec_is_fanout(self) -> None: | ||
| stage = VADSegmentationStage() | ||
| spec = stage.ray_stage_spec() | ||
| from nemo_curator.backends.experimental.utils import RayStageSpecKeys | ||
| assert spec[RayStageSpecKeys.IS_FANOUT_STAGE] is True |
| @property | ||
| def min_interval_ms(self) -> int: | ||
| return _MIN_INTERVAL_MS | ||
|
|
||
| @property | ||
| def threshold(self) -> float: | ||
| return _THRESHOLD | ||
|
|
||
| @property | ||
| def speech_pad_ms(self) -> int: | ||
| return _SPEECH_PAD_MS |
There was a problem hiding this comment.
threshold, speech_pad_ms, and min_interval_ms are non-configurable read-only properties
These three parameters are declared as @property methods returning module-level constants, not as dataclass fields. This means they can never be configured by the caller, and the following all break at runtime:
VADConfig(threshold=0.7)→TypeError: __init__() got an unexpected keyword argument 'threshold'VADConfig.from_dict({"threshold": 0.6})silently ignores"threshold"because__dataclass_fields__only containsmin_duration_secandmax_duration_secto_dict()doesn't serialize them, so a round-trip always loses these values
This directly breaks multiple tests:
test_from_dict(line 36-39): callsfrom_dict({"threshold": 0.6})and assertscfg.threshold == 0.6— always 0.5, failstest_from_dict_ignores_unknown(line 46-47): assertscfg.threshold == 0.3afterfrom_dict({"threshold": 0.3})— always 0.5, failstest_roundtrip(line 56): callsVADConfig(min_duration_sec=1.5, threshold=0.7, speech_pad_ms=200)— raisesTypeErrortest_config_overrides_paramsin the stage tests (line 81): callsVADConfig(threshold=0.7)— raisesTypeError
These fields must be promoted to proper dataclass fields to be configurable:
@dataclass
class VADConfig:
min_duration_sec: float = 2.0
max_duration_sec: float = 60.0
threshold: float = _THRESHOLD
speech_pad_ms: int = _SPEECH_PAD_MS
min_interval_ms: int = _MIN_INTERVAL_MSAnd to_dict() must include them so the round-trip works correctly.
| from dataclasses import dataclass, field | ||
| from typing import Any, List, Dict, Tuple, Optional | ||
|
|
||
| import numpy as np |
There was a problem hiding this comment.
Unused import
numpy (import numpy as np) is imported at line 44 but never referenced anywhere in the file. This will trigger linting warnings and adds unnecessary overhead.
| import numpy as np | |
| import torch |
| def _get_vad_segments(self, waveform: torch.Tensor, sample_rate: int) -> List[Dict[str, float]]: | ||
| """Get speech segments using VAD.""" | ||
| # Ensure waveform is 1D for VAD | ||
| if waveform.dim() > 1: | ||
| waveform = waveform.squeeze(0) |
There was a problem hiding this comment.
squeeze(0) silently no-ops on multi-channel waveforms
The intent here is to reduce a (1, N) waveform to (N,) for Silero VAD, which requires a 1-D tensor. However squeeze(0) only removes a dimension if its size is 1. For a stereo or multi-channel waveform with shape (C, N) where C > 1, squeeze(0) is a no-op and a 2-D tensor is silently passed to get_speech_timestamps. Silero VAD does not handle 2-D input correctly and will raise an error deep in the model.
A more defensive approach is to first convert to mono and then squeeze:
if waveform.dim() > 1:
if waveform.shape[0] > 1:
waveform = waveform.mean(dim=0) # stereo → mono (N,)
else:
waveform = waveform.squeeze(0) # (1, N) → (N,)This is consistent with the existing mono-conversion logic in _load_audio_file.
| resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=SILERO_TARGET_RATE) | ||
| vad_waveform = resampler(waveform_cpu).squeeze(0) |
There was a problem hiding this comment.
_get_vad_segments call
A new torchaudio.transforms.Resample object is constructed each time this method is called for a non-standard sample rate. In a pipeline processing many audio files at the same sample rate (e.g., all at 22050 Hz), this allocates and initialises a fresh transform — including its internal filter kernel — on every single call.
Consider caching the resampler as an instance attribute keyed by (orig_freq, new_freq), for example:
# In __post_init__ / __setstate__:
self._resamplers: dict[int, torchaudio.transforms.Resample] = {}
# In _get_vad_segments:
if sample_rate not in self._resamplers:
self._resamplers[sample_rate] = torchaudio.transforms.Resample(
orig_freq=sample_rate, new_freq=SILERO_TARGET_RATE
)
resampler = self._resamplers[sample_rate]Remember to also exclude _resamplers from __getstate__ (if pickling support is added) since Resample objects may not be picklable.
d084b28 to
a11c176
Compare
| from silero_vad import load_silero_vad, get_speech_timestamps | ||
|
|
||
| # Suppress Silero VAD sample rate warning (48kHz -> 16kHz is expected) | ||
| warnings.filterwarnings('ignore', message='Sampling rate is a multiply of 16000') |
There was a problem hiding this comment.
warnings.filterwarnings is a process-wide side effect
warnings.filterwarnings('ignore', ...) is called at module import time and silences the matched warning globally for the entire Python process — not just within this file. Any other code (tests, user code, third-party libraries) that would legitimately emit a warning matching 'Sampling rate is a multiply of 16000' will also be silenced after this module is first imported.
The idiomatic way to suppress a known-noisy warning only at the specific call site is to use warnings.catch_warnings() as a context manager around the load_silero_vad() call:
import contextlib
with warnings.catch_warnings():
warnings.filterwarnings('ignore', message='Sampling rate is a multiply of 16000')
model = load_silero_vad()This keeps the filter scoped to where it's actually needed.
| # Sample rates like 22050 Hz need to be resampled to 16kHz | ||
| vad_sample_rate = sample_rate | ||
| vad_waveform = waveform | ||
| if sample_rate not in SILERO_SUPPORTED_RATES and sample_rate % 16000 != 0: |
There was a problem hiding this comment.
sample_rate % 16000 != 0 silently bypasses resampling for unsupported rates
The condition sample_rate not in SILERO_SUPPORTED_RATES and sample_rate % 16000 != 0 means that any rate which is a multiple of 16000 but is not in SILERO_SUPPORTED_RATES (e.g. 80 000 Hz, 112 000 Hz) will skip resampling and be fed raw to Silero VAD. Silero's documented supported rates are only 8 kHz and 16 kHz (with the filter-bank supporting 8 k multiples up to 96 k in some versions), but very high rates like 80 kHz are not guaranteed to work and will silently produce incorrect timestamps or a runtime error deep inside the model.
The condition should be tightened to only skip resampling for rates that are explicitly known to be supported:
if sample_rate not in SILERO_SUPPORTED_RATES:
# resample to 16 kHz for all unsupported rates
...This also makes SILERO_SUPPORTED_RATES the single source of truth, instead of having a secondary % 16000 escape hatch.
a11c176 to
76d6af7
Compare
| resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=SILERO_TARGET_RATE) | ||
| vad_waveform = resampler(waveform_cpu).squeeze(0) | ||
| # Move back to original device | ||
| if device.type != 'cpu': | ||
| vad_waveform = vad_waveform.to(device) | ||
| vad_sample_rate = SILERO_TARGET_RATE | ||
|
|
||
| # Get speech timestamps using silero_vad package | ||
| speech_timestamps = get_speech_timestamps( | ||
| vad_waveform, | ||
| self.vad_model, |
There was a problem hiding this comment.
device variable captures pre-move device
The code moves waveform to self._device before determining whether resampling is needed:
if self._device is not None and waveform.device != self._device:
waveform = waveform.to(self._device)
vad_sample_rate = sample_rate
vad_waveform = waveform
if sample_rate not in SILERO_SUPPORTED_RATES and sample_rate % 16000 != 0:
device = waveform.device # ← captures self._device (GPU)
waveform_cpu = waveform.cpu() if waveform.device.type != 'cpu' else waveformdevice = waveform.device correctly captures self._device (since waveform was just moved there). The resampled tensor is then moved back to device (i.e., self._device). This is correct in the current code.
However, if self._device is None (e.g., before _initialize_model has been called, or if the model is not loaded), the waveform.to(self._device) guard is skipped, so device will be whatever device the input tensor lives on. If the caller passes a CPU tensor, everything is fine. But if _initialize_model was never called and a GPU tensor is passed, get_speech_timestamps will fail because the model is None.
Add an assertion or guard to ensure _initialize_model() has been called (and _device is set) before this method is invoked, to make the contract explicit.
| if self._vad_model is not None: | ||
| del self._vad_model | ||
| self._vad_model = None | ||
| torch.cuda.empty_cache() |
There was a problem hiding this comment.
teardown calls torch.cuda.empty_cache() unconditionally
torch.cuda.empty_cache() is always called even when the model ran entirely on CPU (i.e., self._device == torch.device('cpu')). On CPU-only machines this is a no-op but adds unnecessary overhead on teardown. Consider guarding it:
| torch.cuda.empty_cache() | |
| def teardown(self) -> None: | |
| """Clean up resources.""" | |
| if self._vad_model is not None: | |
| del self._vad_model | |
| self._vad_model = None | |
| if self._device is not None and self._device.type == "cuda": | |
| torch.cuda.empty_cache() |
…load_audio_file, add ensure_waveform_2d in VAD processing
db53643 to
a1206b7
Compare
| if waveform is None or sample_rate is None: | ||
| audio_filepath = item.get('audio_filepath') | ||
| if audio_filepath and os.path.exists(audio_filepath): | ||
| try: | ||
| waveform, sample_rate = load_audio_file(audio_filepath) | ||
| item[self.waveform_key] = waveform | ||
| item[self.sample_rate_key] = sample_rate | ||
| except Exception as e: | ||
| logger.error(f"Failed to load audio file {audio_filepath}: {e}") | ||
| continue | ||
| else: | ||
| logger.error("Missing waveform/sample_rate and no valid audio_filepath provided") | ||
| continue |
There was a problem hiding this comment.
sample_rate is missing
When waveform is not None but sample_rate is None, the condition on line 265 (waveform is None or sample_rate is None) is still True, so the code enters the block and calls load_audio_file(audio_filepath) unconditionally. This reloads the entire audio file and overwrites the user's waveform in item[self.waveform_key] with a freshly-decoded copy.
If the user provided a pre-processed waveform (e.g., already normalised, augmented, or synthesised — not matching the raw file on disk) and simply omitted sample_rate, that waveform is silently discarded and replaced by the raw file contents.
Notably, common.py already contains resolve_waveform_from_item, which handles this split case correctly: it reads only the sample rate from the file header (via soundfile.info) when a waveform is already present. The fix is to separate the two fallback branches:
if waveform is None:
# No waveform at all — load from file
audio_filepath = item.get('audio_filepath')
if audio_filepath and os.path.exists(audio_filepath):
try:
waveform, sample_rate = load_audio_file(audio_filepath)
item[self.waveform_key] = waveform
item[self.sample_rate_key] = sample_rate
except Exception as e:
logger.error(f"Failed to load audio file {audio_filepath}: {e}")
continue
else:
logger.error("Missing waveform/sample_rate and no valid audio_filepath provided")
continue
elif sample_rate is None:
# Waveform present but sample_rate missing — read only header
audio_filepath = item.get('audio_filepath')
if audio_filepath and os.path.exists(audio_filepath):
try:
import soundfile
sample_rate = soundfile.info(audio_filepath).samplerate
item[self.sample_rate_key] = sample_rate
except Exception as e:
logger.error(f"Waveform present but cannot read sample_rate from '{audio_filepath}': {e}")
continue
else:
logger.error("Waveform present but 'sample_rate' missing and no audio_filepath available")
continue
| # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| from unittest.mock import MagicMock, patch | ||
|
|
||
| import torch | ||
|
|
||
| from nemo_curator.stages.audio.segmentation.vad_segmentation import VADSegmentationStage | ||
| from nemo_curator.tasks import AudioBatch | ||
|
|
||
|
|
||
| class TestVADSegmentationStage: | ||
| """Tests for VADSegmentationStage.""" | ||
|
|
||
| @patch("nemo_curator.stages.audio.segmentation.vad_segmentation.get_speech_timestamps") | ||
| @patch("nemo_curator.stages.audio.segmentation.vad_segmentation.load_silero_vad") | ||
| def test_process_returns_segments(self, mock_load_vad, mock_get_ts) -> None: | ||
| mock_model = MagicMock() | ||
| mock_load_vad.return_value = mock_model | ||
|
|
||
| sr = 48000 | ||
| mock_get_ts.return_value = [ | ||
| {"start": 0, "end": sr * 3}, | ||
| {"start": sr * 5, "end": sr * 8}, | ||
| ] | ||
|
|
||
| waveform = torch.randn(1, sr * 10) | ||
| batch = AudioBatch( | ||
| data=[{"waveform": waveform, "sample_rate": sr}], | ||
| task_id="test", | ||
| dataset_name="test", | ||
| ) | ||
|
|
||
| stage = VADSegmentationStage(min_duration_sec=1.0, max_duration_sec=30.0) | ||
| stage.setup() | ||
| result = stage.process(batch) | ||
|
|
||
| assert isinstance(result, list) | ||
| assert len(result) == 2 | ||
| for seg in result: | ||
| assert isinstance(seg, AudioBatch) | ||
| item = seg.data[0] | ||
| assert "waveform" in item | ||
| assert "start_ms" in item | ||
| assert "end_ms" in item | ||
| assert "segment_num" in item | ||
| assert "duration_sec" in item | ||
|
|
||
| @patch("nemo_curator.stages.audio.segmentation.vad_segmentation.get_speech_timestamps") | ||
| @patch("nemo_curator.stages.audio.segmentation.vad_segmentation.load_silero_vad") | ||
| def test_process_output_keys(self, mock_load_vad, mock_get_ts) -> None: | ||
| mock_load_vad.return_value = MagicMock() | ||
|
|
||
| sr = 48000 | ||
| mock_get_ts.return_value = [{"start": 0, "end": sr * 5}] | ||
|
|
||
| waveform = torch.randn(1, sr * 10) | ||
| batch = AudioBatch( | ||
| data=[{"waveform": waveform, "sample_rate": sr}], | ||
| task_id="test", | ||
| dataset_name="test", | ||
| ) | ||
|
|
||
| stage = VADSegmentationStage(min_duration_sec=1.0) | ||
| stage.setup() | ||
| result = stage.process(batch) | ||
|
|
||
| item = result[0].data[0] | ||
| assert item["start_ms"] == 0 | ||
| assert item["segment_num"] == 0 | ||
| assert item["duration_sec"] > 0 | ||
| assert item["sample_rate"] == sr | ||
|
|
||
| @patch("nemo_curator.stages.audio.segmentation.vad_segmentation.get_speech_timestamps") | ||
| @patch("nemo_curator.stages.audio.segmentation.vad_segmentation.load_silero_vad") | ||
| def test_empty_speech_returns_empty(self, mock_load_vad, mock_get_ts) -> None: | ||
| mock_load_vad.return_value = MagicMock() | ||
| mock_get_ts.return_value = [] | ||
|
|
||
| waveform = torch.randn(1, 48000 * 5) | ||
| batch = AudioBatch( | ||
| data=[{"waveform": waveform, "sample_rate": 48000}], | ||
| task_id="test", | ||
| dataset_name="test", | ||
| ) | ||
|
|
||
| stage = VADSegmentationStage() | ||
| stage.setup() | ||
| result = stage.process(batch) | ||
|
|
||
| assert isinstance(result, list) | ||
| assert len(result) == 0 | ||
|
|
||
| @patch("nemo_curator.stages.audio.segmentation.vad_segmentation.get_speech_timestamps") | ||
| @patch("nemo_curator.stages.audio.segmentation.vad_segmentation.load_silero_vad") | ||
| def test_segment_numbering(self, mock_load_vad, mock_get_ts) -> None: | ||
| mock_load_vad.return_value = MagicMock() | ||
|
|
||
| sr = 48000 | ||
| mock_get_ts.return_value = [ | ||
| {"start": 0, "end": sr * 2}, | ||
| {"start": sr * 3, "end": sr * 5}, | ||
| {"start": sr * 6, "end": sr * 8}, | ||
| ] | ||
|
|
||
| waveform = torch.randn(1, sr * 10) | ||
| batch = AudioBatch( | ||
| data=[{"waveform": waveform, "sample_rate": sr}], | ||
| task_id="test", | ||
| dataset_name="test", | ||
| ) | ||
|
|
||
| stage = VADSegmentationStage(min_duration_sec=0.5) | ||
| stage.setup() | ||
| result = stage.process(batch) | ||
|
|
||
| assert len(result) == 3 | ||
| for i, seg in enumerate(result): | ||
| assert seg.data[0]["segment_num"] == i | ||
|
|
||
| @patch("nemo_curator.stages.audio.segmentation.vad_segmentation.get_speech_timestamps") | ||
| @patch("nemo_curator.stages.audio.segmentation.vad_segmentation.load_silero_vad") | ||
| def test_missing_waveform_and_filepath_skipped(self, mock_load_vad, mock_get_ts) -> None: | ||
| mock_load_vad.return_value = MagicMock() | ||
|
|
||
| batch = AudioBatch( | ||
| data=[{"some_key": "value"}], | ||
| task_id="test", | ||
| dataset_name="test", | ||
| ) | ||
|
|
||
| stage = VADSegmentationStage() | ||
| stage.setup() | ||
| result = stage.process(batch) | ||
|
|
||
| assert isinstance(result, list) | ||
| assert len(result) == 0 | ||
|
|
||
| def test_pickling(self) -> None: | ||
| """VADSegmentationStage should be picklable (for Ray workers).""" | ||
| import pickle | ||
|
|
||
| stage = VADSegmentationStage(min_duration_sec=2.0, threshold=0.6) | ||
| pickled = pickle.dumps(stage) | ||
| restored = pickle.loads(pickled) | ||
| assert restored.min_duration_sec == 2.0 | ||
| assert restored.threshold == 0.6 | ||
| assert restored._vad_model is None |
There was a problem hiding this comment.
mode="batch" test coverage
The test suite only exercises the default mode="fanout" path. The mode="batch" path in process() has distinct control-flow branches — notably, it returns a single AudioBatch with all segments as items rather than a fan-out list — but none of these branches are tested. Key scenarios that are entirely uncovered include:
- Normal operation: verify that
process(batch)inmode="batch"returns a singleAudioBatchwhosedatais a list of all segment dicts (not a fan-out list). - Empty segments: verify that an empty
AudioBatch(data=[])is returned (not[]) when no speech is detected. - Model unavailable: verify the
mode="batch"early-return path when_vad_model is None.
Without these tests, a regression in the batch mode code path would go undetected.
0ed8d79 to
f53d83c
Compare
| from silero_vad import load_silero_vad, get_speech_timestamps | ||
|
|
||
| # Silero VAD only supports 8kHz, 16kHz, and multiples of 16kHz (32k, 48k, etc.) | ||
| # Sample rates like 22050 Hz need to be resampled to 16kHz | ||
| SILERO_SUPPORTED_RATES = {8000, 16000, 32000, 48000, 64000, 96000} | ||
| SILERO_TARGET_RATE = 16000 | ||
|
|
||
| from nemo_curator.backends.experimental.utils import RayStageSpecKeys | ||
| from nemo_curator.stages.audio.common import load_audio_file, ensure_waveform_2d | ||
| from nemo_curator.stages.base import ProcessingStage | ||
| from nemo_curator.stages.resources import Resources | ||
| from nemo_curator.tasks import AudioBatch | ||
| from nemo_curator.stages.audio.configs.vad import VADConfig |
There was a problem hiding this comment.
SILERO_SUPPORTED_RATES and SILERO_TARGET_RATE (lines 51–52) are defined between two groups of import statements — after the third-party silero_vad import and before the nemo_curator imports. PEP 8 recommends that module-level constants follow all imports, not be interleaved with them. This also confuses linters and import-ordering tools like isort.
| from silero_vad import load_silero_vad, get_speech_timestamps | |
| # Silero VAD only supports 8kHz, 16kHz, and multiples of 16kHz (32k, 48k, etc.) | |
| # Sample rates like 22050 Hz need to be resampled to 16kHz | |
| SILERO_SUPPORTED_RATES = {8000, 16000, 32000, 48000, 64000, 96000} | |
| SILERO_TARGET_RATE = 16000 | |
| from nemo_curator.backends.experimental.utils import RayStageSpecKeys | |
| from nemo_curator.stages.audio.common import load_audio_file, ensure_waveform_2d | |
| from nemo_curator.stages.base import ProcessingStage | |
| from nemo_curator.stages.resources import Resources | |
| from nemo_curator.tasks import AudioBatch | |
| from nemo_curator.stages.audio.configs.vad import VADConfig | |
| from silero_vad import load_silero_vad, get_speech_timestamps | |
| from nemo_curator.backends.experimental.utils import RayStageSpecKeys | |
| from nemo_curator.stages.audio.common import load_audio_file, ensure_waveform_2d | |
| from nemo_curator.stages.base import ProcessingStage | |
| from nemo_curator.stages.resources import Resources | |
| from nemo_curator.tasks import AudioBatch | |
| from nemo_curator.stages.audio.configs.vad import VADConfig | |
| # Silero VAD only supports 8kHz, 16kHz, and multiples of 16kHz (32k, 48k, etc.) | |
| # Sample rates like 22050 Hz need to be resampled to 16kHz | |
| SILERO_SUPPORTED_RATES = {8000, 16000, 32000, 48000, 64000, 96000} | |
| SILERO_TARGET_RATE = 16000 |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| def test_process_returns_segments(self, mock_load_vad, mock_get_ts) -> None: | ||
| mock_model = MagicMock() | ||
| mock_load_vad.return_value = mock_model | ||
|
|
||
| sr = 48000 | ||
| mock_get_ts.return_value = [ | ||
| {"start": 0, "end": sr * 3}, | ||
| {"start": sr * 5, "end": sr * 8}, | ||
| ] | ||
|
|
||
| waveform = torch.randn(1, sr * 10) | ||
| batch = AudioBatch( | ||
| data=[{"waveform": waveform, "sample_rate": sr}], | ||
| task_id="test", | ||
| dataset_name="test", | ||
| ) | ||
|
|
||
| stage = VADSegmentationStage(min_duration_sec=1.0, max_duration_sec=30.0) | ||
| stage.setup() | ||
| result = stage.process(batch) | ||
|
|
||
| assert isinstance(result, list) | ||
| assert len(result) == 2 | ||
| for seg in result: | ||
| assert isinstance(seg, AudioBatch) | ||
| item = seg.data[0] | ||
| assert "waveform" in item | ||
| assert "start_ms" in item | ||
| assert "end_ms" in item | ||
| assert "segment_num" in item | ||
| assert "duration_sec" in item |
There was a problem hiding this comment.
mock_get_ts.return_value uses raw sample-count values (e.g. {"start": 0, "end": sr * 3} where sr = 48000). Since 48 000 Hz is in SILERO_SUPPORTED_RATES, the code does not resample and correctly divides each timestamp by vad_sample_rate (= 48000) to convert to seconds. The resulting seconds are then multiplied back by the original sample_rate to obtain start_sample/end_sample.
However, the round-trip (sr_value / vad_sample_rate) * sample_rate only produces correct slice boundaries because the input sample rate and vad_sample_rate happen to be the same (no resampling). The test does not cover the resampling path (e.g. sr = 22050), where vad_sample_rate becomes SILERO_TARGET_RATE = 16000 and the timestamp-to-sample conversion uses a different ratio. Without a test for this case, a regression in the resampling timestamp conversion (such as computing ts['start'] / vad_sample_rate * sample_rate vs. ts['start'] / sample_rate) would go undetected.
Consider adding a test that uses a non-supported sample rate (e.g. 22050 Hz) to exercise the resampling path and verify that segment start_ms/end_ms and waveform slice boundaries are correct.
|
/ok to test 2de1a9d |
|
/ok to test 4088c20 |
|
/ok to test 46eb1d8 |
|
/ok to test c26094d |
|
/ok to test 6e650d8 |
|
/ok to test e240495 |
|
/ok to test ef09da8 |
…p-level imports Signed-off-by: shbhawsar <shbhawsar@nvidia.com>
Signed-off-by: shbhawsar <shbhawsar@nvidia.com>
Co-authored-by: Sarah Yurick <53962159+sarahyurick@users.noreply.github.com> Signed-off-by: Sarah Yurick <53962159+sarahyurick@users.noreply.github.com>
Signed-off-by: Sarah Yurick <53962159+sarahyurick@users.noreply.github.com>
Signed-off-by: Sarah Yurick <53962159+sarahyurick@users.noreply.github.com>
|
/ok to test 835cd3e |
|
Hi @shubhamNvidia it looks like the audio CPU tests are blocked by needing a GPU dependency. You can mark GPU-only tests with |
Signed-off-by: shbhawsar <shbhawsar@nvidia.com>
Signed-off-by: shbhawsar <shbhawsar@nvidia.com>
|
/ok to test baee33b |
|
|
||
| name: str = "VADSegmentation" | ||
| batch_size: int = 1 | ||
| resources: Resources = field(default_factory=lambda: Resources(cpus=1.0, gpus=0.3)) |
There was a problem hiding this comment.
VADSegmentationStage in nemo_curator/stages/audio/segmentation/vad_segmentation.py defaults to resources=Resources(cpus=1.0, gpus=0.3). On a CPU-only machine, _check_gpu_availability() raises a hard RuntimeError.
Silero VAD is lightweight (~1 MB model) and runs efficiently on CPU. Either:
(a) Default to gpus=0 and let users opt into GPU, or
(b) Change _check_gpu_availability from hard error to graceful CPU fallback:
if gpus > 0 and not torch.cuda.is_available():
logger.warning("GPU requested but CUDA not available, falling back to CPU")
|
|
||
| audio_result = self._resolve_audio(task.data) | ||
| if audio_result is None: | ||
| return [] |
There was a problem hiding this comment.
In nested mode (nested=True), ray_stage_spec() returns {} (no fan-out), implying a 1:1 input→output contract. But when no speech is detected, process() returns [] regardless of the nested flag, dropping the task entirely.
This breaks the 1:1 contract that nested mode implies. Downstream stages (e.g., SegmentConcatenationStage in AudioDataFilterStage) never see the task.
For nested mode with no speech, return the task with empty segments instead:
if not segments:
if self.nested:
task.data["segments"] = []
return task
return []
|
/ok to test 599f658 |
- Default gpus=0.0 (CPU) for VADSegmentationStage since Silero VAD is lightweight; users opt into GPU via .with_(resources=Resources(gpus=X)) - In nested mode, return task with empty segments list instead of [] when no speech is detected, preserving the 1:1 input-output contract - Update test to match new nested no-speech behavior
4 participants
Adds VADSegmentationStage which segments audio into speech chunks using Silero VAD. Produces a fan-out list of AudioBatch objects, one per detected speech segment, with start_ms, end_ms, segment_num, duration_sec, and both PyDub AudioSegment and torch waveform outputs. Supports configurable min_duration_sec, max_duration_sec, threshold, and speech_pad_ms. Works on both CPU and GPU. Includes VADConfig, pickling support for Ray workers, and unit tests with mocked Silero VAD.