diff --git a/nemo_curator/stages/audio/_agent_ready.py b/nemo_curator/stages/audio/_agent_ready.py new file mode 100644 index 0000000000..efeeddc6be --- /dev/null +++ b/nemo_curator/stages/audio/_agent_ready.py @@ -0,0 +1,318 @@ +# 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 __future__ import annotations + +from dataclasses import asdict, dataclass, field, is_dataclass +from enum import Enum +from typing import TYPE_CHECKING, Any, ClassVar, Literal + +if TYPE_CHECKING: + from collections.abc import Mapping + +AudioForm = Literal["file", "waveform"] +ProducedForm = Literal["tensor", "disk"] +Cardinality = Literal["1:1", "1:1 nested-list", "1:N fan-out", "N:1", "filter"] +Dispatch = Literal["process", "process_batch", "auto"] +# How a stage handles per-item failures at runtime. "unknown" means the stage +# has not declared a uniform policy (the default for most stages today). +ErrorPolicy = Literal["skip", "fail", "annotate", "unknown"] + +# Semantic role vocabulary. Because key *names* are agent-configurable +# (config-knobs-only standardization: every read/written key is a ``*_key`` +# constructor field that an agent may rename), two stages can only be chained +# reliably by matching the *role* a key plays, not its literal string. Roles are +# resolved from the invariant ``*_key`` field name (see ``_roles.py``), so they +# survive an agent renaming a key's value. ``"unknown"`` is the safe default and +# never blocks composition. +Role = Literal[ + "audio_filepath", + "waveform", + "sample_rate", + "duration", + "num_samples", + "segments", + "diar_segments", + "vad_segments", + "overlap_segments", + "timestamps", + "start", + "end", + "start_ms", + "end_ms", + "text", + "pred_text", + "reference_text", + "words", + "alignment", + "speaker_id", + "num_speakers", + "score", + "metrics", + "prediction", + "segment_num", + "original_file", + "item_id", + "windows", + "manifest_path", + "output_path", + "unknown", +] + + +@dataclass(frozen=True) +class ParamSpec: + """A single constructor parameter an agent can set on a stage. + + Usually derived automatically from the stage's dataclass fields (or + ``__init__`` signature) via + :func:`nemo_curator.stages.audio._agent_registry.stage_params`, but a stage + may also override/augment entries in ``StageContract.params``. + """ + + name: str + type: str = "Any" + default: Any = None + required: bool = False + choices: list[Any] | None = None # populated for Literal[...] params + description: str | None = None + role: Role | None = None # semantic role of the key this param configures + + +@dataclass(frozen=True) +class IOSpec: + """Task data keys and audio forms read or written by a stage.""" + + data_keys: list[str] = field(default_factory=list) + segment_data_keys: list[str] = field(default_factory=list) + accepts: list[AudioForm] = field(default_factory=list) + produces: list[ProducedForm] = field(default_factory=list) + + +@dataclass(frozen=True) +class Gates: + """Execution gates or side effects an agent should know before wrapping a stage.""" + + writes_to_disk: bool = False + requires_gpu: bool = False + requires_internet_first_run: bool = False + requires_ffmpeg: bool = False + lifecycle_side_effects: bool = False + runtime_secrets: list[str] = field(default_factory=list) + # Serializability contract for sinks (distinguishes the two JSON sinks): + # requires_serializable_input — this stage serializes task.data as-is (e.g. + # raw json.dumps) and will fail on a resident tensor/audio blob. + # sanitizes_output — this stage strips tensors/audio blobs, so anything + # downstream of it is serialization-safe. + requires_serializable_input: bool = False + sanitizes_output: bool = False + + +@dataclass(frozen=True) +class SizeEnvelope: + """Coarse size and memory hints for agent planning.""" + + max_input_sec: float | None = None + allowed_sample_rates: list[int] | None = None + channels: Literal["mono", "stereo", "any"] = "any" + memory_hint: str | None = None + + +@dataclass(frozen=True) +class StaticHints: + """Instance-independent hints a stage may declare for discovery/planning. + + Lets a stage expose ``cardinality_options``, ``gates``, ``dispatch``, + ``error_policy``, ``description`` and ``stage_id`` *without* being + instantiated (see :meth:`AgentReady.describe_static`). Declared on a class + via the ``AGENT_STATIC`` ClassVar; every field defaults so declaring it is + fully optional and additive. + """ + + cardinality_options: list[str] = field(default_factory=list) + gates: Gates = field(default_factory=Gates) + dispatch: Dispatch = "auto" + error_policy: ErrorPolicy = "unknown" + description: str | None = None + stage_id: str | None = None + + +@dataclass(frozen=True) +class StageContract: + """Read-only discovery contract for an agent-ready processing stage.""" + + reads: IOSpec = field(default_factory=IOSpec) + writes: IOSpec = field(default_factory=IOSpec) + reads_one_of: list[IOSpec] = field(default_factory=list) + metadata_reads: list[str] = field(default_factory=list) + metadata_writes: list[str] = field(default_factory=list) + cardinality: Cardinality = "1:1" + cardinality_options: list[str] = field(default_factory=list) + iteration_key: str | None = None + preserves_upstream_keys: bool = True + wrappable: bool = True + size_envelope: SizeEnvelope = field(default_factory=SizeEnvelope) + gates: Gates = field(default_factory=Gates) + # Agent-facing metadata (advisory; defaults keep older describe() calls valid). + stage_id: str | None = None # stable semantic id; defaults to the class name when None + description: str | None = None # one-line human summary for planners/UIs + params: list[ParamSpec] = field(default_factory=list) # usually auto-derived at discovery time + dispatch: Dispatch = "auto" # "auto" => infer from the stage at runtime + error_policy: ErrorPolicy = "unknown" + # Resolved-key-value -> semantic role. Populated at discovery time by + # ``_agent_registry.build_contract``; empty when a contract is built by hand. + key_roles: dict[str, Role] = field(default_factory=dict) + # True when the stage only implements ``process_batch`` (``process`` raises). + # Auto-derived at discovery time; agents must not call ``process`` on these. + batch_only: bool = False + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-safe dict of this contract (``json.dumps`` never raises).""" + return { + "reads": asdict(self.reads), + "writes": asdict(self.writes), + "reads_one_of": [asdict(spec) for spec in self.reads_one_of], + "metadata_reads": list(self.metadata_reads), + "metadata_writes": list(self.metadata_writes), + "cardinality": self.cardinality, + "cardinality_options": list(self.cardinality_options), + "iteration_key": self.iteration_key, + "preserves_upstream_keys": self.preserves_upstream_keys, + "wrappable": self.wrappable, + "size_envelope": asdict(self.size_envelope), + "gates": asdict(self.gates), + "stage_id": self.stage_id, + "description": self.description, + "params": [_paramspec_to_dict(p) for p in self.params], + "dispatch": self.dispatch, + "error_policy": self.error_policy, + "key_roles": dict(self.key_roles), + "batch_only": self.batch_only, + } + + +def _jsonable_default(value: Any) -> Any: # noqa: ANN401, PLR0911 (complexity accepted: one early return per JSON type family) + """Coerce an arbitrary param default to a JSON-serializable value. + + Non-serializable objects (tensors, models, callables) become a short + ``""`` sentinel rather than raising. + """ + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, (list, tuple, set)): + return [_jsonable_default(v) for v in value] + if isinstance(value, dict): + return {str(k): _jsonable_default(v) for k, v in value.items()} + if isinstance(value, Enum): + return value.value + if is_dataclass(value) and not isinstance(value, type): + try: + return {k: _jsonable_default(v) for k, v in asdict(value).items()} + except Exception: # noqa: BLE001 + return f"" + return f"" + + +def _paramspec_to_dict(p: ParamSpec) -> dict[str, Any]: + return { + "name": p.name, + "type": p.type, + "default": _jsonable_default(p.default), + "required": p.required, + "choices": None if p.choices is None else [_jsonable_default(c) for c in p.choices], + "description": p.description, + "role": p.role, + } + + +def _json_type(type_str: str | None) -> str | None: + """Map a rendered ParamSpec.type string to a JSON-Schema type (or None to omit).""" + if not type_str: + return None + t = type_str.replace(" ", "").replace("|None", "") + if t.startswith("Optional["): + t = t[len("Optional[") : -1] if t.endswith("]") else t + scalar = {"str": "string", "int": "integer", "float": "number", "bool": "boolean"} + if t in scalar: + return scalar[t] + lowered = t.lower() + if lowered.startswith(("list", "sequence", "tuple")): + return "array" + if lowered.startswith(("dict", "mapping")): + return "object" + return None + + +def to_json_schema(params: list[ParamSpec]) -> dict[str, Any]: + """Build a JSON-Schema ``object`` describing a stage's configurable params. + + Suitable as the argument schema for an agent's stage-configuration form. + """ + properties: dict[str, Any] = {} + required: list[str] = [] + for p in params: + schema: dict[str, Any] = {} + json_type = _json_type(p.type) + if json_type is not None: + schema["type"] = json_type + if p.choices is not None: + schema["enum"] = [_jsonable_default(c) for c in p.choices] + if p.default is not None: + schema["default"] = _jsonable_default(p.default) + if p.description: + schema["description"] = p.description + if p.role: + schema["x-role"] = p.role + properties[p.name] = schema + if p.required: + required.append(p.name) + out: dict[str, Any] = {"type": "object", "properties": properties} + if required: + out["required"] = required + return out + + +class AgentReady: + """Mixin for stages that expose a read-only agent discovery contract.""" + + # Opt-in, instance-independent discovery hints. Annotated as ClassVar so + # dataclass stages do NOT treat these as fields. All optional/additive. + AGENT_STATIC: ClassVar[StaticHints | None] = None + # Set True on stages whose ``process`` raises (only ``process_batch`` works). + BATCH_ONLY: ClassVar[bool] = False + # Rare per-field role overrides keyed by ``*_key`` field name. Consulted + # before the shared ``_roles.KEY_ROLES`` table. + KEY_ROLE_OVERRIDES: ClassVar[Mapping[str, Role]] = {} + + def describe(self) -> StageContract: + raise NotImplementedError + + @classmethod + def describe_static(cls) -> StageContract: + """Instance-free contract for discovery/planning. + + Uses class defaults + ``AGENT_STATIC`` and never runs ``__init__`` side + effects, so it is safe for stages with required constructor args. Prefer + the instance-level :meth:`describe` (or + ``_agent_registry.build_contract``) when resolved key *values* are + needed. + """ + from nemo_curator.stages.audio._agent_registry import static_contract + + return static_contract(cls) + + +# (resolve_contract was removed: dead code whose signature promised class +# acceptance the body rejected. Instances: use stage.describe() / build_contract; +# classes: use describe_static() / static_contract.) diff --git a/nemo_curator/stages/audio/_residency.py b/nemo_curator/stages/audio/_residency.py new file mode 100644 index 0000000000..135768b3d2 --- /dev/null +++ b/nemo_curator/stages/audio/_residency.py @@ -0,0 +1,157 @@ +# 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 __future__ import annotations + +import contextlib +import os +import tempfile +from typing import TYPE_CHECKING, Any, Literal + +import soundfile as sf + +from nemo_curator.stages.audio.common import ensure_waveform_2d, load_audio_file + +if TYPE_CHECKING: + from collections.abc import Callable + +InputResidency = Literal["file", "waveform", "auto"] + + +def resolve_audio( # noqa: PLR0913 (complexity accepted: keyword-only residency/key knobs mirror the stage fields) + item: dict[str, Any], + *, + residency: InputResidency = "auto", + audio_filepath_key: str = "audio_filepath", + waveform_key: str = "waveform", + sample_rate_key: str = "sample_rate", + mono: bool = True, + loader: Callable[..., tuple[Any, int]] | None = None, +) -> tuple[Any, int] | None: + """Return ``(waveform_2d, sample_rate)`` from tensor keys or a file path. + + ``auto`` prefers an existing waveform, then falls back to file loading. + ``waveform`` never falls back to disk. ``file`` always loads from the + configured path key. + + ``loader`` overrides the file-loading callable (default + :func:`~nemo_curator.stages.audio.common.load_audio_file`); stages pass + their own module-level symbol so callers can patch it at the stage module. + """ + waveform = item.get(waveform_key) + sample_rate = item.get(sample_rate_key) + if residency != "file" and waveform is not None and sample_rate is not None: + return ensure_waveform_2d(waveform), int(sample_rate) + + if residency == "waveform": + return None + + path = item.get(audio_filepath_key) + if path: + expanded = os.path.expanduser(str(path)) + if os.path.exists(expanded): + return (loader or load_audio_file)(expanded, mono=mono) + return None + + +def _as_soundfile_array(waveform: Any) -> Any: # noqa: ANN401 + waveform = ensure_waveform_2d(waveform) + if hasattr(waveform, "detach"): + waveform = waveform.detach() + if hasattr(waveform, "cpu"): + waveform = waveform.cpu() + if hasattr(waveform, "numpy"): + waveform = waveform.numpy() + if getattr(waveform, "ndim", 0) == 2: # noqa: PLR2004 - 2 == a (channels, samples) 2-D array + channels, samples = waveform.shape + if channels == 1: + return waveform[0] + if channels < samples: + return waveform.T + return waveform + + +def resolve_audio_path( # noqa: PLR0913 (complexity accepted: keyword-only residency/key knobs mirror the stage fields) + item: dict[str, Any], + *, + residency: InputResidency = "auto", + audio_filepath_key: str = "audio_filepath", + waveform_key: str = "waveform", + sample_rate_key: str = "sample_rate", + temp_dir: str | None = None, + register_temp: list[str] | None = None, +) -> str | None: + """Return an audio path, writing a temp WAV when only a waveform exists. + + When a temp WAV is materialized from an in-memory waveform and + ``register_temp`` is provided, the temp path is appended to that list so the + caller can delete it after use (see :func:`cleanup_temp_files`). Without + ``register_temp`` the caller is responsible for cleanup itself. + """ + path = item.get(audio_filepath_key) + local_path: str | None = None + if residency != "waveform" and path: + local_path = os.path.expanduser(str(path)) + if os.path.exists(local_path): + return local_path + # Protocol-prefixed paths (file://, http(s)://, s3://, ...) were handled + # by the stages' own fsspec machinery before the residency layer existed; + # keep accepting them when the target exists remotely. + if "://" in str(path): + try: + from fsspec.core import url_to_fs + + fs, fspath = url_to_fs(str(path)) + if fs.exists(fspath): + return path + except Exception: # noqa: BLE001, S110 - unknown protocol/creds -> deliberate fall-through + pass + + if residency == "file": + # Pre-residency stages handed unverified paths straight to their own + # downstream machinery (ffmpeg/NeMo/fsspec) and let it report the + # failure; keep that contract instead of gating on os.path.exists. + return local_path + + waveform = item.get(waveform_key) + sample_rate = item.get(sample_rate_key) + if waveform is None or sample_rate is None: + return local_path + + fd, tmp = tempfile.mkstemp(suffix=".wav", dir=temp_dir) + os.close(fd) + sf.write(tmp, _as_soundfile_array(waveform), int(sample_rate)) + if register_temp is not None: + register_temp.append(tmp) + return tmp + + +def cleanup_temp_files(paths: list[str] | None) -> None: + """Best-effort removal of temp files created by :func:`resolve_audio_path`.""" + for path in paths or (): + with contextlib.suppress(OSError): + os.remove(path) + + +def produce_audio_filepath( + item: dict[str, Any], + new_path: str, + *, + key: str = "audio_filepath", + original_key: str = "original_audio_filepath", +) -> None: + """Update a canonical audio path while preserving the first prior value.""" + if key in item and original_key not in item: + item[original_key] = item[key] + item[key] = new_path diff --git a/nemo_curator/stages/audio/common.py b/nemo_curator/stages/audio/common.py index a27c1f54bd..4cb80af249 100644 --- a/nemo_curator/stages/audio/common.py +++ b/nemo_curator/stages/audio/common.py @@ -25,6 +25,7 @@ from loguru import logger from nemo_curator.backends.base import NodeInfo, WorkerMetadata +from nemo_curator.stages.audio._agent_ready import AgentReady, Gates, IOSpec, StageContract from nemo_curator.stages.base import CompositeStage, ProcessingStage from nemo_curator.stages.file_partitioning import FilePartitioningStage from nemo_curator.tasks import AudioTask, EmptyTask, FileGroupTask @@ -41,7 +42,7 @@ def get_audio_duration(audio_filepath: str) -> float: @dataclass -class GetAudioDurationStage(ProcessingStage[AudioTask, AudioTask]): +class GetAudioDurationStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): """Compute audio duration from the file at *audio_filepath_key* and store the result under *duration_key*. @@ -65,6 +66,12 @@ def inputs(self) -> tuple[list[str], list[str]]: def outputs(self) -> tuple[list[str], list[str]]: return [], [self.duration_key] + def describe(self) -> StageContract: + return StageContract( + reads=IOSpec(data_keys=[self.audio_filepath_key], accepts=["file"]), + writes=IOSpec(data_keys=[self.duration_key]), + ) + def process(self, task: AudioTask) -> AudioTask: t0 = time.perf_counter() audio_filepath = task.data[self.audio_filepath_key] @@ -74,7 +81,7 @@ def process(self, task: AudioTask) -> AudioTask: return task -class PreserveByValueStage(ProcessingStage[AudioTask, AudioTask]): +class PreserveByValueStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): """Filter entries by comparing *input_value_key* against *target_value*. Returns ``None`` from ``process()`` to drop entries that fail the @@ -87,6 +94,7 @@ class PreserveByValueStage(ProcessingStage[AudioTask, AudioTask]): """ name: str = "PreserveByValueStage" + BATCH_ONLY = True # process() raises; only process_batch is implemented (agent-discovery hint) def __init__( self, @@ -108,6 +116,13 @@ def inputs(self) -> tuple[list[str], list[str]]: def outputs(self) -> tuple[list[str], list[str]]: return [], [self.input_value_key] + def describe(self) -> StageContract: + return StageContract( + reads=IOSpec(data_keys=[self.input_value_key]), + writes=IOSpec(data_keys=[self.input_value_key]), + cardinality="filter", + ) + def process(self, task: AudioTask) -> AudioTask | None: msg = "PreserveByValueStage only supports process_batch" raise NotImplementedError(msg) @@ -133,7 +148,7 @@ def process_batch(self, tasks: list[AudioTask]) -> list[AudioTask]: @dataclass -class ManifestReaderStage(ProcessingStage[FileGroupTask, AudioTask]): +class ManifestReaderStage(AgentReady, ProcessingStage[FileGroupTask, AudioTask]): """Read JSONL manifest files from a FileGroupTask and emit one AudioTask per line. Uses line-by-line streaming via fsspec (no Pandas) to keep memory at ~1x file size. @@ -177,9 +192,16 @@ def ray_stage_spec(self) -> dict[str, Any]: def num_workers(self) -> int | None: return 1 + def describe(self) -> StageContract: + return StageContract( + writes=IOSpec(data_keys=["audio_filepath"]), + cardinality="1:N fan-out", + gates=Gates(lifecycle_side_effects=True), + ) + @dataclass -class ManifestReader(CompositeStage[EmptyTask, AudioTask]): +class ManifestReader(AgentReady, CompositeStage[EmptyTask, AudioTask]): """Composite stage for reading JSONL manifests. Decomposes into: @@ -227,9 +249,12 @@ def get_description(self) -> str: parts.append(f"with target blocksize {self.blocksize}") return ", ".join(parts) + def describe(self) -> StageContract: + return StageContract(cardinality="1:N fan-out", wrappable=False) + @dataclass -class ManifestWriterStage(ProcessingStage[AudioTask, AudioTask]): +class ManifestWriterStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): """Append a single AudioTask to a JSONL manifest file. The output file is truncated once in ``setup()`` (called on the driver) @@ -290,6 +315,18 @@ def process(self, task: AudioTask) -> AudioTask: def num_workers(self) -> int | None: return 1 + def describe(self) -> StageContract: + return StageContract( + gates=Gates( + writes_to_disk=True, + lifecycle_side_effects=True, + # Serializes task.data as-is via json.dumps; a resident tensor + # (e.g. a waveform) will crash it. Route through + # AudioToDocumentStage (which sanitizes) if one may be present. + requires_serializable_input=True, + ), + ) + def load_audio_file(audio_path: str, mono: bool = True) -> tuple[torch.Tensor, int]: """Load audio file and return waveform tensor (channels, samples) and sample rate.""" @@ -327,6 +364,13 @@ def resolve_waveform_from_item( item['audio_filepath'], resolves missing sample_rate from file header. Updates item in-place when loading from file. Returns None if resolution fails. + + .. note:: + The canonical resolver is :func:`nemo_curator.stages.audio._residency.resolve_audio`. + This helper is retained for its unique behavior — reading ``sample_rate`` from the + file header *without* reloading an already-present waveform, and writing the loaded + waveform/sample_rate back into ``item`` — which ``resolve_audio`` does not replicate. + Prefer ``resolve_audio`` in new code. """ waveform = item.get("waveform") sample_rate = item.get("sample_rate") diff --git a/nemo_curator/stages/audio/segmentation/speaker_separation.py b/nemo_curator/stages/audio/segmentation/speaker_separation.py index 06e1434015..3497759f6a 100755 --- a/nemo_curator/stages/audio/segmentation/speaker_separation.py +++ b/nemo_curator/stages/audio/segmentation/speaker_separation.py @@ -45,7 +45,8 @@ from nemo_curator.backends.base import WorkerMetadata from nemo_curator.backends.utils import RayStageSpecKeys -from nemo_curator.stages.audio.common import resolve_waveform_from_item +from nemo_curator.stages.audio._agent_ready import AgentReady, Gates, IOSpec, StageContract +from nemo_curator.stages.audio._residency import resolve_audio 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 @@ -62,7 +63,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. @@ -76,8 +77,21 @@ 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). Note: + Per-speaker child tasks DROP the parent's audio_filepath (it points at the + full multi-speaker file) and carry ``original_file`` for provenance instead; + downstream stages consume the per-speaker waveform. + GPU assignment is handled by the executor via _resources. Use .with_(resources=Resources(gpus=X)) to configure GPU allocation. """ @@ -87,6 +101,14 @@ 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: str = "auto" name: str = "SpeakerSeparation" batch_size: int = 1 @@ -100,7 +122,41 @@ 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"] + return [], [ + self.waveform_key, + self.sample_rate_key, + self.speaker_id_key, + self.num_speakers_key, + self.duration_key, + self.diar_segments_key, + ] + + def describe(self) -> StageContract: + return StageContract( + reads_one_of=[ + IOSpec(data_keys=[self.waveform_key, self.sample_rate_key], accepts=["waveform"]), + IOSpec(data_keys=[self.audio_filepath_key], accepts=["file"]), + ], + writes=IOSpec( + data_keys=[ + self.waveform_key, + self.sample_rate_key, + self.speaker_id_key, + self.num_speakers_key, + self.duration_key, + self.diar_segments_key, + "original_file", + ], + produces=["tensor"], + ), + # 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.gpus > 0, requires_internet_first_run=True), + ) def ray_stage_spec(self) -> dict[str, Any]: return {RayStageSpecKeys.IS_FANOUT_STAGE: True} @@ -178,21 +234,39 @@ 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.waveform_key: spk_waveform, + self.sample_rate_key: spk_sr, + 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", } 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 @@ -212,7 +286,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 d259ead3b9..99f2549bc1 100755 --- a/nemo_curator/stages/audio/segmentation/vad_segmentation.py +++ b/nemo_curator/stages/audio/segmentation/vad_segmentation.py @@ -48,7 +48,8 @@ from nemo_curator.backends.base import WorkerMetadata from nemo_curator.backends.utils import RayStageSpecKeys -from nemo_curator.stages.audio.common import ensure_waveform_2d, load_audio_file +from nemo_curator.stages.audio._agent_ready import AgentReady, Gates, IOSpec, StageContract +from nemo_curator.stages.audio._residency import resolve_audio from nemo_curator.stages.base import ProcessingStage from nemo_curator.stages.resources import Resources from nemo_curator.tasks import AudioTask @@ -58,7 +59,7 @@ @dataclass -class VADSegmentationStage(ProcessingStage[AudioTask, AudioTask]): +class VADSegmentationStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): """ Stage to segment audio using Voice Activity Detection (VAD). @@ -76,6 +77,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. @@ -87,9 +102,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: str = "auto" + keep_segment_waveform_in_task: bool = True name: str = "VADSegmentation" batch_size: int = 1 @@ -99,12 +123,51 @@ 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] + return StageContract( + reads_one_of=[ + IOSpec(data_keys=[self.waveform_key, self.sample_rate_key], accepts=["waveform"]), + IOSpec(data_keys=[self.audio_filepath_key], accepts=["file"]), + ], + 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, + gates=Gates(requires_gpu=self.resources.gpus > 0), + ) def ray_stage_spec(self) -> dict[str, Any]: if self.nested: @@ -180,51 +243,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. @@ -238,7 +293,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 @@ -248,11 +307,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( @@ -260,11 +319,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] = [] @@ -273,9 +332,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