nemo_curator.stages.audio.alm.pretrain.planning

Segment-level planning: overlap drop, greedy packing, repetition filter.

These three stages all operate on a task’s segments / _snippet_plan in memory, before the extractor reads any audio. The pure helpers (filter_empty_segments, find_overlapping_indices, plan_snippets, relativize_segments, n-gram counters, color highlighting) are unit- testable without Ray / soundfile / torch.

Module Contents

Classes

Name	Description
`OverlapFilterStage`	Drop empty segments and overlapping segment pairs.
`SnippetCutPlannerStage`	Compute snippet cut boundaries for one input audio.
`SnippetRepetitionFilterStage`	Drop planned snippets whose text shows suspicious n-gram repetition.

Functions

Name	Description
`_count_ngrams`	Count contiguous n-gram frequencies in a token id sequence.
`_find_offending_ngrams`	Return n-grams whose frequency strictly exceeds `max_count`.
`_format_red`	Wrap each char range in loguru `<red>...</red>` markup.
`_locate_ngram_char_ranges`	Char-range spans for every position where an offending n-gram starts.
`_merge_char_ranges`	Merge overlapping or touching char ranges; input may be unsorted.
`filter_empty_segments`	Drop segments with no text and no words.
`find_overlapping_indices`	Indices of segments that overlap any other segment.
`plan_snippets`	Greedy contiguous packing of segments into snippets.
`relativize_segments`	Return shallow-copied segments with timestamps shifted to snippet-relative.

Data

_MIN_SEGMENTS_FOR_OVERLAP

API

class nemo_curator.stages.audio.alm.pretrain.planning.OverlapFilterStage(
    min_overlap_sec: float = 0.5,
    name: str = 'OverlapFilter',
    batch_size: int = 1,
    resources: nemo_curator.stages.resources.Resources = (lambda: Resources(cpus=1.0...
)

Dataclass

Bases: ProcessingStage[AudioTask, AudioTask]

Drop empty segments and overlapping segment pairs.

First filters segments that have neither text nor words. Then drops every segment that overlaps any other surviving segment, where “overlap” means intersection ≥ min_overlap_sec OR one fully contains the other. Both members of an overlapping pair are discarded — this version keeps no overlap-resolution heuristic.

Per-original counters are stamped onto task._metadata under the pretrain_long_form key so the final aggregator can build a per-original metrics breakdown.

batch_size

int = 1

min_overlap_sec

float = 0.5

name

str = 'OverlapFilter'

resources

Resources

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

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

nemo_curator.stages.audio.alm.pretrain.planning.OverlapFilterStage.process(
    task: nemo_curator.tasks.AudioTask
) -> nemo_curator.tasks.AudioTask

class nemo_curator.stages.audio.alm.pretrain.planning.SnippetCutPlannerStage(
    max_duration_sec: float = 600.0,
    min_duration_sec: float = 0.5,
    max_segment_gap_in_snippet: float = 30.0,
    name: str = 'SnippetCutPlanner',
    batch_size: int = 1,
    resources: nemo_curator.stages.resources.Resources = (lambda: Resources(cpus=1.0...
)

Dataclass

Bases: ProcessingStage[AudioTask, AudioTask]

Compute snippet cut boundaries for one input audio.

Pure planning — no audio I/O. Produces a list of snippet specs each holding start, end (absolute seconds in the source audio) and the contained segments. The plan is stored under task.data["_snippet_plan"] for the downstream extractor to act on. Drop counts (too_long, too_short, no_text) are written to task._metadata['pretrain_long_form'].

batch_size

int = 1

max_duration_sec

float = 600.0

max_segment_gap_in_snippet

float = 30.0

min_duration_sec

float = 0.5

name

str = 'SnippetCutPlanner'

resources

Resources

nemo_curator.stages.audio.alm.pretrain.planning.SnippetCutPlannerStage.__post_init__() -> None

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

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

nemo_curator.stages.audio.alm.pretrain.planning.SnippetCutPlannerStage.process(
    task: nemo_curator.tasks.AudioTask
) -> nemo_curator.tasks.AudioTask

class nemo_curator.stages.audio.alm.pretrain.planning.SnippetRepetitionFilterStage(
    tokenizer_path: str,
    ngram_n: int = 10,
    ngram_max_count: int = 3,
    cache_dir: str | None = None,
    hf_token: str | None = None,
    name: str = 'SnippetRepetitionFilter',
    batch_size: int = 1,
    resources: nemo_curator.stages.resources.Resources = (lambda: Resources(cpus=1.0...
)

Dataclass

Bases: ProcessingStage[AudioTask, AudioTask]

Drop planned snippets whose text shows suspicious n-gram repetition.

Whisper-style ASR sometimes degenerates into repeating the same short phrase for many seconds; the resulting transcript looks fine locally but contains the same n-gram of token ids dozens of times. Such snippets are unsuitable for pretraining.

For every planned snippet (read from task.data["_snippet_plan"]) we join the segment text fields with the same formula the extractor uses, tokenize with the configured HuggingFace fast tokenizer, count n-gram frequencies over the resulting token-id sequence, and drop the snippet if any n-gram appears strictly more than ngram_max_count times. Filtered snippets are logged with the offending occurrences highlighted in red (loguru color tags).

Snippets whose tokenized text has fewer than ngram_n tokens are kept unchanged (no n-grams to evaluate; the planner already enforces a minimum-duration threshold).

Sits between :class:SnippetCutPlannerStage and :class:SnippetExtractionStage so filtered snippets never incur audio decode / resample / file-write cost.

tokenizer_path is either a local directory loadable by AutoTokenizer.from_pretrained or a HuggingFace Hub repository id (e.g. openai/whisper-large-v3). When it’s a repo id, the tokenizer is fetched once per node in :meth:setup_on_node so workers in :meth:setup only ever read from the local cache.

batch_size

int = 1

cache_dir

str | None = None

hf_token

str | None = None

name

str = 'SnippetRepetitionFilter'

ngram_max_count

int = 3

ngram_n

int = 10

resources

Resources

tokenizer_path

str

nemo_curator.stages.audio.alm.pretrain.planning.SnippetRepetitionFilterStage.__post_init__() -> None

nemo_curator.stages.audio.alm.pretrain.planning.SnippetRepetitionFilterStage._snippet_is_repetitive(
    text: str,
    snippet: dict,
    task_id: str
) -> bool

Tokenize text and decide whether to drop the snippet.

On drop, emit a colorized warning showing the offending n-gram occurrences highlighted in red.

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

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

nemo_curator.stages.audio.alm.pretrain.planning.SnippetRepetitionFilterStage.process(
    task: nemo_curator.tasks.AudioTask
) -> nemo_curator.tasks.AudioTask

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

nemo_curator.stages.audio.alm.pretrain.planning.SnippetRepetitionFilterStage.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.planning._count_ngrams(
    token_ids: list[int],
    n: int
) -> collections.Counter[tuple[int, ...]]

Count contiguous n-gram frequencies in a token id sequence.

nemo_curator.stages.audio.alm.pretrain.planning._find_offending_ngrams(
    counts: collections.Counter[tuple[int, ...]],
    max_count: int
) -> set[tuple[int, ...]]

Return n-grams whose frequency strictly exceeds max_count.

nemo_curator.stages.audio.alm.pretrain.planning._format_red(
    text: str,
    ranges: list[tuple[int, int]]
) -> str

Wrap each char range in loguru <red>...</red> markup.

Literal < in the surrounding text is escaped to \< so loguru’s tag parser leaves it alone. ranges must be merged and sorted (use :func:_merge_char_ranges).

nemo_curator.stages.audio.alm.pretrain.planning._locate_ngram_char_ranges(
    token_ids: list[int],
    offsets: list[tuple[int, int]],
    offending: set[tuple[int, ...]],
    n: int
) -> list[tuple[int, int]]

Char-range spans for every position where an offending n-gram starts.

nemo_curator.stages.audio.alm.pretrain.planning._merge_char_ranges(
    ranges: list[tuple[int, int]]
) -> list[tuple[int, int]]

Merge overlapping or touching char ranges; input may be unsorted.

nemo_curator.stages.audio.alm.pretrain.planning.filter_empty_segments(
    segments: list[dict]
) -> tuple[list[dict], int]

Drop segments with no text and no words.

Returns (kept, dropped_count). Order is preserved.

nemo_curator.stages.audio.alm.pretrain.planning.find_overlapping_indices(
    segments: list[dict],
    min_overlap_sec: float
) -> set[int]

Indices of segments that overlap any other segment.

Two segments are considered overlapping (and both indices are returned) iff they share at least min_overlap_sec seconds of intersection OR one fully contains the other. Brief touch-ups smaller than min_overlap_sec where neither covers the other are not flagged.

Implementation is a sweep-line scan over segments sorted by (start, end). An end-time-keyed min-heap holds the currently active intervals (those whose end is still beyond the cursor’s start); each new segment evicts the heap prefix it can no longer intersect and is then compared only against the survivors. For typical diarized audio (a handful of overlapping speakers at any instant) this is effectively O(n log n), vs the pairwise O(n^2) of comparing every pair; the worst case where all intervals overlap each other is still O(n^2) because the overlap relation itself is dense in that case.

nemo_curator.stages.audio.alm.pretrain.planning.plan_snippets(
    segments: list[dict],
    max_duration_sec: float,
    min_duration_sec: float,
    max_segment_gap_in_snippet: float
) -> tuple[list[dict], dict[str, int]]

Greedy contiguous packing of segments into snippets.

Walks segments (assumed sorted by start) and grows a current snippet while:

its span [first.start, last.end] stays within max_duration_sec, AND
the gap from the last accepted segment’s end to the next segment’s start is at most max_segment_gap_in_snippet.

Either constraint failing closes the current snippet and opens a new one with the current segment. Single segments longer than max_duration_sec are emitted as a one-segment candidate and then dropped under too_long.

The gap constraint matters for ALM pretraining: two segments separated by a long silence often belong to semantically distinct conversations (e.g. a topic change, an ad break, two takes recorded back to back), and a snippet that bridges them would teach the model to associate unrelated content. Closing the snippet at long gaps keeps each training example semantically coherent.

Returns (snippets, drop_counts) where each snippet is a dict with keys start, end, segments (the actual segment dicts) and drop counts keys are too_long, too_short, no_text.

Precondition: segments must be non-overlapping (sorted by start with each end <= next.start). OverlapFilterStage guarantees this upstream in the pipeline. If overlapping segments are passed in, gap becomes negative and the gap constraint is silently bypassed, grouping content that should belong to separate snippets.

nemo_curator.stages.audio.alm.pretrain.planning.relativize_segments(
    segments: list[dict],
    snippet_start: float,
    snippet_end: float
) -> list[dict]

Return shallow-copied segments with timestamps shifted to snippet-relative.

Each segment-level and word-level start/end is shifted by -snippet_start and clamped to [0, snippet_end - snippet_start]. Real diarization data has small (~10 ms) jitter where words are annotated as starting fractionally before their parent segment or ending fractionally after, so unclamped values can slip outside [0, duration] even though the snippet boundaries themselves align with segment boundaries; clamping keeps downstream consumers from having to handle that.

Other fields are reused by reference — treat the returned segments as read-only.

nemo_curator.stages.audio.alm.pretrain.planning._MIN_SEGMENTS_FOR_OVERLAP = 2