From d019ecf8bdf13baefd5430b89df9ac780790eb9d Mon Sep 17 00:00:00 2001 From: shbhawsar Date: Tue, 7 Jul 2026 06:30:38 -0700 Subject: [PATCH 1/2] =?UTF-8?q?feat(audio):=20agent-ready=20foundation=20?= =?UTF-8?q?=E2=80=94=20AgentReady=20contract=20mixin=20+=20residency=20res?= =?UTF-8?q?olver?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Shared base imported by the audio stage modules to make them agent-ready: - _agent_ready.py: AgentReady mixin, StageContract/IOSpec/Gates/Role - _residency.py: input residency resolver (file/waveform/auto) - common.py: agent-ready helpers (ensure_mono/ensure_waveform_2d) Excludes the agent orchestration layer (registry/catalog/conformance/planning/agent). --- nemo_curator/stages/audio/_agent_ready.py | 318 ++++++++++++++++++++++ nemo_curator/stages/audio/_residency.py | 157 +++++++++++ nemo_curator/stages/audio/common.py | 54 +++- 3 files changed, 524 insertions(+), 5 deletions(-) create mode 100644 nemo_curator/stages/audio/_agent_ready.py create mode 100644 nemo_curator/stages/audio/_residency.py 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") From ef6640df63509ae87a1f47e345f45d15ef3488f7 Mon Sep 17 00:00:00 2001 From: shbhawsar Date: Tue, 7 Jul 2026 06:32:41 -0700 Subject: [PATCH 2/2] feat(audio-tagging): make tagging stages agent-ready (describe/StageContract + residency) --- .../audio/tagging/inference/nemo_asr_align.py | 42 +++-- .../tagging/merge_alignment_diarization.py | 19 ++- .../audio/tagging/prepare_module_segments.py | 31 ++-- .../stages/audio/tagging/resample_audio.py | 112 +++++++++++--- nemo_curator/stages/audio/tagging/split.py | 144 ++++++++++++------ .../audio/tagging/text/chinese_conversion.py | 19 ++- nemo_curator/stages/audio/tagging/text/itn.py | 19 ++- 7 files changed, 282 insertions(+), 104 deletions(-) diff --git a/nemo_curator/stages/audio/tagging/inference/nemo_asr_align.py b/nemo_curator/stages/audio/tagging/inference/nemo_asr_align.py index 7d28122f39..bb894f5af8 100644 --- a/nemo_curator/stages/audio/tagging/inference/nemo_asr_align.py +++ b/nemo_curator/stages/audio/tagging/inference/nemo_asr_align.py @@ -35,13 +35,14 @@ from nemo.collections.asr.parts.submodules.rnnt_decoding import RNNTDecodingConfig 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 ProcessingStage from nemo_curator.stages.resources import Resources from nemo_curator.tasks import AudioTask @dataclass -class BaseASRProcessorStage(ProcessingStage[AudioTask, AudioTask]): +class BaseASRProcessorStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): """Base class for ASR stages with shared config and segment preparation. Provides common fields and _prepare_segment_batch_with_metadata for @@ -72,9 +73,14 @@ class BaseASRProcessorStage(ProcessingStage[AudioTask, AudioTask]): # Output keys text_key: str = "text" words_key: str = "words" + alignment_key: str = "alignment" compute_timestamps: bool = True segments_key: str = "segments" + audio_filepath_key: str = "audio_filepath" + resampled_audio_filepath_key: str = "resampled_audio_filepath" + split_filepaths_key: str = "split_filepaths" + split_metadata_key: str = "split_metadata" # Stage metadata (subclasses can override) name: str = "BaseASRProcessor" @@ -110,7 +116,7 @@ def _prepare_segment_batch_with_metadata( if cut_audio_segments: for metadata_idx, metadata in enumerate(metadata_batch): - audio_path = metadata.get("resampled_audio_filepath", metadata.get("audio_filepath")) + audio_path = metadata.get(self.resampled_audio_filepath_key, metadata.get(self.audio_filepath_key)) if not audio_path: continue audio, sr = torchaudio.load(audio_path) @@ -131,10 +137,10 @@ def _prepare_segment_batch_with_metadata( else: for metadata_idx, metadata in enumerate(metadata_batch): for segment_idx, segment in enumerate(metadata.get(segments_key, [])): - if "resampled_audio_filepath" in segment: + if self.resampled_audio_filepath_key in segment: segment_metadata_list.append( { - "resampled_audio_filepath": segment["resampled_audio_filepath"], + self.resampled_audio_filepath_key: segment[self.resampled_audio_filepath_key], "metadata_idx": metadata_idx, "segment_idx": segment_idx, } @@ -188,6 +194,9 @@ class NeMoASRAlignerStage(BaseASRProcessorStage): # input keys segments_key: str = "segments" + split_filepaths_key: str = "split_filepaths" + split_metadata_key: str = "split_metadata" + alignment_key: str = "alignment" # Output keys text_key: str = "text" @@ -263,10 +272,17 @@ def setup(self, _: WorkerMetadata | None = None) -> None: logger.info(f"[{self.name}] Initialized ASR model on {self._device}") def inputs(self) -> tuple[list[str], list[str]]: - return ["data"], ["duration", self.segments_key, "split_filepaths", "split_metadata"] + return ["data"], ["duration", self.segments_key, self.split_filepaths_key, self.split_metadata_key] def outputs(self) -> tuple[list[str], list[str]]: - return ["data"], ["duration", self.segments_key, "split_filepaths", "split_metadata"] + return ["data"], ["duration", self.segments_key, self.split_filepaths_key, self.split_metadata_key] + + def describe(self) -> StageContract: + return StageContract( + reads=IOSpec(data_keys=["duration", self.segments_key, self.split_filepaths_key, self.split_metadata_key]), + writes=IOSpec(data_keys=[self.text_key, self.alignment_key], segment_data_keys=[self.text_key, self.words_key]), + gates=Gates(requires_gpu=self.resources.gpus > 0, requires_internet_first_run=self.model_path is None), + ) def get_alignments_text(self, hypotheses: Any) -> tuple[list, str]: # noqa: ANN401 """Extract word alignments and text from model hypotheses.""" @@ -338,7 +354,7 @@ def process_full_audio(self, tasks: list[AudioTask]) -> list[AudioTask]: # noqa skip_indices = [] meta_indices = [] for i, data in enumerate(entries): - split_filepaths = data.get("split_filepaths") + split_filepaths = data.get(self.split_filepaths_key) has_splits = isinstance(split_filepaths, list) and len(split_filepaths) > 0 if has_splits or split_filepaths is None: meta_indices.append(i) @@ -347,14 +363,14 @@ def process_full_audio(self, tasks: list[AudioTask]) -> list[AudioTask]: # noqa for i in skip_indices: entries[i][self.text_key] = "" - entries[i]["alignment"] = [] + entries[i][self.alignment_key] = [] # collect all split paths of all entries in the batch all_paths = [] path_to_entry_and_split = [] for entry_idx in meta_indices: meta_entry = entries[entry_idx] - split_filepaths = meta_entry.get("split_filepaths") + split_filepaths = meta_entry.get(self.split_filepaths_key) if not split_filepaths: logger.warning(f"[{self.name}] Entry at index {entry_idx} has no split_filepaths, skipping.") continue @@ -396,13 +412,13 @@ def process_full_audio(self, tasks: list[AudioTask]) -> list[AudioTask]: # noqa else: alignments, text = [], "" - split_metadata = meta_entry.get("split_metadata") + split_metadata = meta_entry.get(self.split_metadata_key) if split_metadata and split_idx < len(split_metadata): split_metadata[split_idx][self.text_key] = text - split_metadata[split_idx]["alignment"] = alignments + split_metadata[split_idx][self.alignment_key] = alignments else: meta_entry[self.text_key] = text - meta_entry["alignment"] = alignments + meta_entry[self.alignment_key] = alignments return tasks @@ -426,7 +442,7 @@ def process_segments(self, tasks: list[AudioTask]) -> list[AudioTask]: with torch.no_grad(): hypotheses_list = self._asr_model.transcribe(all_segments, override_config=self._override_cfg) except Exception as e: - files_list = [x.get("resampled_audio_filepath", x.get("audio_filepath")) for x in entries] + files_list = [x.get(self.resampled_audio_filepath_key, x.get(self.audio_filepath_key)) for x in entries] msg = f"[{self.name}] Exception for audio list: {files_list}, error: {e}" raise ValueError(msg) from e diff --git a/nemo_curator/stages/audio/tagging/merge_alignment_diarization.py b/nemo_curator/stages/audio/tagging/merge_alignment_diarization.py index 986c7fba6f..6029af7535 100644 --- a/nemo_curator/stages/audio/tagging/merge_alignment_diarization.py +++ b/nemo_curator/stages/audio/tagging/merge_alignment_diarization.py @@ -21,12 +21,13 @@ from loguru import logger +from nemo_curator.stages.audio._agent_ready import AgentReady, IOSpec, StageContract from nemo_curator.stages.base import ProcessingStage from nemo_curator.tasks import AudioTask @dataclass -class MergeAlignmentDiarizationStage(ProcessingStage[AudioTask, AudioTask]): +class MergeAlignmentDiarizationStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): """ Stage that merges alignment and diarization information. @@ -51,15 +52,23 @@ class MergeAlignmentDiarizationStage(ProcessingStage[AudioTask, AudioTask]): # Output keys text_key: str = "text" words_key: str = "words" + alignment_key: str = "alignment" + segments_key: str = "segments" # Stage metadata name: str = "MergeAlignmentDiarization" def inputs(self) -> tuple[list[str], list[str]]: - return [], ["alignment", "segments"] + return [], [self.alignment_key, self.segments_key] def outputs(self) -> tuple[list[str], list[str]]: - return [], ["alignment", "segments"] + return [], [self.alignment_key, self.segments_key] + + def describe(self) -> StageContract: + return StageContract( + reads=IOSpec(data_keys=[self.alignment_key, self.segments_key]), + writes=IOSpec(segment_data_keys=[self.text_key, self.words_key]), + ) @staticmethod def align_words_to_segments( @@ -183,8 +192,8 @@ def process(self, task: AudioTask) -> AudioTask: """Process entry to merge alignment and diarization.""" t0 = time.perf_counter() data_entry = task.data - alignment = data_entry.get("alignment", []) - segments = data_entry.get("segments", []) + alignment = data_entry.get(self.alignment_key, []) + segments = data_entry.get(self.segments_key, []) if alignment and segments: self.align_words_to_segments(alignment, segments, self.text_key, self.words_key) diff --git a/nemo_curator/stages/audio/tagging/prepare_module_segments.py b/nemo_curator/stages/audio/tagging/prepare_module_segments.py index 4fef91e261..b8afd3851c 100644 --- a/nemo_curator/stages/audio/tagging/prepare_module_segments.py +++ b/nemo_curator/stages/audio/tagging/prepare_module_segments.py @@ -24,6 +24,7 @@ from loguru import logger +from nemo_curator.stages.audio._agent_ready import AgentReady, IOSpec, StageContract from nemo_curator.stages.base import ProcessingStage from nemo_curator.tasks import AudioTask @@ -32,7 +33,7 @@ @dataclass -class PrepareModuleSegmentsStage(ProcessingStage[AudioTask, AudioTask]): +class PrepareModuleSegmentsStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): """ Stage that prepares segments for TTS or ASR by merging and splitting based on duration, punctuation, and bandwidth. @@ -65,11 +66,23 @@ class PrepareModuleSegmentsStage(ProcessingStage[AudioTask, AudioTask]): terminal_punct_marks: str = ".!?。？？！。" # noqa: RUF001 full_utterance_ratio: float = 1.0 punctuation_split_only: bool = False + segments_key: str = "segments" + duration_key: str = "duration" + metrics_key: str = "metrics" name: str = "PrepareModuleSegments" def inputs(self) -> tuple[list[str], list[str]]: - return [], ["segments", "duration"] + return [], [self.segments_key, self.duration_key] + + def outputs(self) -> tuple[list[str], list[str]]: + return [], [self.segments_key] + + def describe(self) -> StageContract: + return StageContract( + reads=IOSpec(data_keys=[self.segments_key, self.duration_key]), + writes=IOSpec(data_keys=[self.segments_key]), + ) def __post_init__(self): if self.module not in ("tts", "asr"): @@ -94,8 +107,8 @@ def get_words_list_from_all_segments(self, metadata: dict[str, Any]) -> list[dic - sisdr_squim: The SI-SDR score of the word if available - bandwidth: The bandwidth of the word if available """ - segments = metadata["segments"] - audio_duration = metadata.get("duration", 0.0) + segments = metadata[self.segments_key] + audio_duration = metadata.get(self.duration_key, 0.0) if "overlap_segments" not in metadata: add_non_speaker_segments(segments, audio_duration) @@ -111,8 +124,8 @@ def get_words_list_from_all_segments(self, metadata: dict[str, Any]) -> list[dic for word in segment[self.words_key]: new_word = dict(word) new_word["speaker"] = segment["speaker"] - if "metrics" in segment: - m = segment["metrics"] + if self.metrics_key in segment: + m = segment[self.metrics_key] new_word["stoi_squim"] = m.get("stoi_squim") if isinstance(m, dict) else None new_word["sisdr_squim"] = m.get("sisdr_squim") if isinstance(m, dict) else None new_word["pesq_squim"] = m.get("pesq_squim") if isinstance(m, dict) else None @@ -339,7 +352,7 @@ def add_new_segments_to_metadata(self, metadata: dict[str, Any], new_segments: l {"word": w.get("word", ""), "start": w.get("start", 0.0), "end": w.get("end", 0.0)} for w in new_segment["words"] ], - "metrics": { + self.metrics_key: { "pesq_squim": [w.get("pesq_squim") for w in new_segment["words"]], "stoi_squim": [w.get("stoi_squim") for w in new_segment["words"]], "sisdr_squim": [w.get("sisdr_squim") for w in new_segment["words"]], @@ -348,7 +361,7 @@ def add_new_segments_to_metadata(self, metadata: dict[str, Any], new_segments: l } segments.append(seg) - metadata["segments"] = segments + metadata[self.segments_key] = segments def prepare_asr_segments(self, words: list[dict[str, Any]], metadata: dict[str, Any]) -> None: """Prepare ASR segments (multi-speaker per segment allowed).""" @@ -406,7 +419,7 @@ def process(self, task: AudioTask) -> AudioTask: seed = int(hashlib.md5(entry_id.encode()).hexdigest()[:8], 16) # noqa: S324 self._rng.seed(seed) try: - if "segments" not in data_entry: + if self.segments_key not in data_entry: logger.info(f"[{self.name}] No segments in metadata for: {data_entry.get('audio_filepath', '')}") return task diff --git a/nemo_curator/stages/audio/tagging/resample_audio.py b/nemo_curator/stages/audio/tagging/resample_audio.py index eb83505d41..4b803afb37 100644 --- a/nemo_curator/stages/audio/tagging/resample_audio.py +++ b/nemo_curator/stages/audio/tagging/resample_audio.py @@ -25,19 +25,22 @@ import os import shutil import subprocess +import tempfile import time from dataclasses import dataclass from fsspec.core import url_to_fs +from nemo_curator.stages.audio._agent_ready import AgentReady, Gates, IOSpec, StageContract +from nemo_curator.stages.audio._residency import cleanup_temp_files, produce_audio_filepath, resolve_audio_path from nemo_curator.backends.base import NodeInfo, WorkerMetadata -from nemo_curator.stages.audio.common import get_audio_duration +from nemo_curator.stages.audio.common import get_audio_duration, load_audio_file from nemo_curator.stages.base import ProcessingStage from nemo_curator.tasks import AudioTask @dataclass -class ResampleAudioStage(ProcessingStage[AudioTask, AudioTask]): +class ResampleAudioStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): """ Stage for resampling audio files in a TTS/ALM dataset. @@ -57,12 +60,25 @@ class ResampleAudioStage(ProcessingStage[AudioTask, AudioTask]): # Key names audio_filepath_key: str = "audio_filepath" resampled_audio_filepath_key: str = "resampled_audio_filepath" + waveform_key: str = "waveform" + sample_rate_key: str = "sample_rate" duration_key: str = "duration" audio_item_id_key: str = "audio_item_id" + original_audio_filepath_key: str = "original_audio_filepath" + + input_residency: str = "file" + keep_waveform_in_task: bool = False + write_to_disk: bool = True + update_audio_filepath: bool = False # Stage metadata name: str = "ResampleAudio" + def __post_init__(self) -> 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) + def setup_on_node( self, _node_info: NodeInfo | None = None, _worker_metadata: WorkerMetadata | None = None ) -> None: @@ -76,12 +92,31 @@ def inputs(self) -> tuple[list[str], list[str]]: return [], [self.audio_filepath_key] def outputs(self) -> tuple[list[str], list[str]]: - return [], [ - self.audio_filepath_key, - self.audio_item_id_key, - self.resampled_audio_filepath_key, - self.duration_key, - ] + outputs = [self.audio_item_id_key, self.duration_key] + if self.write_to_disk: + outputs.append(self.resampled_audio_filepath_key) + if self.keep_waveform_in_task: + outputs.extend([self.waveform_key, self.sample_rate_key]) + if self.update_audio_filepath: + outputs.append(self.audio_filepath_key) + return [], outputs + + def describe(self) -> StageContract: + writes = [self.audio_item_id_key, self.duration_key] + produces = [] + if self.write_to_disk: + writes.append(self.resampled_audio_filepath_key) + produces.append("disk") + if self.keep_waveform_in_task: + writes.extend([self.waveform_key, self.sample_rate_key]) + produces.append("tensor") + if self.update_audio_filepath: + writes.append(self.audio_filepath_key) + return StageContract( + reads=IOSpec(data_keys=[self.audio_filepath_key], accepts=["file", "waveform"]), + writes=IOSpec(data_keys=writes, produces=produces), + gates=Gates(writes_to_disk=self.write_to_disk, requires_ffmpeg=True), + ) def process(self, task: AudioTask) -> AudioTask: """ @@ -96,29 +131,42 @@ def process(self, task: AudioTask) -> AudioTask: t0 = time.perf_counter() data_entry = task.data - if self.audio_filepath_key not in data_entry: - msg = "Absolute audio filepath is required" + temp_paths: list[str] = [] + input_audio_path = resolve_audio_path( + data_entry, + 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, + register_temp=temp_paths, + ) + if input_audio_path is None: + msg = "Audio file path or waveform/sample_rate is required" raise ValueError(msg) - original_audio_filepath = data_entry[self.audio_filepath_key] - _, local_audio_path = url_to_fs(original_audio_filepath) + original_audio_filepath = data_entry.get(self.audio_filepath_key) + _, local_audio_path = url_to_fs(input_audio_path) if self.audio_item_id_key not in data_entry: stem = os.path.splitext(os.path.basename(local_audio_path))[0] path_hash = hashlib.sha256(local_audio_path.encode()).hexdigest()[:8] data_entry[self.audio_item_id_key] = f"{stem}_{path_hash}" - input_audio_path = local_audio_path - output_audio_path = os.path.join( - self.resampled_audio_dir, - data_entry[self.audio_item_id_key] + "." + self.target_format, - ) + if self.write_to_disk: + output_audio_path = os.path.join( + self.resampled_audio_dir, + data_entry[self.audio_item_id_key] + "." + self.target_format, + ) + else: + fd, output_audio_path = tempfile.mkstemp(suffix=f".{self.target_format}") + os.close(fd) # Convert audio file if not already done fs, output_path = url_to_fs(output_audio_path) - skipped_conversion = fs.exists(output_path) + skipped_conversion = self.write_to_disk and fs.exists(output_path) if not skipped_conversion: cmd = [ "ffmpeg", + "-y", "-v", "error", "-i", @@ -135,14 +183,36 @@ def process(self, task: AudioTask) -> AudioTask: try: subprocess.run(cmd, check=True, capture_output=True, text=True) # noqa: S603 except subprocess.CalledProcessError as e: + cleanup_temp_files(temp_paths) msg = f"Error converting {input_audio_path}: {e}" raise RuntimeError(msg) from e - # Update metadata — preserve original URL for cloud paths - data_entry[self.audio_filepath_key] = original_audio_filepath - data_entry[self.resampled_audio_filepath_key] = output_audio_path + # Input temp WAV (materialized from a waveform) is no longer needed after conversion. + cleanup_temp_files(temp_paths) + + # Update metadata — preserve original URL for cloud paths. + if original_audio_filepath is not None: + data_entry[self.audio_filepath_key] = original_audio_filepath + if self.write_to_disk: + data_entry[self.resampled_audio_filepath_key] = output_audio_path + if self.update_audio_filepath: + produce_audio_filepath( + data_entry, + output_audio_path, + key=self.audio_filepath_key, + original_key=self.original_audio_filepath_key, + ) + if self.keep_waveform_in_task: + waveform, sample_rate = load_audio_file(output_audio_path, mono=False) + data_entry[self.waveform_key] = waveform + data_entry[self.sample_rate_key] = sample_rate duration = get_audio_duration(output_audio_path) data_entry[self.duration_key] = duration + if not self.write_to_disk: + try: + os.remove(output_audio_path) + except OSError: + pass self._log_metrics( { diff --git a/nemo_curator/stages/audio/tagging/split.py b/nemo_curator/stages/audio/tagging/split.py index 27a9b0ff68..42a84e7011 100644 --- a/nemo_curator/stages/audio/tagging/split.py +++ b/nemo_curator/stages/audio/tagging/split.py @@ -25,13 +25,14 @@ from fsspec.core import url_to_fs from loguru import logger +from nemo_curator.stages.audio._agent_ready import AgentReady, Gates, IOSpec, StageContract from nemo_curator.stages.audio.tagging.inference.nemo_asr_align import NeMoASRAlignerStage from nemo_curator.stages.base import CompositeStage, ProcessingStage from nemo_curator.tasks import AudioTask @dataclass -class SplitLongAudioStage(ProcessingStage[AudioTask, AudioTask]): +class SplitLongAudioStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): """ Stage that splits long audio files into smaller segments. @@ -46,31 +47,56 @@ class SplitLongAudioStage(ProcessingStage[AudioTask, AudioTask]): # Split parameters suggested_max_len: float = 3600.0 min_len: float = 1.0 + duration_key: str = "duration" + segments_key: str = "segments" + audio_filepath_key: str = "resampled_audio_filepath" + audio_item_id_key: str = "audio_item_id" + split_filepaths_key: str = "split_filepaths" + split_metadata_key: str = "split_metadata" + split_offsets_key: str = "split_offsets" + split_timestamps_key: str = "split_timestamps" # Stage metadata name: str = "SplitLongAudio" def inputs(self) -> tuple[list[str], list[str]]: - return [], ["duration", "segments", "resampled_audio_filepath"] + return [], [self.duration_key, self.segments_key, self.audio_filepath_key] def outputs(self) -> tuple[list[str], list[str]]: return [], [ - "duration", - "segments", - "resampled_audio_filepath", - "split_filepaths", - "split_metadata", - "split_offsets", - "split_timestamps", + self.duration_key, + self.segments_key, + self.audio_filepath_key, + self.split_filepaths_key, + self.split_metadata_key, + self.split_offsets_key, + self.split_timestamps_key, ] + def describe(self) -> StageContract: + return StageContract( + reads=IOSpec(data_keys=[self.duration_key, self.segments_key, self.audio_filepath_key]), + writes=IOSpec( + data_keys=[ + self.split_filepaths_key, + self.split_metadata_key, + self.split_offsets_key, + self.split_timestamps_key, + ], + produces=["disk"], + ), + cardinality="1:1 nested-list", + iteration_key=self.split_metadata_key, + gates=Gates(writes_to_disk=True), + ) + def get_split_points(self, metadata: dict) -> list[float]: """Get the split points for the audio file based on segments.""" splits = [] split_start = 0 prev_end = 0 - segments = sorted(metadata.get("segments", []), key=lambda s: s.get("start", 0)) + segments = sorted(metadata.get(self.segments_key, []), key=lambda s: s.get("start", 0)) for segment in segments: end = segment.get("end", 0) @@ -90,25 +116,25 @@ def process(self, task: AudioTask) -> AudioTask: def _do_split(self, task: AudioTask) -> AudioTask: """Core splitting logic, separated to keep statement count within limits.""" data_entry = task.data - duration = data_entry["duration"] + duration = data_entry[self.duration_key] if duration < self.suggested_max_len: - data_entry["split_filepaths"] = [data_entry["resampled_audio_filepath"]] - data_entry["split_metadata"] = [ + data_entry[self.split_filepaths_key] = [data_entry[self.audio_filepath_key]] + data_entry[self.split_metadata_key] = [ { - "audio_item_id": data_entry.get("audio_item_id", "unknown"), - "resampled_audio_filepath": data_entry["resampled_audio_filepath"], - "duration": duration, + self.audio_item_id_key: data_entry.get(self.audio_item_id_key, "unknown"), + self.audio_filepath_key: data_entry[self.audio_filepath_key], + self.duration_key: duration, } ] - data_entry["split_offsets"] = [0.0] - data_entry["split_timestamps"] = [0.0] + data_entry[self.split_offsets_key] = [0.0] + data_entry[self.split_timestamps_key] = [0.0] self._log_metrics({"input_duration": duration, "splits_produced": 1}) return task splits = self.get_split_points(data_entry) - audio_path = data_entry["resampled_audio_filepath"] + audio_path = data_entry[self.audio_filepath_key] _fs, resolved_path = url_to_fs(audio_path) # parent_url preserves protocol prefix (e.g. "s3://bucket/dir") for stored paths; @@ -147,7 +173,7 @@ def _do_split(self, task: AudioTask) -> AudioTask: split_durations.append(remaining_frames / sr) actual_splits.append(split_start / sr) - audio_item_id, split_filepaths_before = data_entry.get("audio_item_id", "unknown"), bool(split_filepaths) + audio_item_id, split_filepaths_before = data_entry.get(self.audio_item_id_key, "unknown"), bool(split_filepaths) if not split_filepaths: logger.warning( @@ -159,20 +185,20 @@ def _do_split(self, task: AudioTask) -> AudioTask: split_durations = [duration] actual_splits = [0.0] - data_entry["split_metadata"] = self._build_split_metadata( + data_entry[self.split_metadata_key] = self._build_split_metadata( audio_item_id, split_filepaths, split_durations, fallback=not split_filepaths_before, ) - data_entry["split_filepaths"] = split_filepaths - data_entry["split_offsets"] = actual_splits - data_entry["split_timestamps"] = splits + data_entry[self.split_filepaths_key] = split_filepaths + data_entry[self.split_offsets_key] = actual_splits + data_entry[self.split_timestamps_key] = splits self._log_metrics({"input_duration": duration, "splits_produced": len(split_filepaths)}) return task - @staticmethod def _build_split_metadata( + self, audio_item_id: str, split_filepaths: list[str], split_durations: list[float], @@ -183,23 +209,23 @@ def _build_split_metadata( if fallback: return [ { - "audio_item_id": audio_item_id, - "resampled_audio_filepath": split_filepaths[0], - "duration": split_durations[0], + self.audio_item_id_key: audio_item_id, + self.audio_filepath_key: split_filepaths[0], + self.duration_key: split_durations[0], } ] return [ { - "audio_item_id": f"{audio_item_id}_{idx}", - "resampled_audio_filepath": path, - "duration": split_durations[idx], + self.audio_item_id_key: f"{audio_item_id}_{idx}", + self.audio_filepath_key: path, + self.duration_key: split_durations[idx], } for idx, path in enumerate(split_filepaths) ] @dataclass -class JoinSplitAudioMetadataStage(ProcessingStage[AudioTask, AudioTask]): +class JoinSplitAudioMetadataStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): """ Stage for joining metadata of previously split audio files. @@ -213,13 +239,36 @@ class JoinSplitAudioMetadataStage(ProcessingStage[AudioTask, AudioTask]): """ text_key: str = "text" + split_filepaths_key: str = "split_filepaths" + split_metadata_key: str = "split_metadata" + split_offsets_key: str = "split_offsets" + split_timestamps_key: str = "split_timestamps" + alignment_key: str = "alignment" name: str = "JoinSplitAudioMetadata" def inputs(self) -> tuple[list[str], list[str]]: - return [], ["split_filepaths", "split_metadata", "split_offsets", "split_timestamps"] + return [], [ + self.split_filepaths_key, + self.split_metadata_key, + self.split_offsets_key, + self.split_timestamps_key, + ] def outputs(self) -> tuple[list[str], list[str]]: - return [], [self.text_key, "alignment"] + return [], [self.text_key, self.alignment_key] + + def describe(self) -> StageContract: + return StageContract( + reads=IOSpec( + data_keys=[ + self.split_filepaths_key, + self.split_metadata_key, + self.split_offsets_key, + self.split_timestamps_key, + ] + ), + writes=IOSpec(data_keys=[self.text_key, self.alignment_key]), + ) def process(self, task: AudioTask) -> AudioTask: """ @@ -234,13 +283,13 @@ def process(self, task: AudioTask) -> AudioTask: words_aligned = 0 # Check if this is a meta-entry with split information - if "split_filepaths" in data_entry: - if data_entry["split_filepaths"] is None: - del data_entry["split_filepaths"] + if self.split_filepaths_key in data_entry: + if data_entry[self.split_filepaths_key] is None: + del data_entry[self.split_filepaths_key] else: - splits_joined = len(data_entry.get("split_metadata", [])) + splits_joined = len(data_entry.get(self.split_metadata_key, [])) self._join_split_metadata(data_entry) - words_aligned = len(data_entry.get("alignment", [])) + words_aligned = len(data_entry.get(self.alignment_key, [])) self._log_metrics( { @@ -253,11 +302,11 @@ def process(self, task: AudioTask) -> AudioTask: def _join_split_metadata(self, meta_entry: dict) -> None: """Join metadata from split audio files.""" - split_metadata = meta_entry.get("split_metadata", []) - split_offsets = meta_entry.get("split_offsets", []) + split_metadata = meta_entry.get(self.split_metadata_key, []) + split_offsets = meta_entry.get(self.split_offsets_key, []) if not split_metadata: - del meta_entry["split_filepaths"] + del meta_entry[self.split_filepaths_key] return transcripts = [] @@ -269,7 +318,7 @@ def _join_split_metadata(self, meta_entry: dict) -> None: if text: transcripts.append(text) - alignment = split_entry.get("alignment", []) + alignment = split_entry.get(self.alignment_key, []) offset = split_offsets[idx] if idx < len(split_offsets) else 0 for word in alignment: @@ -280,15 +329,15 @@ def _join_split_metadata(self, meta_entry: dict) -> None: # Create joined entry meta_entry[self.text_key] = " ".join(transcripts) - meta_entry["alignment"] = alignments + meta_entry[self.alignment_key] = alignments # Remove split-related fields - for key in ["split_filepaths", "split_metadata"]: + for key in [self.split_filepaths_key, self.split_metadata_key]: meta_entry.pop(key, None) @dataclass -class SplitASRAlignJoinStage(CompositeStage[AudioTask, AudioTask]): +class SplitASRAlignJoinStage(AgentReady, CompositeStage[AudioTask, AudioTask]): """Composite stage: Split long audio -> ASR align -> Join results. Decomposes into three sequential stages that always run together: @@ -353,6 +402,9 @@ class SplitASRAlignJoinStage(CompositeStage[AudioTask, AudioTask]): def __post_init__(self) -> None: super().__init__() + def describe(self) -> StageContract: + return StageContract(wrappable=False) + def decompose(self) -> list[ProcessingStage]: return [ SplitLongAudioStage( diff --git a/nemo_curator/stages/audio/tagging/text/chinese_conversion.py b/nemo_curator/stages/audio/tagging/text/chinese_conversion.py index ea07008bad..3d0971ac71 100644 --- a/nemo_curator/stages/audio/tagging/text/chinese_conversion.py +++ b/nemo_curator/stages/audio/tagging/text/chinese_conversion.py @@ -21,12 +21,13 @@ from opencc import OpenCC from nemo_curator.backends.base import WorkerMetadata +from nemo_curator.stages.audio._agent_ready import AgentReady, IOSpec, StageContract from nemo_curator.stages.base import ProcessingStage from nemo_curator.tasks import AudioTask @dataclass -class ChineseConversionStage(ProcessingStage[AudioTask, AudioTask]): +class ChineseConversionStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): """Convert Traditional Chinese text to Simplified Chinese (or other OpenCC conversions). Iterates over the ``segments`` list of each entry and writes the converted @@ -40,6 +41,8 @@ class ChineseConversionStage(ProcessingStage[AudioTask, AudioTask]): text_key: str = "text" convert_type: str = "t2s" + segments_key: str = "segments" + output_suffix: str = "_simplified" # Stage metadata name: str = "ChineseConversion" @@ -48,10 +51,16 @@ class ChineseConversionStage(ProcessingStage[AudioTask, AudioTask]): _converter: Any = field(default=None, repr=False) def inputs(self) -> tuple[list[str], list[str]]: - return [], ["segments"] + return [], [self.segments_key] def outputs(self) -> tuple[list[str], list[str]]: - return [], ["segments"] + return [], [self.segments_key] + + def describe(self) -> StageContract: + return StageContract( + reads=IOSpec(data_keys=[self.segments_key]), + writes=IOSpec(segment_data_keys=[f"{self.text_key}{self.output_suffix}"]), + ) def setup(self, _worker_metadata: WorkerMetadata | None = None) -> None: """Setup stage.""" @@ -61,8 +70,8 @@ def setup(self, _worker_metadata: WorkerMetadata | None = None) -> None: def process(self, task: AudioTask) -> AudioTask: data_entry = task.data - output_key = f"{self.text_key}_simplified" - for segment in data_entry.get("segments", []): + output_key = f"{self.text_key}{self.output_suffix}" + for segment in data_entry.get(self.segments_key, []): if self.text_key in segment: try: segment[output_key] = self._converter.convert(segment[self.text_key]) diff --git a/nemo_curator/stages/audio/tagging/text/itn.py b/nemo_curator/stages/audio/tagging/text/itn.py index 1b0863b5fc..65efb2eb37 100644 --- a/nemo_curator/stages/audio/tagging/text/itn.py +++ b/nemo_curator/stages/audio/tagging/text/itn.py @@ -23,12 +23,13 @@ ) from nemo_curator.backends.base import WorkerMetadata +from nemo_curator.stages.audio._agent_ready import AgentReady, IOSpec, StageContract from nemo_curator.stages.base import ProcessingStage from nemo_curator.tasks import AudioTask @dataclass -class InverseTextNormalizationStage(ProcessingStage[AudioTask, AudioTask]): +class InverseTextNormalizationStage(AgentReady, ProcessingStage[AudioTask, AudioTask]): """ Stage that performs inverse text normalization on text data. @@ -45,6 +46,8 @@ class InverseTextNormalizationStage(ProcessingStage[AudioTask, AudioTask]): # Text key text_key: str = "text" + segments_key: str = "segments" + output_suffix: str = "_ITN" # Stage metadata name: str = "InverseTextNormalization" @@ -52,10 +55,16 @@ class InverseTextNormalizationStage(ProcessingStage[AudioTask, AudioTask]): _normalizer: Any = field(default=None, repr=False) def inputs(self) -> tuple[list[str], list[str]]: - return [], ["segments"] + return [], [self.segments_key] def outputs(self) -> tuple[list[str], list[str]]: - return [], ["segments"] + return [], [self.segments_key] + + def describe(self) -> StageContract: + return StageContract( + reads=IOSpec(data_keys=[self.segments_key]), + writes=IOSpec(segment_data_keys=[f"{self.text_key}{self.output_suffix}"]), + ) def setup(self, _worker_metadata: WorkerMetadata | None = None) -> None: """Load the inverse normalizer once per worker.""" @@ -66,13 +75,13 @@ def setup(self, _worker_metadata: WorkerMetadata | None = None) -> None: def process(self, task: AudioTask) -> AudioTask: """Process entry for inverse text normalization.""" data_entry = task.data - segments = data_entry.get("segments", []) + segments = data_entry.get(self.segments_key, []) for segment in segments: if self.text_key in segment: text = segment[self.text_key] if text: sentences = self._normalizer.split_text_into_sentences(text) text_itn = " ".join(self._normalizer.normalize_list(sentences)) - segment[f"{self.text_key}_ITN"] = text_itn + segment[f"{self.text_key}{self.output_suffix}"] = text_itn return task