nemo_curator.stages.audio.alm.pretrain.extraction

Audio extraction stage: slice, mono-resample, write into a tar shard.

Module Contents

Classes

Name	Description
`SnippetExtractionStage`	Slice the source audio per snippet plan, mono-resample, and write into a tar.

API

class nemo_curator.stages.audio.alm.pretrain.extraction.SnippetExtractionStage(
    output_dir: str,
    output_audio_tar_path: str,
    target_sample_rate: int = 16000,
    output_format: str = 'flac',
    audio_filepath_key: str = 'audio_filepath',
    dry_run: bool = False,
    name: str = 'SnippetExtraction',
    batch_size: int = 1,
    resources: nemo_curator.stages.resources.Resources = (lambda: Resources(cpus=1.0...
)

Dataclass

Bases: ProcessingStage[AudioTask, AudioTask]

Slice the source audio per snippet plan, mono-resample, and write into a tar.

For each planned snippet:

Read just the slice [start, end] from the source file.
Channel-average to mono if the source has > 1 channel.
Resample to target_sample_rate using torchaudio if the source rate differs.
Encode the mono waveform in-memory (via soundfile to a BytesIO) and append it as <snippet_id>.<output_format> to this replica’s tar shard (output_audio_tar_path.shard-...); all replicas’ shards are merged into output_audio_tar_path by :func:finalize_audio_pretrain_outputs.
Emit one AudioTask per snippet with the source row’s metadata carried over (minus alignment), the new snippet_id, audio_filepath set to the tar-internal basename (<snippet_id>.<output_format>), updated duration, and segments relativized to the snippet start.

The tar-internal basename matches webdataset / Energon convention: sample key is <snippet_id> (everything before the first .), extension is <output_format>. make_snippet_id already avoids . characters so the snippet id never spuriously splits.

If the input produced zero snippets, a single “stub” AudioTask is emitted (snippet_id=None, no audio written) so that per-original metrics can still flow to the aggregator.

Dry-run mode (dry_run=True): skips steps 1-4 entirely (no soundfile reads, no resampling, no tar writes — not even a tar shard is opened), and step 5 uses the planned end - start as the snippet duration instead of the post-resample frame count. The emitted audio_filepath still uses the basename form for parity with real runs. Useful for previewing the manifest and metrics on real data before committing to a full run.

audio_filepath_key

str = 'audio_filepath'

batch_size

int = 1

dry_run

bool = False

name

str = 'SnippetExtraction'

output_audio_tar_path

str

output_dir

str

output_format

str = 'flac'

resources

Resources

target_sample_rate

int = 16000

nemo_curator.stages.audio.alm.pretrain.extraction.SnippetExtractionStage.__post_init__() -> None

nemo_curator.stages.audio.alm.pretrain.extraction.SnippetExtractionStage._dry_run_emit(
    task: nemo_curator.tasks.AudioTask,
    plan: list[dict],
    original_id: str
) -> list[nemo_curator.tasks.AudioTask]

Emit snippet metadata only, without reading or writing audio.

audio_filepath is the tar-internal basename <snippet_id>.<output_format> for parity with real runs — the tar itself is not opened in dry-run. Snippet duration is the planned end - start (vs. the resampled-frame-count duration the real path would compute — the difference is at most one frame at target_sample_rate).

nemo_curator.stages.audio.alm.pretrain.extraction.SnippetExtractionStage._extract_emit(
    task: nemo_curator.tasks.AudioTask,
    plan: list[dict],
    original_id: str
) -> list[nemo_curator.tasks.AudioTask]

nemo_curator.stages.audio.alm.pretrain.extraction.SnippetExtractionStage._extract_one_snippet(
    task: nemo_curator.tasks.AudioTask,
    snippet: dict,
    source_path: str,
    info: typing.Any,
    original_id: str
) -> nemo_curator.tasks.AudioTask | None

nemo_curator.stages.audio.alm.pretrain.extraction.SnippetExtractionStage._make_snippet_task(
    task: nemo_curator.tasks.AudioTask,
    snippet: dict,
    snippet_id: str,
    out_path: str,
    duration: float
) -> nemo_curator.tasks.AudioTask

nemo_curator.stages.audio.alm.pretrain.extraction.SnippetExtractionStage._make_stub_task(
    task: nemo_curator.tasks.AudioTask
) -> nemo_curator.tasks.AudioTask

nemo_curator.stages.audio.alm.pretrain.extraction.SnippetExtractionStage.inputs() -> tuple[list[str], list[str]]

nemo_curator.stages.audio.alm.pretrain.extraction.SnippetExtractionStage.outputs() -> tuple[list[str], list[str]]

nemo_curator.stages.audio.alm.pretrain.extraction.SnippetExtractionStage.process(
    task: nemo_curator.tasks.AudioTask
) -> list[nemo_curator.tasks.AudioTask]

nemo_curator.stages.audio.alm.pretrain.extraction.SnippetExtractionStage.setup(
    _worker_metadata: nemo_curator.backends.base.WorkerMetadata | None = None
) -> None

nemo_curator.stages.audio.alm.pretrain.extraction.SnippetExtractionStage.setup_on_node(
    _node_info: nemo_curator.backends.base.NodeInfo | None = None,
    _worker_metadata: nemo_curator.backends.base.WorkerMetadata | None = None
) -> None

nemo_curator.stages.audio.alm.pretrain.extraction.SnippetExtractionStage.teardown() -> None