> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/curator/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/curator/_mcp/server.

# 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`](#nemo_curator-stages-audio-alm-pretrain-planning-OverlapFilterStage)                     | Drop empty segments and overlapping segment pairs.                   |
| [`SnippetCutPlannerStage`](#nemo_curator-stages-audio-alm-pretrain-planning-SnippetCutPlannerStage)             | Compute snippet cut boundaries for one input audio.                  |
| [`SnippetRepetitionFilterStage`](#nemo_curator-stages-audio-alm-pretrain-planning-SnippetRepetitionFilterStage) | Drop planned snippets whose text shows suspicious n-gram repetition. |

### Functions

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

### Data

[`_MIN_SEGMENTS_FOR_OVERLAP`](#nemo_curator-stages-audio-alm-pretrain-planning-_MIN_SEGMENTS_FOR_OVERLAP)

### API

```python
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\]](/nemo-curator/nemo_curator/stages/base#nemo_curator-stages-base-ProcessingStage)

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`

---

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

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

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

```python
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\]](/nemo-curator/nemo_curator/stages/base#nemo_curator-stages-base-ProcessingStage)

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`

---

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

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

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

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

```python
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\]](/nemo-curator/nemo_curator/stages/base#nemo_curator-stages-base-ProcessingStage)

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`

---

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

```python
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.

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

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

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

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

```python
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
```

```python
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.

```python
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`.

```python
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`).

```python
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.

```python
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.

```python
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.

```python
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.

```python
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:

1. its span `[first.start, last.end]` stays within
   `max_duration_sec`, AND
2. 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.

```python
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.

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