Add audio preprocessing stages (MonoConversion, SegmentConcatenation, TimestampMapper) by shubhamNvidia · Pull Request #1575 · NVIDIA-NeMo/Curator

shubhamNvidia · 2026-03-05T14:43:46Z

Adds foundational audio preprocessing stages for the ADV pipeline:
MonoConversionStage: Converts multi-channel audio to mono and verifies sample rate, with strict/non-strict sample rate enforcement.
SegmentConcatenationStage: Concatenates multiple audio segments with configurable silence gaps between them.
Config dataclasses (MonoConversionConfig, SegmentConcatenationConfig) with from_dict/to_dict support.
Unit tests with mocked I/O covering config, stage properties, process logic, and edge cases.

…ing stages

copy-pr-bot · 2026-03-05T14:43:51Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

greptile-apps · 2026-03-05T14:48:05Z

Greptile Summary

This PR adds three foundational audio preprocessing/postprocessing stages (MonoConversionStage, SegmentConcatenationStage, TimestampMapperStage) along with shared utilities in common.py and unit tests. The overall architecture is clean and the previous review concerns around channel-count mismatches, sample-rate mismatches, and test mock scoping have all been addressed.

Two logic issues remain in TimestampMapperStage:

item.get(\"end_ms\", 0) silently drops any task where end_ms is absent and start_ms is 0, because 0 <= 0 triggers the "invalid range" early-return.
_build_output_item_no_mapping can return duration_ms=0 (or negative) without logging a warning when all three duration fallbacks fail.

Confidence Score: 4/5

Safe to merge after addressing the two TimestampMapperStage logic issues; the preprocessing stages themselves are solid.

Two P1 logic bugs in TimestampMapperStage: tasks with start_ms=0 and no end_ms key are silently dropped, and zero/negative durations can be emitted without a warning. These affect correctness on real data where end_ms may be absent or where no duration fallback succeeds. The preprocessing stages (MonoConversion, SegmentConcatenation) are well-implemented and all prior review concerns have been resolved.

nemo_curator/stages/audio/postprocessing/timestamp_mapper.py — two logic issues in process() and _build_output_item_no_mapping.

Important Files Changed

Filename	Overview
nemo_curator/stages/audio/common.py	Adds `load_audio_file`, `ensure_waveform_2d`, `ensure_mono`, `resolve_waveform_from_item`, and `resolve_model_path` utilities; `resolve_waveform_from_item` and `resolve_model_path` are not referenced by any stage in this PR (dead code), and `resolve_waveform_from_item` mutates its `item` argument in-place as documented.
nemo_curator/stages/audio/preprocessing/mono_conversion.py	New stage; converts multi-channel audio to mono and enforces sample-rate policy. Clean, well-tested implementation with proper error handling.
nemo_curator/stages/audio/preprocessing/concatenation.py	New stage; concatenates nested VAD segments with configurable silence gaps and writes segment mappings to `_metadata`. Correctly guards channel-count mismatches and warns on sample-rate mismatches.
nemo_curator/stages/audio/postprocessing/timestamp_mapper.py	New stage; maps concatenated-waveform positions back to original-file timestamps. Two logic issues: tasks with `start_ms=0` and no `end_ms` key are silently dropped due to `end_ms` defaulting to 0, and `_build_output_item_no_mapping` can return `duration_ms=0` (or negative) without logging a warning.
tests/stages/audio/preprocessing/test_mono_conversion.py	Good test coverage; mocks are correctly scoped to the module under test (`MOCK_EXISTS = "...mono_conversion.os.path.exists"`).
tests/stages/audio/preprocessing/test_concatenation.py	Covers basic concatenation, silence insertion, empty/missing inputs, and error paths. No mocking issues.
tests/stages/audio/postprocessing/test_timestamp_mapper.py	Thorough coverage of mapping paths, cross-boundary rejection, passthrough filtering, and edge cases.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[AudioTask\naudio_filepath] --> B[MonoConversionStage\nload file → mono waveform]
    B -->|strict_sample_rate check| C{sample rate OK?}
    C -->|no| D[drop task]
    C -->|yes| E[AudioTask\nwaveform + sample_rate + is_mono]

    E --> F[VADSegmentationStage\nnested=True]
    F --> G[AudioTask\ndata.segments list]

    G --> H[SegmentConcatenationStage\nsort → cat with silence]
    H -->|no valid parts| I[drop task]
    H -->|ok| J[AudioTask\ncombined waveform\n_metadata.segment_mappings]

    J --> K[ASR / Quality\n/ Speaker stages]
    K --> L[AudioTask\nstart_ms, end_ms per segment]

    L --> M[TimestampMapperStage\ntranslate concat coords\nto original file coords]
    M -->|spans 2+ mappings| N[drop task]
    M -->|in silence gap| N
    M -->|ok| O[AudioTask\noriginal_file\noriginal_start/end_ms\nduration_sec]

Comments Outside Diff (2)

nemo_curator/stages/audio/postprocessing/timestamp_mapper.py, line 110-116 (link)

Missing end_ms silently drops valid tasks

item.get("end_ms", 0) defaults to 0 when the key is absent. A task whose start_ms is also 0 (or absent) will hit concat_end <= concat_start (0 <= 0) and be discarded with just a warning, even though start_ms=0 is a perfectly valid segment start. Any task coming in without an end_ms key — including the first segment of a file that genuinely starts at offset 0 — will be silently dropped.

Use a sentinel default (e.g. None) so that the "missing key" case can be distinguished from the "key is explicitly 0" case:
```
concat_start = item.get("start_ms", 0)
concat_end = item.get("end_ms")
if concat_end is None:
    logger.warning(
        f"[TimestampMapper] Skipping task with missing 'end_ms'"
    )
    return []
if concat_end <= concat_start:
    logger.warning(
        f"[TimestampMapper] Skipping task with invalid range: start_ms={concat_start}, end_ms={concat_end}"
    )
    return []
```
nemo_curator/stages/audio/postprocessing/timestamp_mapper.py, line 159-180 (link)

Silent zero/negative duration_ms in no-mapping fallback path

When all three fallbacks fail — end_ms - start_ms <= 0, no duration/duration_sec, and no waveform — the function still returns a dict with duration_ms=0 (or even a negative value if start_ms > end_ms) and duration_sec=0.0, without logging any warning. Downstream consumers that divide by duration or filter on duration_sec > 0 will silently receive corrupt output.

At minimum, add a warning and guard when no fallback succeeds:
```
if duration_ms <= 0:
    logger.warning(
        f"[TimestampMapper] Could not determine a positive duration for task; "
        f"start_ms={start_ms}, end_ms={end_ms}. Emitting duration_ms=0."
    )
```

_{Reviews (29): Last reviewed commit: "Merge branch 'main' into pr/audio-founda..." | Re-trigger Greptile}

sarahyurick · 2026-03-09T22:54:41Z

+
+
+@dataclass
+class SegmentConcatenationConfig:


Discussed offline. We do not need a config class per stage.

sarahyurick · 2026-03-09T22:55:03Z

+
+
+@dataclass
+class MonoConversionConfig:


Discussed offline. We do not need a config class per stage.

sarahyurick · 2026-03-09T22:56:37Z

+                dataset_name=first_task.dataset_name,
+                _metadata=first_task._metadata,
+                _stage_perf=list(first_task._stage_perf),


Discussed offline. We should not use only the first task for this.

sarahyurick · 2026-03-09T22:57:08Z

+            logger.error(f"Error concatenating segments: {e}")
+            return []
+
+    def process(self, task: AudioBatch) -> Optional[AudioBatch]:


Discussed offline. We should not need to implement process if we are implementing process_batch.

sarahyurick · 2026-03-09T23:00:14Z

+        expected = 1.0 + 2.0 + 1.0
+        assert abs(combined_duration_sec - expected) < 0.1
+
+    def test_process_single_delegates_to_batch(self) -> None:


This can be removed.

sarahyurick · 2026-03-09T23:00:24Z

+
+        assert result == []
+
+    def test_preserves_dataset_name(self) -> None:


This can be removed.

sarahyurick · 2026-03-09T23:00:41Z

+    def test_defaults(self) -> None:
+        cfg = MonoConversionConfig()
+        assert cfg.output_sample_rate == 48000
+        assert cfg.audio_filepath_key == "audio_filepath"
+        assert cfg.strict_sample_rate is True
+
+    def test_from_dict(self) -> None:
+        cfg = MonoConversionConfig.from_dict(
+            {"output_sample_rate": 16000, "strict_sample_rate": False}
+        )
+        assert cfg.output_sample_rate == 16000
+        assert cfg.strict_sample_rate is False
+        assert cfg.audio_filepath_key == "audio_filepath"
+
+    def test_from_dict_ignores_unknown_keys(self) -> None:
+        cfg = MonoConversionConfig.from_dict({"unknown_key": 42})
+        assert cfg.output_sample_rate == 48000


These can be removed.

sarahyurick · 2026-03-09T23:00:54Z

+    def test_stage_properties(self) -> None:
+        stage = MonoConversionStage()
+        assert stage.name == "MonoConversion"
+        assert stage.inputs() == (["data"], [])
+        assert stage.outputs() == ([], ["waveform", "sample_rate", "is_mono", "duration", "num_samples"])
+
+    def test_config_overrides_params(self) -> None:
+        cfg = MonoConversionConfig(output_sample_rate=16000, strict_sample_rate=False)
+        stage = MonoConversionStage(config=cfg)
+        assert stage.output_sample_rate == 16000
+        assert stage.strict_sample_rate is False


These can be removed.

sarahyurick · 2026-03-09T23:01:23Z

+
+        assert len(result.data) == 0
+
+    def test_preserves_task_metadata(self, tmp_path: Path) -> None:


This can be removed.

…atenation for canonical waveform format

…catenation

greptile-apps · 2026-03-18T11:58:18Z

+            w = waveform.squeeze() if waveform.dim() > 1 else waveform
+            num_samples = w.shape[-1] if w.dim() > 0 else 0
+            segment_duration_ms = int(1000 * num_samples / sample_rate)
+
+            orig_start = item.get('start_ms', 0)
+            orig_end = item.get('end_ms', 0)
+            if orig_end <= orig_start:
+                orig_end = orig_start + segment_duration_ms
+
+            mapping = SegmentMapping(
+                original_file=item.get('original_file', item.get('audio_filepath', 'unknown')),
+                original_start_ms=orig_start,
+                original_end_ms=orig_end,
+                concat_start_ms=current_pos_ms,
+                concat_end_ms=current_pos_ms + segment_duration_ms,
+                segment_index=idx,
+            )
+            mappings.append(mapping.to_dict())
+
+            parts.append(w.unsqueeze(0) if w.dim() == 1 else w)
+            current_pos_ms += segment_duration_ms
+
+            parts.append(torch.zeros(1, silence_samples, dtype=w.dtype, device=w.device))


Multi-channel waveform causes RuntimeError in torch.cat

waveform.squeeze() (no dimension argument) squeezes all singleton dimensions. For a mono (1, N) tensor it correctly collapses to (N,), but for a multi-channel (C, N) tensor where C > 1, squeeze() is a no-op and w keeps shape (C, N).

Because w.dim() == 2 in that case, line 162 appends w as (C, N). The silence tensor on line 165, however, is always hardcoded to (1, silence_samples). When torch.cat(parts[:-1], dim=-1) is later called, the shapes (C, N) and (1, silence_samples) are incompatible along dim 0 (when C ≠ 1), raising a RuntimeError at runtime.

Even though SegmentConcatenationStage is intended to run after MonoConversionStage, there is no guard enforcing this, so receiving multi-channel data is an entirely reachable error path.

A minimal fix is to validate — or explicitly enforce mono — before appending:

# Ensure waveform is 1D (squeeze the channel dim only if it's 1) if waveform.dim() == 2 and waveform.shape[0] != 1: logger.warning( f"[SegmentConcat] Segment {idx} has {waveform.shape[0]} channels; " "expected mono. Averaging channels." ) waveform = waveform.mean(dim=0, keepdim=True) w = waveform.squeeze(0) # now safely (N,)

greptile-apps · 2026-03-18T11:58:19Z

+                item['waveform'] = mono_waveform
+                item['sample_rate'] = sr
+                item['is_mono'] = True
+                item['duration'] = mono_waveform.shape[1] / sr
+                item['num_samples'] = mono_waveform.shape[1]
+
+                results.append(item)


In-place mutation of input item dict

The stage writes directly into the original item dict from task.data and then appends that same object to results. This means:

After process() returns, the original task.data items already carry the mutated fields (waveform, sample_rate, etc.) — the caller's data is silently modified.

If the same batch is ever replayed or inspected, the item will appear as already-processed even though no copy was made.

Prefer creating a copy of the item to keep the output immutable with respect to the input:

Suggested change

item['waveform'] = mono_waveform

item['sample_rate'] = sr

item['is_mono'] = True

item['duration'] = mono_waveform.shape[1] / sr

item['num_samples'] = mono_waveform.shape[1]

results.append(item)

result_item = {

**item,

'waveform': mono_waveform,

'sample_rate': sr,

'is_mono': True,

'duration': mono_waveform.shape[1] / sr,

'num_samples': mono_waveform.shape[1],

}

results.append(result_item)

greptile-apps · 2026-03-21T10:57:06Z

+        stereo = np.random.randn(48000, 2).astype(np.float32)
+
+        with patch("nemo_curator.stages.audio.preprocessing.mono_conversion.sf.read", return_value=(stereo, 48000)):
+            with patch("os.path.exists", return_value=True):


Global os.path.exists mock leaks across all modules

patch("os.path.exists", return_value=True) replaces the os.path.exists function in the os module itself, meaning every call to os.path.exists from any imported module during the test will also return True. This creates a test-isolation risk — e.g., framework code that legitimately checks for a file's absence would be silently bypassed.

The correct approach is to scope the patch to the module under test:

Suggested change

with patch("os.path.exists", return_value=True):

with patch("nemo_curator.stages.audio.preprocessing.mono_conversion.os.path.exists", return_value=True):

The same global-patch issue applies at lines 55, 71, 85, and 110 — all five patch("os.path.exists", ...) calls should be updated to use the fully qualified module path.

greptile-apps · 2026-03-21T10:57:08Z

+    def outputs(self) -> Tuple[List[str], List[str]]:
+        return [], ["waveform", "sample_rate", "num_segments", "total_duration_sec"]


outputs() declaration is missing "original_file"

The returned declaration ["waveform", "sample_rate", "num_segments", "total_duration_sec"] does not include "original_file", which is unconditionally present in the output dict built at line 198. If the pipeline framework uses outputs() for contract validation or schema registration, downstream stages expecting "original_file" would not discover it through the declared contract.

Suggested change

def outputs(self) -> Tuple[List[str], List[str]]:

return [], ["waveform", "sample_rate", "num_segments", "total_duration_sec"]

def outputs(self) -> Tuple[List[str], List[str]]:

return [], ["waveform", "sample_rate", "num_segments", "total_duration_sec", "original_file"]

…ers in preprocessing - common.py: add load_audio_file, ensure_waveform_2d, ensure_mono, resolve_waveform_from_item, resolve_model_path utility functions - postprocessing/timestamp_mapper.py: new TimestampMapperStage with cross-boundary rejection and passthrough_keys filtering - configs/timestamp_mapper.py: new TimestampMapperConfig - mono_conversion.py: replace inline sf.read with load_audio_file - concatenation.py: replace inline tensor normalization with ensure_waveform_2d from common - Update __init__.py and configs/__init__.py exports - Add unit tests for TimestampMapperStage

sarahyurick

Thanks, kicking off tests and the Ruff formatter as additional feedback.

sarahyurick · 2026-03-23T19:06:37Z

/ok to test c5e326b

- Q000: Replace single quotes with double quotes across all audio stage files - ANN401: Add noqa suppression for intentional Any type in ensure_waveform_2d - BLE001: Replace blind Exception catches with specific exceptions (OSError, RuntimeError, soundfile.SoundFileError)

…nctions

… imports

sarahyurick · 2026-03-31T16:42:20Z

/ok to test 7ed5fc8

sarahyurick · 2026-03-31T16:43:13Z

Hi @shubhamNvidia please update: https://github.com/NVIDIA-NeMo/Curator/actions/runs/23808747074/job/69389685007?pr=1575

sarahyurick · 2026-03-31T18:08:11Z

/ok to test 4130f6b

sarahyurick · 2026-03-31T19:22:10Z

/ok to test be48818

sarahyurick · 2026-04-01T17:54:34Z

/ok to test 2ff712c

shubhamNvidia · 2026-04-01T20:34:50Z

/ok to test 4136b49

shubhamNvidia · 2026-04-02T11:39:10Z

/ok to test ba466fa

shubhamNvidia · 2026-04-02T15:51:03Z

/ok to test 86c2569

sarahyurick · 2026-04-02T18:15:50Z

+        assert result == []
+
+    def test_missing_segments_key_raises(self) -> None:
+        import pytest


Very nit: please move to top-level import.

shubhamNvidia · 2026-04-02T18:54:11Z

/ok to test d7c7c68

sarahyurick

LGTM, thanks!

sarahyurick · 2026-04-02T18:57:30Z

/ok to test 4a43d49

Signed-off-by: shbhawsar <shbhawsar@nvidia.com>

sarahyurick · 2026-04-02T20:01:03Z

/ok to test 3deae25

shubhamNvidia added 2 commits March 5, 2026 02:14

Add audio preprocessing stages and config base

f4345b2

Add unit tests for MonoConversion and SegmentConcatenation preprocess…

b189ac0

…ing stages

github-actions Bot added the community-request label Mar 5, 2026

greptile-apps Bot reviewed Mar 5, 2026

View reviewed changes

Comment thread nemo_curator/stages/audio/preprocessing/mono_conversion.py

sarahyurick self-requested a review March 5, 2026 18:10

ayushdg removed the community-request label Mar 5, 2026

sarahyurick requested changes Mar 9, 2026

View reviewed changes

shubhamNvidia added 2 commits March 11, 2026 03:45

Simplified concatenation and mono_conversion configs, refactored conc…

118a1b6

…atenation for canonical waveform format

Fix mono_conversion unused import and concatenation edge cases

ab9ea66

greptile-apps Bot reviewed Mar 12, 2026

View reviewed changes

Comment thread nemo_curator/stages/audio/configs/mono_conversion.py Outdated

Comment thread nemo_curator/stages/audio/preprocessing/concatenation.py Outdated

github-actions Bot added the community-request label Mar 12, 2026

shubhamNvidia added 2 commits March 18, 2026 04:13

fix: resolve sample_rate from item instead of hardcoding 48000 in con…

88e75a1

…catenation

fix: update copyright, rewrite tests, add sample rate warning

f9369bf

greptile-apps Bot reviewed Mar 18, 2026

View reviewed changes

fix: validate channel consistency in SegmentConcatenationStage

1a229c8

greptile-apps Bot reviewed Mar 21, 2026

View reviewed changes

shubhamNvidia force-pushed the pr/audio-foundation branch from eefedd8 to 46710d6 Compare March 22, 2026 08:00

shubhamNvidia force-pushed the pr/audio-foundation branch from 46710d6 to 0d369c8 Compare March 22, 2026 10:07

sarahyurick reviewed Mar 23, 2026

View reviewed changes

copy-pr-bot Bot temporarily deployed to test March 23, 2026 19:06 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci March 23, 2026 19:07 Failure

shubhamNvidia and others added 6 commits March 25, 2026 19:36

Migrate to AudioTask single-dict pattern

4e33206

Merge upstream/main, resolve conflicts keeping AudioTask + utility fu…

406c5e9

…nctions

Add nested segments handling to SegmentConcatenationStage

b95de7c

Simplify SegmentConcatenationStage for nested-only mode, use absolute…

855c8eb

… imports

Merge branch 'main' into pr/audio-foundation

7ed5fc8

upstream main merged

4130f6b

shubhamNvidia added 2 commits April 1, 2026 19:04

fix ruff linting errors in audio foundation files

94c97cf

Merge upstream/main into pr/audio-foundation

4136b49

Merge branch 'main' into pr/audio-foundation

ba466fa

Merge remote-tracking branch 'upstream/main' into pr/audio-foundation

86c2569

sarahyurick reviewed Apr 2, 2026

View reviewed changes

sarahyurick approved these changes Apr 2, 2026

View reviewed changes

shubhamNvidia and others added 3 commits April 2, 2026 19:01

fix: move pytest import to top-level in test_concatenation.py

d7c7c68

Signed-off-by: shbhawsar <shbhawsar@nvidia.com>

Merge branch 'main' into pr/audio-foundation

4a43d49

Merge branch 'main' into pr/audio-foundation

3deae25

lbliii mentioned this pull request Apr 3, 2026

docs: add fern docs for 26.04 PRs (#1160, #1575, #1576, #1603, #1612, #1652) #1730

Closed

3 tasks

This was referenced Apr 5, 2026

Add SIGMOSFilterStage for audio quality assessment #1577

Merged

Add AudioDataFilterStage composite pipeline for end-to-end audio curation #1640

Merged


		assert result == []

		def test_preserves_dataset_name(self) -> None:


		assert len(result.data) == 0

		def test_preserves_task_metadata(self, tmp_path: Path) -> None:

-                item['waveform'] = mono_waveform
-                item['sample_rate'] = sr
-                item['is_mono'] = True
-                item['duration'] = mono_waveform.shape[1] / sr
-                item['num_samples'] = mono_waveform.shape[1]
-                results.append(item)
+                result_item = {
+                    **item,
+                    'waveform': mono_waveform,
+                    'sample_rate': sr,
+                    'is_mono': True,
+                    'duration': mono_waveform.shape[1] / sr,
+                    'num_samples': mono_waveform.shape[1],
+                }
+                results.append(result_item)

	with patch("os.path.exists", return_value=True):
	with patch("nemo_curator.stages.audio.preprocessing.mono_conversion.os.path.exists", return_value=True):

		def outputs(self) -> Tuple[List[str], List[str]]:
		return [], ["waveform", "sample_rate", "num_segments", "total_duration_sec"]

Uh oh!

Conversation

shubhamNvidia commented Mar 5, 2026

Uh oh!

copy-pr-bot Bot commented Mar 5, 2026

Uh oh!

greptile-apps Bot commented Mar 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 4/5

Important Files Changed

Flowchart

Comments Outside Diff (2)

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

greptile-apps Bot Mar 18, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Mar 18, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

greptile-apps Bot Mar 21, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

greptile-apps Bot Mar 21, 2026

Choose a reason for hiding this comment

Uh oh!

sarahyurick left a comment

Choose a reason for hiding this comment

Uh oh!

sarahyurick commented Mar 23, 2026

Uh oh!

sarahyurick commented Mar 31, 2026

Uh oh!

sarahyurick commented Mar 31, 2026

Uh oh!

sarahyurick commented Mar 31, 2026

Uh oh!

sarahyurick commented Mar 31, 2026

Uh oh!

sarahyurick commented Apr 1, 2026

Uh oh!

shubhamNvidia commented Apr 1, 2026

Uh oh!

shubhamNvidia commented Apr 2, 2026

Uh oh!

shubhamNvidia commented Apr 2, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shubhamNvidia commented Apr 2, 2026

Uh oh!

sarahyurick left a comment

greptile-apps Bot commented Mar 5, 2026 •

edited

Loading