diff --git a/nemo_curator/stages/audio/segmentation/speaker_separation.py b/nemo_curator/stages/audio/segmentation/speaker_separation.py index 7ba80c33ba..6876d3b307 100755 --- a/nemo_curator/stages/audio/segmentation/speaker_separation.py +++ b/nemo_curator/stages/audio/segmentation/speaker_separation.py @@ -30,6 +30,7 @@ ) """ +import os from dataclasses import dataclass, field from typing import Any @@ -44,7 +45,13 @@ SortformerEncLabelModel = None from nemo_curator.backends.base import WorkerMetadata -from nemo_curator.stages.audio.common import resolve_waveform_from_item +from nemo_curator.stages.audio._agent._agent_ready import AgentReady, Gates, IOSpec, StageContract +from nemo_curator.stages.audio._agent._residency import ( + InputResidency, + accepts_for_residency, + resolve_audio, + write_audio_stable, +) from nemo_curator.stages.audio.segmentation.speaker_separation_module.speaker_sep import SpeakerSeparator from nemo_curator.stages.base import ProcessingStage from nemo_curator.stages.resources import Resources @@ -61,7 +68,7 @@ def _pydub_to_waveform_sr(seg: AudioSegment) -> tuple[torch.Tensor, int]: @dataclass -class SpeakerSeparationStage(ProcessingStage[AudioTask, AudioTask]): +class SpeakerSeparationStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): """ Speaker separation stage using NeMo SortFormer diarization model. @@ -75,8 +82,28 @@ class SpeakerSeparationStage(ProcessingStage[AudioTask, AudioTask]): min_duration: Minimum segment duration in seconds gap_threshold: Gap threshold for merging speaker segments buffer_time: Buffer time around speaker segments + audio_filepath_key: Key in data dict for the input audio file path. + waveform_key: Key in data dict for the in-memory waveform tensor. + sample_rate_key: Key in data dict for the waveform sample rate. + speaker_id_key: Key where each child task's speaker id is written. + num_speakers_key: Key where the detected speaker count is written. + duration_key: Key where each child's speech duration in seconds is written. + diar_segments_key: Key where each child's diarization segments are written. + input_residency: Which input to use — "waveform" (in-memory only), "file" + (audio_filepath only), or "auto" (waveform first, file fallback; default). + keep_waveform_in_task: Keep each per-speaker waveform in the task (default True, + today's behavior). Set False to emit only on-disk paths (requires write_to_disk). + write_to_disk: Also write each per-speaker track to ``separated_audio_dir`` and set + ``audio_filepath_key`` to it, so file-based downstream stages can consume it. + Defaults to False (in-memory only, unchanged). + separated_audio_dir: Directory for per-speaker WAVs (required when write_to_disk=True). Note: + By default (write_to_disk=False) per-speaker child tasks DROP the parent's + audio_filepath (it points at the full multi-speaker file) and carry + ``original_file`` for provenance; downstream consumes the per-speaker waveform. + With write_to_disk=True, each child instead gets its own audio_filepath. + GPU assignment is handled by the executor via _resources. Use .with_(resources=Resources(gpus=X)) to configure GPU allocation. """ @@ -86,6 +113,18 @@ class SpeakerSeparationStage(ProcessingStage[AudioTask, AudioTask]): min_duration: float = 0.8 gap_threshold: float = 0.1 buffer_time: float = 0.5 + audio_filepath_key: str = "audio_filepath" + waveform_key: str = "waveform" + sample_rate_key: str = "sample_rate" + speaker_id_key: str = "speaker_id" + num_speakers_key: str = "num_speakers" + duration_key: str = "duration" + diar_segments_key: str = "diar_segments" + input_residency: InputResidency = "auto" + # Output residency (both default to today's behavior: in-memory waveform only, no disk). + keep_waveform_in_task: bool = True + write_to_disk: bool = False + separated_audio_dir: str | None = None name: str = "SpeakerSeparation" batch_size: int = 1 @@ -94,12 +133,64 @@ class SpeakerSeparationStage(ProcessingStage[AudioTask, AudioTask]): def __post_init__(self): super().__init__() self._separator = None + if not (self.keep_waveform_in_task or self.write_to_disk): + msg = "At least one of keep_waveform_in_task or write_to_disk must be True" + raise ValueError(msg) + if self.write_to_disk and not self.separated_audio_dir: + msg = "separated_audio_dir is required when write_to_disk=True" + raise ValueError(msg) def inputs(self) -> tuple[list[str], list[str]]: return [], [] def outputs(self) -> tuple[list[str], list[str]]: - return [], ["waveform", "sample_rate", "speaker_id", "num_speakers", "duration"] + outs: list[str] = [] + if self.keep_waveform_in_task: + outs.extend([self.waveform_key, self.sample_rate_key]) + outs.extend([self.speaker_id_key, self.num_speakers_key, self.duration_key, self.diar_segments_key]) + if self.write_to_disk: + outs.append(self.audio_filepath_key) + return [], outs + + def describe(self) -> StageContract: + forms = accepts_for_residency(self.input_residency) + reads_one_of = [] + if "waveform" in forms: + reads_one_of.append(IOSpec(data_keys=[self.waveform_key, self.sample_rate_key], accepts=["waveform"])) + if "file" in forms: + reads_one_of.append(IOSpec(data_keys=[self.audio_filepath_key], accepts=["file"])) + writes: list[str] = [] + produces: list[str] = [] + if self.keep_waveform_in_task: + writes.extend([self.waveform_key, self.sample_rate_key]) + produces.append("tensor") + writes.extend( + [self.speaker_id_key, self.num_speakers_key, self.duration_key, self.diar_segments_key, "original_file"] + ) + if self.write_to_disk: + writes.append(self.audio_filepath_key) + produces.append("disk") + return StageContract( + reads_one_of=reads_one_of, + writes=IOSpec(data_keys=writes, produces=produces), + # children drop the parent's audio_filepath (and blob keys) + preserves_upstream_keys=False, + cardinality="1:N fan-out", + # One child per detected speaker; speaker_id is the per-child key that + # identifies which slice of the iteration a child is (role-resolvable). + iteration_key=self.speaker_id_key, + gates=Gates( + requires_gpu=self.resources.requires_gpu, + requires_internet_first_run=True, + writes_to_disk=self.write_to_disk, + output_path_params=["separated_audio_dir"], + # Diarization runs on one file's audio, and ``num_speakers`` counts the speakers + # found in THAT file. Unlike SplitLongAudioStage, a shared output directory is + # still safe here: ``write_audio_stable`` names each per-speaker WAV after a + # digest of its own samples rather than after the source basename. + per_row_independent=True, + ), + ) def setup_on_node(self, _node_info: Any = None, _worker_metadata: Any = None) -> None: # noqa: ANN401 try: @@ -160,6 +251,17 @@ def _initialize_separator(self) -> None: # each child gets its own duration from the diarization result. _INHERITED_DROP_KEYS = frozenset({"audio", "waveform", "duration", "num_samples"}) + def _write_speaker_wav(self, waveform: torch.Tensor, sr: int, original_file: str, speaker_id: str) -> str: + """Write one per-speaker waveform to ``separated_audio_dir`` and return the path.""" + stem = os.path.splitext(os.path.basename(str(original_file)))[0] or "audio" + return write_audio_stable( + waveform, + sr, + output_dir=self.separated_audio_dir, + stem=stem, + tag=str(speaker_id), + ) + def _build_speaker_tasks( self, speaker_audio_data: dict, @@ -174,21 +276,46 @@ def _build_speaker_tasks( logger.debug(f"Skipping {speaker_id}: duration {result.duration:.2f}s < {self.min_duration}s") continue spk_waveform, spk_sr = _pydub_to_waveform_sr(result.audio) + # Drop the parent's file path(s) too: they point at the FULL + # multi-speaker file, so a file-preferring downstream stage would + # process the whole file per speaker instead of this speaker's + # extracted waveform. With the path gone, downstream resolves the + # per-speaker waveform (input_residency="auto") instead. + drop_keys = { + *self._INHERITED_DROP_KEYS, + self.waveform_key, + self.duration_key, + self.audio_filepath_key, + "audio_filepath", + } speaker_data = { - **{k: v for k, v in item.items() if k not in self._INHERITED_DROP_KEYS}, - "waveform": spk_waveform, - "sample_rate": spk_sr, - "speaker_id": speaker_id, - "num_speakers": num_speakers, - "duration": result.duration, - "diar_segments": result.diar_segments, + **{k: v for k, v in item.items() if k not in drop_keys}, + self.speaker_id_key: speaker_id, + self.num_speakers_key: num_speakers, + self.duration_key: result.duration, + self.diar_segments_key: result.diar_segments, + # Source identity must survive the audio_filepath drop above — + # TimestampMapper (and any provenance consumer) reads original_file. + "original_file": item.get("original_file") + or item.get(self.audio_filepath_key) + or item.get("audio_filepath") + or "unknown", } + # Output residency: keep the in-memory waveform (default) and/or persist a + # per-speaker WAV and point audio_filepath at it (opt-in write_to_disk). + if self.keep_waveform_in_task: + speaker_data[self.waveform_key] = spk_waveform + speaker_data[self.sample_rate_key] = spk_sr + if self.write_to_disk: + speaker_data[self.audio_filepath_key] = self._write_speaker_wav( + spk_waveform, spk_sr, speaker_data["original_file"], speaker_id + ) spk_task = AudioTask( data=speaker_data, dataset_name=task.dataset_name, + _metadata=dict(task._metadata or {}), + _stage_perf=list(task._stage_perf), ) - if task._metadata: - spk_task._metadata = dict(task._metadata) results.append(spk_task) return results @@ -208,7 +335,13 @@ def process(self, task: AudioTask) -> list[AudioTask]: results: list[AudioTask] = [] try: - audio_result = resolve_waveform_from_item(item, task.task_id) + audio_result = resolve_audio( + item, + residency=self.input_residency, # type: ignore[arg-type] + audio_filepath_key=self.audio_filepath_key, + waveform_key=self.waveform_key, + sample_rate_key=self.sample_rate_key, + ) if audio_result is None: return [] waveform, sample_rate = audio_result diff --git a/nemo_curator/stages/audio/segmentation/vad_segmentation.py b/nemo_curator/stages/audio/segmentation/vad_segmentation.py index 430691ce2a..264d7adac4 100755 --- a/nemo_curator/stages/audio/segmentation/vad_segmentation.py +++ b/nemo_curator/stages/audio/segmentation/vad_segmentation.py @@ -47,7 +47,8 @@ from silero_vad import get_speech_timestamps, load_silero_vad from nemo_curator.backends.base import WorkerMetadata -from nemo_curator.stages.audio.common import ensure_waveform_2d, load_audio_file +from nemo_curator.stages.audio._agent._agent_ready import AgentReady, Gates, IOSpec, StageContract +from nemo_curator.stages.audio._agent._residency import InputResidency, accepts_for_residency, resolve_audio from nemo_curator.stages.base import ProcessingStage from nemo_curator.stages.resources import Resources from nemo_curator.tasks import AudioTask @@ -57,7 +58,7 @@ @dataclass -class VADSegmentationStage(ProcessingStage[AudioTask, AudioTask]): +class VADSegmentationStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): """ Stage to segment audio using Voice Activity Detection (VAD). @@ -75,6 +76,20 @@ class VADSegmentationStage(ProcessingStage[AudioTask, AudioTask]): speech_pad_ms: Padding in ms to add before/after speech segments. waveform_key: Key to get waveform data. sample_rate_key: Key to get sample rate. + audio_filepath_key: Key in data dict for the input audio file path. + segments_key: Key where the nested segments list is written (nested=True). + start_ms_key: Key where each segment's start time in milliseconds is written. + end_ms_key: Key where each segment's end time in milliseconds is written. + segment_num_key: Key where each segment's index is written. + duration_key: Key where each segment's duration in seconds is written. + original_file_key: Key carrying the source file path for provenance. + nested: If True, return one task with all segment dicts under segments_key + instead of fanning out one task per segment (default False). + input_residency: Which input to use — "waveform" (in-memory only), "file" + (audio_filepath only), or "auto" (waveform first, file fallback; default). + keep_segment_waveform_in_task: If True (default), store each segment's waveform + in the segment item. If False, nested segments are metadata-only — waveform + consumers such as SegmentConcatenation will skip them. Note: Default resources: cpus=1.0, gpus=0.0 (CPU). Silero VAD is lightweight. @@ -86,9 +101,18 @@ class VADSegmentationStage(ProcessingStage[AudioTask, AudioTask]): max_duration_sec: float = 60.0 threshold: float = 0.5 speech_pad_ms: int = 300 + audio_filepath_key: str = "audio_filepath" waveform_key: str = "waveform" sample_rate_key: str = "sample_rate" + segments_key: str = "segments" + start_ms_key: str = "start_ms" + end_ms_key: str = "end_ms" + segment_num_key: str = "segment_num" + duration_key: str = "duration" + original_file_key: str = "original_file" nested: bool = False + input_residency: InputResidency = "auto" + keep_segment_waveform_in_task: bool = True name: str = "VADSegmentation" batch_size: int = 1 @@ -98,12 +122,59 @@ def __post_init__(self): super().__init__() self._vad_model = None self._device = None + if self.nested and not self.keep_segment_waveform_in_task: + logger.warning( + "[VADSegmentation] nested=True with keep_segment_waveform_in_task=False: " + "segments will carry no audio — SegmentConcatenation (and any waveform " + "consumer) will silently drop every segment. Metadata-only use intended?" + ) def inputs(self) -> tuple[list[str], list[str]]: return [], [] def outputs(self) -> tuple[list[str], list[str]]: - return [], ["waveform", "sample_rate", "start_ms", "end_ms", "segment_num", "duration"] + if self.nested: + return [], [self.segments_key] + outputs = [self.sample_rate_key, self.start_ms_key, self.end_ms_key, self.segment_num_key, self.duration_key] + if self.keep_segment_waveform_in_task: + outputs.append(self.waveform_key) + outputs.append(self.original_file_key) + return [], outputs + + def describe(self) -> StageContract: + writes = [ + self.sample_rate_key, + self.start_ms_key, + self.end_ms_key, + self.segment_num_key, + self.duration_key, + self.original_file_key, # _build_segment_item always writes it + ] + produces = [] + if self.keep_segment_waveform_in_task: + writes.append(self.waveform_key) + produces.append("tensor") + if self.nested: + writes = [self.segments_key] + forms = accepts_for_residency(self.input_residency) + reads_one_of = [] + if "waveform" in forms: + reads_one_of.append(IOSpec(data_keys=[self.waveform_key, self.sample_rate_key], accepts=["waveform"])) + if "file" in forms: + reads_one_of.append(IOSpec(data_keys=[self.audio_filepath_key], accepts=["file"])) + return StageContract( + reads_one_of=reads_one_of, + writes=IOSpec(data_keys=writes, produces=produces), + cardinality="1:1 nested-list" if self.nested else "1:N fan-out", + cardinality_options=["fan_out", "nested"], + iteration_key=self.segments_key, + # nested mode pops the top-level waveform after building segments, so a + # downstream stage reading the top-level waveform would find it gone. + removes_keys=[self.waveform_key] if self.nested else [], + # Silero decides speech from this file's own samples against the configured + # threshold, and every segment it emits is a slice of that same file. + gates=Gates(requires_gpu=self.resources.requires_gpu, per_row_independent=True), + ) def ray_stage_spec(self) -> dict[str, Any]: if self.nested: @@ -179,51 +250,43 @@ def _build_segment_item( not in ( self.waveform_key, self.sample_rate_key, - "start_ms", - "end_ms", - "segment_num", - "duration", + self.start_ms_key, + self.end_ms_key, + self.segment_num_key, + self.duration_key, "num_samples", ) } + if not self.keep_segment_waveform_in_task: + segment_waveform = None segment_data.update( { - "waveform": segment_waveform, - "sample_rate": sample_rate, - "start_ms": start_ms, - "end_ms": end_ms, - "segment_num": segment_num, - "duration": (end_ms - start_ms) / 1000.0, - "original_file": item.get("original_file", item.get("audio_filepath", "unknown")), + self.sample_rate_key: sample_rate, + self.start_ms_key: start_ms, + self.end_ms_key: end_ms, + self.segment_num_key: segment_num, + self.duration_key: (end_ms - start_ms) / 1000.0, + self.original_file_key: item.get(self.original_file_key, item.get(self.audio_filepath_key, "unknown")), } ) + if segment_waveform is not None: + segment_data[self.waveform_key] = segment_waveform return segment_data def _resolve_audio(self, item: dict[str, Any]) -> tuple[torch.Tensor, int] | None: """Resolve waveform and sample_rate from task data. Returns None on failure.""" - waveform = item.get(self.waveform_key) - sample_rate = item.get(self.sample_rate_key) - - if waveform 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: # noqa: BLE001 - logger.error(f"Failed to load audio file {audio_filepath}: {e}") - return None - else: - logger.error("Missing waveform and no valid audio_filepath provided") - return None - elif sample_rate is None: - logger.warning("Waveform present but sample_rate missing - task skipped") - return None - - return ensure_waveform_2d(waveform), sample_rate + resolved = resolve_audio( + item, + residency=self.input_residency, # type: ignore[arg-type] + audio_filepath_key=self.audio_filepath_key, + waveform_key=self.waveform_key, + sample_rate_key=self.sample_rate_key, + ) + if resolved is None: + logger.error("Missing waveform/sample_rate and no valid audio path provided") + return resolved - def process(self, task: AudioTask) -> AudioTask | list[AudioTask]: + def process(self, task: AudioTask) -> AudioTask | list[AudioTask]: # noqa: PLR0911 (complexity accepted: one early return per input/error condition) """ Process a single AudioTask. @@ -237,7 +300,11 @@ def process(self, task: AudioTask) -> AudioTask | list[AudioTask]: msg = "VAD model failed to initialize. Cannot process audio." raise RuntimeError(msg) - audio_result = self._resolve_audio(task.data) + try: + audio_result = self._resolve_audio(task.data) + except (OSError, RuntimeError) as e: # corrupt/unreadable audio -> skip the row, don't crash the batch + logger.error(f"Failed to load audio for {task.data.get(self.audio_filepath_key)!r}: {e}") + return [] if audio_result is None: return [] waveform, sample_rate = audio_result @@ -247,11 +314,11 @@ def process(self, task: AudioTask) -> AudioTask | list[AudioTask]: if not segments: logger.warning("No speech segments detected by VAD") if self.nested: - task.data["segments"] = [] + task.data[self.segments_key] = [] return task return [] - original_file = task.data.get("audio_filepath", "unknown") + original_file = task.data.get(self.audio_filepath_key, "unknown") file_name = os.path.basename(original_file) if original_file != "unknown" else task.task_id total_duration = sum((s["end"] - s["start"]) for s in segments) logger.info( @@ -259,11 +326,11 @@ def process(self, task: AudioTask) -> AudioTask | list[AudioTask]: ) if self.nested: - task.data["segments"] = [ + task.data[self.segments_key] = [ self._build_segment_item(task.data, waveform, sample_rate, seg, i) for i, seg in enumerate(segments) ] - del task.data[self.waveform_key] + task.data.pop(self.waveform_key, None) return task output_tasks: list[AudioTask] = [] @@ -272,9 +339,9 @@ def process(self, task: AudioTask) -> AudioTask | list[AudioTask]: seg_task = AudioTask( data=seg_data, dataset_name=task.dataset_name, + _metadata=dict(task._metadata or {}), + _stage_perf=list(task._stage_perf), ) - if task._metadata: - seg_task._metadata = dict(task._metadata) output_tasks.append(seg_task) except Exception as e: # noqa: BLE001 diff --git a/tests/stages/audio/segmentation/test_speaker_separation.py b/tests/stages/audio/segmentation/test_speaker_separation.py index 961d5f6820..94c18e1590 100644 --- a/tests/stages/audio/segmentation/test_speaker_separation.py +++ b/tests/stages/audio/segmentation/test_speaker_separation.py @@ -12,9 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os import pickle +from types import SimpleNamespace +from typing import Any from unittest.mock import MagicMock, patch +import pytest import torch from pydub import AudioSegment @@ -23,6 +27,7 @@ SpeakerResult, SpeakerSeparator, ) +from nemo_curator.stages.resources import Resources from nemo_curator.tasks import AudioTask @@ -87,6 +92,54 @@ def test_process_output_keys(self, mock_init: MagicMock) -> None: assert "waveform" in item assert "sample_rate" in item + # --- output residency (write-to-disk extension) --- + + def test_default_output_is_in_memory_only(self) -> None: + """Regression: default config emits a tensor and sets no disk gate/path.""" + contract = SpeakerSeparationStage().describe() + assert contract.writes.produces == ["tensor"] + assert contract.gates.writes_to_disk is False + assert "audio_filepath" not in contract.writes.data_keys + + @patch("nemo_curator.stages.audio.segmentation.speaker_separation.SpeakerSeparationStage._initialize_separator") + def test_write_to_disk_persists_and_sets_path(self, mock_init: MagicMock, tmp_path) -> None: # noqa: ANN001 + stage = SpeakerSeparationStage(min_duration=0.5, write_to_disk=True, separated_audio_dir=str(tmp_path / "sep")) + separator = MagicMock() + separator.get_speaker_audio_data.return_value = { + "spk_0": SpeakerResult(_make_audio_segment(3000), 3.0, [(0.0, 3.0)]), + } + stage._separator = separator + item = stage.process(_make_task())[0].data + # default keep_waveform_in_task=True -> waveform AND a written per-speaker file + assert "waveform" in item + assert "audio_filepath" in item + assert os.path.exists(item["audio_filepath"]) + + @patch("nemo_curator.stages.audio.segmentation.speaker_separation.SpeakerSeparationStage._initialize_separator") + def test_write_to_disk_only_drops_waveform(self, mock_init: MagicMock, tmp_path) -> None: # noqa: ANN001 + stage = SpeakerSeparationStage( + min_duration=0.5, + write_to_disk=True, + separated_audio_dir=str(tmp_path / "sep"), + keep_waveform_in_task=False, + ) + separator = MagicMock() + separator.get_speaker_audio_data.return_value = { + "spk_0": SpeakerResult(_make_audio_segment(3000), 3.0, [(0.0, 3.0)]), + } + stage._separator = separator + item = stage.process(_make_task())[0].data + assert "waveform" not in item + assert os.path.exists(item["audio_filepath"]) + + def test_requires_dir_when_write_to_disk(self) -> None: + with pytest.raises(ValueError, match="separated_audio_dir"): + SpeakerSeparationStage(write_to_disk=True) + + def test_requires_at_least_one_output_sink(self) -> None: + with pytest.raises(ValueError, match="keep_waveform_in_task or write_to_disk"): + SpeakerSeparationStage(keep_waveform_in_task=False) + @patch("nemo_curator.stages.audio.segmentation.speaker_separation.SpeakerSeparationStage._initialize_separator") def test_min_duration_filters_short_speakers(self, mock_init: MagicMock) -> None: stage = SpeakerSeparationStage(min_duration=2.0) @@ -318,3 +371,57 @@ def test_empty_input(self) -> None: sep = _make_separator() result = sep.exclude_overlapping_segments({}, buffer_time=0.0) assert result == {} + + +class _TinyAudioSegment: + """A pydub-shaped stub: just enough for the separator to hand back audio.""" + + sample_width = 2 + channels = 1 + frame_rate = 16000 + + def get_array_of_samples(self) -> list[int]: + return [0, 500, -500, 0] * 100 + + +# Lifted from tests/stages/audio/test_agent_simulation_pipelines.py: it drives only +# SpeakerSeparationStage, and was the sole coverage of fan-out metadata isolation. +def test_agent_fanout_children_have_isolated_metadata() -> None: + """Fan-out children must own independent _metadata / _stage_perf copies. + + Pins the de-aliasing fix behaviorally: mutating one child must not leak into a + sibling or the parent (the shared-reference bug class). + """ + + def fake_speaker_audio_data(*_args: Any, **_kwargs: Any) -> dict[str, Any]: # noqa: ANN401 + return { + "spk0": SimpleNamespace(audio=_TinyAudioSegment(), duration=0.25, diar_segments=[(0.0, 0.25)]), + "spk1": SimpleNamespace(audio=_TinyAudioSegment(), duration=0.30, diar_segments=[(0.25, 0.55)]), + } + + stage = SpeakerSeparationStage( + input_residency="waveform", + waveform_key="agent_waveform", + sample_rate_key="agent_sr", + min_duration=0.1, + resources=Resources(gpus=0.0), + ) + stage._separator = SimpleNamespace(get_speaker_audio_data=fake_speaker_audio_data) + + parent = AudioTask( + dataset_name="t", + data={"agent_waveform": torch.randn(1, 9600), "agent_sr": 16000}, + _metadata={"trace": "kept"}, + _stage_perf=["fanout-input"], + ) + children = stage.process(parent) + assert len(children) == 2 + + children[0]._metadata["mutated"] = True + children[0]._stage_perf.append("child0-only") + + assert "mutated" not in children[1]._metadata + assert "mutated" not in parent._metadata + assert "child0-only" not in children[1]._stage_perf + assert "child0-only" not in parent._stage_perf + assert children[1]._metadata["trace"] == "kept" diff --git a/tests/stages/audio/segmentation/test_vad_segmentation.py b/tests/stages/audio/segmentation/test_vad_segmentation.py index 2b35544b7d..23046ea9f7 100644 --- a/tests/stages/audio/segmentation/test_vad_segmentation.py +++ b/tests/stages/audio/segmentation/test_vad_segmentation.py @@ -216,3 +216,39 @@ def test_pickling(self) -> None: assert restored.min_duration_sec == 2.0 assert restored.threshold == 0.6 assert restored._vad_model is None + + +class TestNestedAndFanoutAgree: + """The two packagings of a VAD result must describe the same speech. + + Deliberately outside ``TestVADSegmentationStage``, which is marked ``gpu``: this drives no + model -- both timestamps and the loader are patched -- so gating it behind a GPU would mean + the CPU suite never checks that the two modes agree. Lifted from + tests/stages/audio/test_agent_simulation_pipelines.py, its only previous home. + """ + + @patch("nemo_curator.stages.audio.segmentation.vad_segmentation.get_speech_timestamps") + @patch("nemo_curator.stages.audio.segmentation.vad_segmentation.load_silero_vad") + def test_nested_and_fanout_produce_the_same_boundaries( + self, mock_load_vad: MagicMock, mock_get_ts: MagicMock + ) -> None: + sr = 16000 + mock_load_vad.return_value = MagicMock() + mock_get_ts.return_value = [{"start": 0, "end": int(sr * 0.4)}, {"start": int(sr * 0.4), "end": int(sr * 0.9)}] + waveform = torch.randn(1, sr) + + def _task() -> AudioTask: + return AudioTask(dataset_name="t", data={"waveform": waveform.clone(), "sample_rate": sr}) + + nested = VADSegmentationStage(nested=True, input_residency="waveform") + nested.setup() + nested_segments = nested.process(_task()).data["segments"] + + fanout = VADSegmentationStage(nested=False, input_residency="waveform") + fanout.setup() + children = fanout.process(_task()) + + assert len(nested_segments) == len(children) == 2, "both modes must find the same speech" + nested_bounds = [(seg["start_ms"], seg["end_ms"]) for seg in nested_segments] + fanout_bounds = [(child.data["start_ms"], child.data["end_ms"]) for child in children] + assert nested_bounds == fanout_bounds == [(0, 400), (400, 900)], "only the packaging may differ"