feat(audio-tagging): output control + in-memory handoff + agent-ready contracts

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._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,
                             }
@@ -190,6 +196,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"
@@ -268,13 +277,36 @@ def setup(self, _: WorkerMetadata | None = None) -> None:
 
     def inputs(self) -> tuple[list[str], list[str]]:
         if self.infer_segment_only:
-            return ["data"], ["resampled_audio_filepath", self.segments_key]
-        return ["data"], ["duration", self.segments_key, "split_filepaths", "split_metadata"]
+            return ["data"], [self.resampled_audio_filepath_key, self.segments_key]
+        return ["data"], ["duration", self.segments_key, self.split_filepaths_key, self.split_metadata_key]
 
     def outputs(self) -> tuple[list[str], list[str]]:
         if self.infer_segment_only:
-            return ["data"], ["resampled_audio_filepath", self.segments_key]
-        return ["data"], ["duration", self.segments_key, "split_filepaths", "split_metadata"]
+            return ["data"], [self.resampled_audio_filepath_key, self.segments_key]
+        return ["data"], ["duration", self.segments_key, self.split_filepaths_key, self.split_metadata_key]
+
+    def describe(self) -> StageContract:
+        if self.infer_segment_only:
+            reads = IOSpec(data_keys=[self.resampled_audio_filepath_key, self.segments_key])
+            writes = IOSpec(segment_data_keys=[self.text_key, self.words_key])
+        else:
+            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],
+            )
+        return StageContract(
+            reads=reads,
+            writes=writes,
+            gates=Gates(
+                requires_gpu=self.resources.requires_gpu,
+                requires_internet_first_run=self.model_path is None,
+                # Batched for GPU throughput, but each row is transcribed and aligned on its own.
+                per_row_independent=True,
+            ),
+        )
 
     def get_alignments_text(self, hypotheses: Any) -> tuple[list, str]:  # noqa: ANN401
         """Extract word alignments and text from model hypotheses."""
@@ -346,7 +378,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)
@@ -355,14 +387,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
@@ -404,13 +436,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
 
@@ -434,7 +466,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
 

nemo_curator/stages/audio/tagging/merge_alignment_diarization.py

@@ -21,12 +21,13 @@
 
 from loguru import logger
 
+from nemo_curator.stages.audio._agent._agent_ready import AgentReady, Gates, 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,25 @@ 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]),
+            # Merges this row's own alignment with this row's own segments.
+            gates=Gates(per_row_independent=True),
+        )
 
     @staticmethod
     def align_words_to_segments(
@@ -183,8 +194,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)

nemo_curator/stages/audio/tagging/prepare_module_segments.py

@@ -24,6 +24,7 @@
 
 from loguru import logger
 
+from nemo_curator.stages.audio._agent._agent_ready import AgentReady, Gates, 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,28 @@ 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]),
+            # The ``asr`` module draws its per-segment length limit from ``self._rng``, which
+            # ``process`` reseeds from a hash of the row's own id before touching it. So the draws
+            # a row gets depend on that row alone, unlike PyAnnoteDiarizationStage's unseeded
+            # generator, whose sequence position is decided by how many rows preceded it.
+            gates=Gates(per_row_independent=True),
+        )
 
     def __post_init__(self):
         if self.module not in ("tts", "asr"):
@@ -94,8 +112,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 +129,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
@@ -345,7 +363,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"]],
@@ -354,7 +372,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)."""
@@ -412,7 +430,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