diff --git a/.github/workflows/config/.secrets.baseline b/.github/workflows/config/.secrets.baseline index 7529f85089..b146dc231b 100644 --- a/.github/workflows/config/.secrets.baseline +++ b/.github/workflows/config/.secrets.baseline @@ -475,14 +475,14 @@ "filename": "tutorials/audio/fleurs/fleurs_tutorial.ipynb", "hashed_secret": "b158f69d26847139deedc024c3270d8b8fc79d81", "is_verified": false, - "line_number": 232 + "line_number": 347 }, { "type": "Base64 High Entropy String", "filename": "tutorials/audio/fleurs/fleurs_tutorial.ipynb", "hashed_secret": "45b1cbaefaa19a0f8d57a797a742ff8e5681b922", "is_verified": false, - "line_number": 323 + "line_number": 512 } ], "tutorials/audio/readspeech/readspeech_tutorial.ipynb": [ @@ -543,5 +543,5 @@ } ] }, - "generated_at": "2026-09-11T14:18:56Z" + "generated_at": "2026-09-21T18:09:26Z" } diff --git a/benchmarking/scripts/audio_librispeech_benchmark.py b/benchmarking/scripts/audio_librispeech_benchmark.py index 92873e2e94..e4f7673623 100644 --- a/benchmarking/scripts/audio_librispeech_benchmark.py +++ b/benchmarking/scripts/audio_librispeech_benchmark.py @@ -177,6 +177,9 @@ def run_audio_librispeech_benchmark( # noqa: PLR0913 model_id=model_name, audio_filepath_key="audio_filepath", batch_size=asr_batch_size, + max_audio_sec_per_actor=240.0, + max_inference_duration_s=120.0, + local_bucketing=True, fail_on_audio_error=True, adapter_kwargs={"use_cuda_graph_decoder": False}, ) diff --git a/nemo_curator/models/asr/base.py b/nemo_curator/models/asr/base.py index eef44f292c..0b0f08eea4 100644 --- a/nemo_curator/models/asr/base.py +++ b/nemo_curator/models/asr/base.py @@ -70,6 +70,7 @@ class ASRAdapter(Protocol): * ``waveform``: contiguous, mono, 1-D float32 NumPy samples normalized by ``ASRStage`` from a file or a reader-provided in-memory waveform. * ``sample_rate`` (``int``): the stage's configured target sample rate. + * ``audio_seconds`` (``float``): duration of this model-input chunk. * ``language`` (``str | None``): human-readable name (e.g. ``"English"``). * ``language_code`` (``str | None``): original language code from the configured stage input column. diff --git a/nemo_curator/stages/audio/README.md b/nemo_curator/stages/audio/README.md index 3f986a5362..4717a87215 100644 --- a/nemo_curator/stages/audio/README.md +++ b/nemo_curator/stages/audio/README.md @@ -135,9 +135,10 @@ Key differences from a CPU stage: ### Setting `batch_size` for GPU inference The `batch_size` field on a GPU stage controls how many `AudioTask` tasks -the backend groups into a single `process_batch()` call. This directly -determines how many files are passed to your model in one batched GPU -inference call. +the backend groups into a single `process_batch()` call. The stage can then +make one or more model calls from that finite candidate window. For example, +`ASRStage` can segment long parents, locally regroup model items by duration, +and enforce a padded-audio budget for each adapter call. **Defining batch_size in the stage class:** @@ -158,25 +159,31 @@ pipeline.add_stage( ASRStage( adapter_target="nemo_curator.models.asr.nemo_asr.NeMoASRAdapter", model_id="nvidia/parakeet-tdt-0.6b-v2", + max_audio_sec_per_actor=240, + max_inference_duration_s=120, + local_bucketing=True, audio_filepath_key="audio_filepath", ) .with_(resources=Resources(gpus=1), batch_size=32) ) ``` -The `.with_()` method sets any stage field. Here it bumps `batch_size` -from the default `16` to `32` and assigns 1 GPU. +The `.with_()` method supports common execution overrides such as `resources` +and `batch_size`. Here it sets `batch_size` to `32` and assigns 1 GPU. Put +stage- or model-specific fields in the stage constructor. **Overriding batch_size via Hydra YAML:** ```yaml -pipeline: - stages: - - _target_: nemo_curator.stages.audio.inference.asr.stage.ASRStage - adapter_target: nemo_curator.models.asr.nemo_asr.NeMoASRAdapter - model_id: nvidia/parakeet-tdt-0.6b-v2 - audio_filepath_key: audio_filepath - batch_size: 32 +stages: + - _target_: nemo_curator.stages.audio.inference.asr.stage.ASRStage + adapter_target: nemo_curator.models.asr.nemo_asr.NeMoASRAdapter + model_id: nvidia/parakeet-tdt-0.6b-v2 + max_audio_sec_per_actor: 240 + max_inference_duration_s: 120 + local_bucketing: true + audio_filepath_key: audio_filepath + batch_size: 32 ``` For Hydra to accept `batch_size` from YAML, it must be a dataclass field @@ -189,19 +196,25 @@ Backend reads stage.batch_size → groups N tasks into batches of batch_size → sends each batch to a worker → worker calls stage.process_batch(tasks) - → your override receives exactly batch_size tasks - (or fewer for the last batch) + → your override receives that finite candidate window + → stage-specific code makes one or more model calls ``` +For variable-duration audio, `batch_size` is not necessarily the exact number +of items in one model call. See +[Local Duration Bucketing for Audio GPU Inference](inference/README.md) for the +current `ASRStage` behavior, the duration-packing theory, and the integration +contract for other GPU stages. + **Choosing a good batch_size:** -- **Too small** (e.g. `1`) — GPU is underutilised; kernel launch overhead - dominates. Each call processes one file, losing the benefit of batching. -- **Too large** (e.g. `1024`) — may exceed GPU memory (OOM), especially - with long audio files or large models. -- **Sweet spot** — depends on model size, audio length, and GPU VRAM. - Start with `16` and increase until you see OOM or throughput plateaus. - For NeMo ASR FastConformer models, `16–64` is typical on a single GPU. +- **Too small** (e.g. `1`) — gives the stage little opportunity to form useful + model batches or duration-coherent groups. +- **Too large** (e.g. `1024`) — can increase waveform preparation and host + memory pressure before the stage makes any model calls. +- **Sweet spot** — depends on the model, audio distribution, GPU memory, and + the stage-level audio budget. Tune with representative inputs rather than + treating the backend window as the model batch size. ## What you must always declare @@ -252,8 +265,9 @@ process_batch(list[AudioTask]) -> list[AudioTask] with N tasks. - `process` is the natural single-task hook for CPU stages — no boilerplate to handle lists. -- GPU/IO stages override `process_batch` to receive the full batch for - one batched kernel call. Their `process()` raises +- GPU/IO stages override `process_batch` to receive the full backend batch and + organize its work efficiently. A stage may issue one or more bounded model + calls from that candidate window. Their `process()` raises `NotImplementedError`, matching the dedup-stage convention (`ConnectedComponentsStage`, `KMeansReadFitWriteStage`, etc.). @@ -346,10 +360,12 @@ For a GPU stage with `resources=Resources(cpus=1.0, gpus=1.0)` and └─────────────────────────────────────────┘ ``` -Each `process_batch([16 tasks])` call goes directly to: -`ASRStage.process_batch` → validate and load the current task waveforms → -`NeMoASRAdapter.transcribe_batch` → **one** batched NeMo call → mutate each -task in-place. +Each `process_batch([16 tasks])` call goes through: +`ASRStage.process_batch` → validate, load, and model-safely segment the current +waveforms → plan capacity-bounded adapter calls across all segments in this +window → `NeMoASRAdapter.transcribe_batch` once per planned call → stitch +segments and restore parent-task order. `batch_size=16` therefore defines the +planning window, not a guarantee of exactly one 16-item NeMo call. ### Xenna specifics @@ -398,10 +414,11 @@ pipeline.run(executor) | Level | What it controls | Who sets it | |---|---|---| | **Worker count** | How many parallel copies of your stage run (one per CPU core or GPU) | The backend, based on `stage.resources` and available hardware | -| **`batch_size`** | How many tasks each worker processes per call | The stage author (domain knowledge about optimal GPU batch size) | +| **`batch_size`** | Maximum candidate tasks supplied to each worker's `process_batch` call | The stage author | -Total in-flight = `num_workers x batch_size`. For 4 GPUs with -`batch_size=16`, that is 64 audio files being processed concurrently. +Maximum candidate tasks in flight = `num_workers x batch_size`. For 4 GPUs +with `batch_size=16`, up to 64 audio files can be inside stage calls at once. +The stage may split each candidate window into smaller model calls. ## How `batch_size` travels from your stage to the backend @@ -440,7 +457,9 @@ Key takeaways: - Subclasses override it as a field (e.g. `batch_size: int = 16`). - Pipeline authors can further override via `.with_(batch_size=32)` or Hydra YAML. - The backend adapter reads `stage.batch_size` and groups tasks *before* - calling `process_batch`. Your stage never has to split or batch tasks itself. + calling `process_batch`. +- A stage can still split model-unsafe inputs and plan one or more adapter + calls within that finite backend-provided batch, as `ASRStage` does. ## Exact call chains @@ -527,7 +546,7 @@ pipeline.run(executor) │ → models/asr/nemo_asr.py NeMoASRAdapter.load_model(num_gpus=1) │ ASRModel.from_pretrained(model_name=model_id, map_location=cuda) │ -├─ Per batch (batch_size=16, so 16 AudioTask tasks per call): +├─ Per backend batch (batch_size=16, so up to 16 candidate tasks per call): │ backends/xenna/adapter.py XennaStageAdapter.process_data(tasks) │ → backends/base.py BaseStageAdapter.process_batch(tasks) │ ├─ start perf timer @@ -539,14 +558,14 @@ pipeline.run(executor) │ │ ASRStage.process_batch() (generic batched GPU stage) │ │ stages/audio/inference/asr/stage.py │ │ validate_input(task) per task schema check -│ │ load and normalize the current 16 waveforms -│ │ adapter.transcribe_batch(items) -│ │ → models/asr/nemo_asr.py -│ │ self._model.transcribe(audio=waveforms, batch_size=16) -│ │ → ONE batched NeMo inference call -│ │ return list[ASRResult] -│ │ for task, result in zip(tasks, results): -│ │ task.data[self.pred_text_key] = result.text +│ │ load and normalize up to 16 current waveforms +│ │ split every waveform at max_inference_duration_s +│ │ plan calls bounded by max_audio_sec_per_actor +│ │ for each planned call: +│ │ adapter.transcribe_batch(items) +│ │ → models/asr/nemo_asr.py +│ │ self._model.transcribe(audio=waveforms, batch_size=len(items)) +│ │ restore segment order, stitch parent transcripts, and update tasks │ └─ return tasks → same 16 AudioTask objects ``` @@ -656,10 +675,11 @@ AudioTask( ### Stage 2: `ASRStage` + `NeMoASRAdapter` (GPU) -The generic stage loads and normalizes each current-batch waveform; the NeMo -adapter loads `nvidia/parakeet-tdt-0.6b-v2` onto the GPU and runs one batched -`transcribe()` call for 16 `AudioTask`s. The stage writes predictions back -**in-place**. +The generic stage loads and normalizes each current-batch waveform, performs +model-safe segmentation, and plans capacity-bounded calls from the current +candidate window. The NeMo adapter loads `nvidia/parakeet-tdt-0.6b-v2` onto the +GPU and runs one batched `transcribe()` per planned call. The stage stitches +segments, restores parent order, and writes predictions back **in-place**. **Output** — `data` gains `pred_text`: @@ -1016,3 +1036,501 @@ Appends the entry as a single JSON line to `NotImplementedError` (matching the dedup-stage convention). 7. Declare GPU resources via `.with_(resources=Resources(gpus=1.0))` 8. Add tests in `tests/stages/audio/` using `AudioTask` for fixtures + +--- + +## Local duration bucketing: complete theory and current contract + +This section describes the local duration-bucketing algorithm currently +implemented by [`ASRStage`](inference/asr/stage.py), why it is structured this +way, and the contract another audio GPU inference stage must preserve to adopt +the same approach. This is the top-level conceptual contract; the dedicated +[audio inference guide](inference/README.md) adds ASR-specific configuration, +implementation, and test references. + +Local bucketing solves one specific problem: candidate audio inputs in one GPU +batch can have very different durations, while the model usually pads every +input to the longest input in that batch. The stage therefore uses duration to +reorder only the finite set of model inputs already present in one +`process_batch()` call. It does not buffer globally or change which parent +rows belong to that call. + +### The four controls are independent + +| Control | Current `ASRStage` default | Boundary it controls | +|---|---:|---| +| `batch_size` | `32` | Number of parent `AudioTask` rows the backend normally offers to one `process_batch()` planning window | +| `max_inference_duration_s` | `2400` | User-supplied maximum duration of one segment derived from a parent waveform | +| `max_audio_sec_per_actor` | required | Maximum padded-audio proxy cost of one adapter call | +| `local_bucketing` | `false` | Whether segments are stable-sorted by duration and DP-partitioned over contiguous spans instead of greedily packed in input order | + +These settings are deliberately not aliases for one another: + +- `batch_size` bounds the backend candidate window; it is not an adapter item + limit. `ASRStage.process_batch()` itself does not reject a directly supplied + list merely because it contains more than `batch_size` parents. +- `max_inference_duration_s` enforces the configured single-input ceiling; it + applies with bucketing both on and off. The stage does not discover or + verify the selected model's true limit. +- `max_audio_sec_per_actor` bounds every adapter invocation produced from the + current window; it is not the total audio accepted by `process_batch()`. +- `local_bucketing` changes planning order and boundary selection only; it + does not enable or disable segmentation. + +There is no separate `adapter_batch_size`, maximum-items-per-bucket map, +static set of bucket edges, timer, queue, flush hook, or generic item-cost +feature. Despite its name, `max_audio_sec_per_actor` is a per-adapter-call +budget planned by the actor. It is not accumulated over the actor's lifetime +or across multiple `process_batch()` calls. + +Configuration validation matches those boundaries: + +- `max_audio_sec_per_actor` is required, must be a numeric `Real` other than + `bool`, and must be finite and greater than zero; +- `local_bucketing` must be an actual Boolean, so integer `1` is rejected; +- `max_inference_duration_s` must be finite and greater than zero, and cannot + exceed `max_audio_sec_per_actor`; +- `batch_size` and `target_sample_rate` are converted to integers and must be + greater than zero. + +The planner consumes durations generated internally from prepared segments; +it does not expose an API for arbitrary caller-provided cost features. + +### Parent rows, segments, and adapter calls + +Three different units pass through the stage: + +1. A **parent row** is one `AudioTask` received in `process_batch(tasks)`. +2. A **segment** is one model-safe waveform prepared from a parent. A short + parent normally produces one segment; a long parent produces several. +3. An **adapter call** is one capacity-bounded group of segments passed to + `adapter.transcribe_batch()`. Its segments may come from multiple parents. + +The complete data flow is: + +```text +backend candidate rows (`batch_size`) + -> skip/eligibility checks + -> load, downmix, and resample each eligible waveform + -> model-safe segmentation of every eligible parent + -> one flat segment list for the current `process_batch()` call + -> optional stable duration sort + -> padded-seconds partitioning + -> one adapter invocation per planned group + -> scatter results to original segment positions + -> stitch each parent's segments in temporal order + -> return parent rows in their original order +``` + +Rows with a reused output, rows rejected by a configured language allowlist +(including a missing code while that allowlist is active), and rows whose +waveforms fail preparation without raising do not enter the segment plan. +Every other segment from every eligible row in the same +`process_batch()` participates in one shared local planning window. No segment +is retained for a later backend batch, actor, or worker. + +### Why the budget measures padded audio + +For a nonempty adapter call with segment durations +`d[0], ..., d[n - 1]`, define: + +```text +useful audio seconds = sum(d) +padded audio seconds = n * max(d) +padding waste = n * max(d) - sum(d) +padding efficiency = sum(d) / (n * max(d)), when max(d) > 0 +``` + +For an all-zero-duration call, padded cost and padding waste are both zero, +while this efficiency ratio is undefined rather than `100%`. + +The feasibility condition for every planned call is: + +```text +n * max(d) <= max_audio_sec_per_actor +``` + +This is intentionally not `sum(d)`. If durations `[1, 10]` are sent together, +the shorter waveform is normally padded to 10 seconds and the model processes +a tensor representing approximately `2 * 10 = 20` audio seconds, not 11. +Grouping similar durations reduces work that carries no useful audio. + +Padded seconds are a useful proxy, not a proof of GPU-memory safety. Actual +memory and runtime also depend on model architecture, precision, decoder +state, framework workspaces, fixed per-item costs, and adapter-specific +behavior. The budget must therefore be measured and tuned for each model, +adapter, and hardware configuration. The proxy is most meaningful when the +adapter forms a jointly padded batch. A serial, ragged, or packed adapter may +receive the same planned groups without realizing the expected padding +benefit. + +### Model-safe segmentation always precedes bucketing + +After waveform preparation, `ASRStage` calls +[`plan_audio_segments`](model_input_segmentation.py) for every eligible parent. +For target sample rate `s` and duration ceiling `D`, the segment limit in +samples is: + +```text +max_samples = int(D * s) +``` + +`D` must cover at least one sample. The current planner creates contiguous, +nonoverlapping ranges from sample zero to the end of the waveform. Every +remainder becomes a segment, an exact multiple creates no empty tail, and a +zero-sample waveform remains representable as one zero-duration segment. +`audio_seconds` is calculated from the actual prepared segment length and +target sample rate, rather than trusted from manifest metadata. + +Construction requires: + +```text +max_inference_duration_s <= max_audio_sec_per_actor +``` + +This guarantees that every permitted individual segment can at least fit in a +singleton adapter call. + +The current ASR segmentation has an important semantic limitation: the cuts +are hard and have no shared audio context. Each segment is decoded +independently, and the current stitcher joins nonempty transcript strings with +a space. A word crossing a cut can consequently be omitted, duplicated, or +decoded differently. Local bucketing neither causes nor fixes that behavior; +adding overlap would also require a model-appropriate transcript +reconciliation algorithm. Other audio stages must define their own safe +segmentation and reconstruction semantics rather than copying ASR text +concatenation blindly. + +### Planning when `local_bucketing=true` + +Each prepared segment enters the planner as: + +```text +(original_segment_index, adapter_item, audio_seconds) +``` + +The enabled planner then performs these steps: + +1. Stable-sort all segments by ascending `audio_seconds`. Equal-duration + segments retain their original relative order. +2. Restrict each adapter call to a contiguous span of that sorted sequence. +3. Consider every span whose padded cost fits the actor budget. +4. Use dynamic programming to select the complete partition with the fewest + adapter calls. +5. Among partitions with that minimum call count, select the one with the + fewest total padded seconds. +6. Execute calls in planned order and scatter every result through the saved + original segment index. + +For sorted durations `d[0] <= ... <= d[N - 1]`, the cost of span +`[start, stop)` is: + +```text +span_count = stop - start +span_max = d[stop - 1] +span_cost = span_count * span_max +``` + +A span is feasible when `span_cost` is within the budget. The implementation +accepts representation-level equality using `math.isclose` with +`rel_tol=1e-12` and `abs_tol=1e-9`, avoiding a spurious split for values such +as three 0.1-second segments under a 0.3-second budget. + +The dynamic-programming suffix state can be written as: + +```text +best[N] = (0 calls, 0 padded seconds) + +best[start] = lexicographic minimum over every feasible stop: + ( + 1 + best[stop].calls, + cost(start, stop) + best[stop].padded_seconds, + ) +``` + +The selected `stop` is stored for each `start`, then those boundaries are +followed from zero to reconstruct the calls. Python tuple comparison gives the +required calls-first, padded-seconds-second ordering. Exact score ties retain +the first feasible boundary encountered. Score comparison uses the raw float +totals; the small tolerance described above applies only to budget +feasibility. + +This is an exact optimum for the implementation's stable duration-sorted, +contiguous-span planning space. It is not an optimizer over arbitrary +noncontiguous subsets, other `process_batch()` calls, actors, workers, or the +whole dataset. + +### Why sorted greedy filling is not enough + +Greedily extending a sorted call until the next segment would breach the +budget makes the current call as full as possible, but can make the complete +plan do unnecessary padded work. + +For sorted durations `[1, 2, 2]` and budget `4`, greedy produces: + +```text +[1, 2] -> 2 * 2 = 4 padded seconds +[2] -> 1 * 2 = 2 padded seconds +score -> 2 calls, 6 padded seconds +``` + +The current dynamic program chooses: + +```text +[1] -> 1 * 1 = 1 padded second +[2, 2] -> 2 * 2 = 4 padded seconds +score -> 2 calls, 5 padded seconds +``` + +It deliberately leaves room unused in the first call because doing so keeps +the same two-call count and eliminates one padded second overall. Three +singleton calls would also use five padded seconds, but lose on the primary +call-count objective. Prioritizing call count prevents the padding objective +from degenerating into one adapter call per segment and minimizes adapter +launches before optimizing the padded-work proxy; it is an objective, not a +guarantee of the lowest wall-clock time. + +### Example across several parent rows + +Suppose one `process_batch()` receives three parents with prepared durations +`250`, `10`, and `1` seconds, with: + +```text +max_inference_duration_s = 120 +max_audio_sec_per_actor = 240 +local_bucketing = true +``` + +Segmentation happens first: + +```text +parent A -> [120, 120, 10] +parent B -> [10] +parent C -> [1] +``` + +The shared sorted planning sequence is `[1, 10, 10, 120, 120]`, and the +selected adapter calls are: + +```text +[1, 10, 10] -> 3 * 10 = 30 padded seconds +[120, 120] -> 2 * 120 = 240 padded seconds +``` + +The 10-second remainder from parent A can share a call with segments from B +and C. After inference, saved indices restore temporal and parent ownership: +parent A receives its `120`, `120`, and `10` results in that order, while B +and C each receive their own result. Adapter-call order never becomes output +row order. + +### Planning when `local_bucketing=false` + +The disabled mode is the input-order control. It keeps the flattened segment +order and greedily extends the current call while its candidate padded cost +fits. When the next segment would breach the budget, it emits the current call +and starts another with that segment. + +Disabling local bucketing therefore disables the duration sort and the +whole-window dynamic-programming boundary optimization, but it does not +disable either segmentation or the +`max_audio_sec_per_actor` constraint. The same adapter and result-restoration +contract applies in both modes. + +For example, original durations `[8, 2, 7, 3]` under budget `16` produce: + +```text +input-order greedy: + [8, 2] -> 16 padded seconds + [7, 3] -> 14 padded seconds + total -> 30 padded seconds + +local bucketing, sorted order [2, 3, 7, 8]: + [2, 3] -> 6 padded seconds + [7, 8] -> 16 padded seconds + total -> 22 padded seconds +``` + +Both plans process the same 20 useful seconds in two calls, but local +bucketing presents substantially less padding to the model. + +### Adapter execution, result restoration, and ASR stitching + +Every selected group causes exactly one `adapter.transcribe_batch(items)` +invocation. `NeMoASRAdapter` filters zero-length waveforms while preserving +their result slots. If any nonempty waveforms remain, it then makes one +`model.transcribe(audio=waveforms, batch_size=len(waveforms), ...)` call for +that planned group. Consequently, one `process_batch()` can make zero, one, or +many model calls. + +The adapter must return exactly one result for every submitted item. The stage +places each result into an array indexed by the segment's position before +bucketing and rejects incomplete or wrong-sized result sets. It then groups +those aligned segment results by parent, preserving the parent's temporal +segment order. This restoration relies on the adapter's ordered contract of +one result per input; matching result counts alone cannot detect an adapter +that internally permutes its outputs. + +For a multi-segment ASR parent, the current stitcher: + +- strips and space-joins nonempty segment transcripts; +- marks the parent skipped if any segment was skipped; +- retains the first available skip reason and unsupported-language value; +- merges adapter `extras` dictionaries in segment order, with a later value + replacing an earlier value for the same key. + +Finally, predictions are written to the same parent `AudioTask` objects and +the original parent-row order is returned. Local execution order is therefore +an internal optimization, not an externally visible row reorder. + +### Correctness invariants + +The implementation is correct only if all of the following remain true: + +1. Model-safe segmentation runs whether bucketing is enabled or disabled. +2. Duration is calculated once from each final prepared segment. +3. Every eligible segment enters exactly one adapter call. +4. Every nonempty call satisfies + `len(call) * max(segment_duration) <= max_audio_sec_per_actor`, allowing + only the documented representation-level floating-point tolerance. +5. Disabled mode greedily preserves original segment order. +6. Enabled mode uses stable ascending-duration order and chooses the + contiguous partition with minimum call count and then minimum total padded + seconds. +7. Every adapter call returns one result per submitted segment. +8. Results are scattered to original segment positions before parent + assembly. +9. Each parent's segment results are reconstructed in temporal order. +10. For successfully returned, ordered adapter results, scatter preserves + parent association and parent-row order. The planner cannot guarantee that + a stateful adapter's outputs or exceptions are independent of call + composition or execution order. +11. No queued audio, timer, flush obligation, or planner state survives the + current `process_batch()` call. + +The adapter and loaded model remain worker-local and may persist across calls; +that lifecycle is separate from the deliberately stateless local planner. + +### Applying the same design to another audio GPU stage + +The planner currently lives in `ASRStage`; it is not a generic base-class +hook. It is appropriate for another stage only when independent audio inputs +can be grouped safely, each result can be mapped back unambiguously, and +`item_count * longest_duration` meaningfully approximates that adapter's +jointly padded work. Serial, ragged, or packed execution may need a different +cost model or may gain nothing from duration bucketing. + +To apply the pattern: + +1. Keep the backend `batch_size`, required padded-seconds budget, and + `local_bucketing` Boolean as distinct controls. +2. Validate and prepare all eligible parents within one finite + `process_batch()` call. +3. Apply any model-specific maximum-input segmentation unconditionally. +4. Flatten the resulting model items while saving original item index, parent + index, segment ordinal, timestamps, output paths, and any other metadata + required to reconstruct results. +5. Measure duration from the final resampled or sliced waveform. +6. Preserve input order and greedily pack when bucketing is disabled. +7. Stable-sort by duration and use the calls-first, padded-seconds-second + dynamic program over contiguous spans when bucketing is enabled. +8. Call the adapter once per planned group, require a complete ordered + one-to-one result mapping, scatter by saved index, and only then assemble + parent outputs. How many native model calls the adapter makes is + adapter-specific. +9. Test both modes against the same correctness oracle before measuring + performance. + +Reordering independent whole inputs is often safe; segmentation is +model-specific: + +- ASR must define how boundary transcripts are reconciled. +- SED must restore valid-frame counts, frame offsets, timestamps, and sidecar + paths. +- VAD splits can change onset/offset decisions and merged speech regions. +- Diarization may require whole-recording speaker clustering and identity. +- Alignment must keep every hypothesis paired with the right segment + metadata. +- File-producing stages must derive paths from saved parent metadata rather + than reordered call positions. + +Do not add buffering across `process_batch()` calls merely to improve the +duration distribution. That would change latency, ownership, failure, and +end-of-stream semantics and would require a separate queue-and-flush design. + +### Tuning and measurement + +There is no universal safe or optimal audio-seconds budget: + +1. Choose `max_inference_duration_s` from the model's semantic and technical + single-input limit first. +2. Start `max_audio_sec_per_actor` from a known-safe uniform model call. If + `k` inputs of duration `d` are safe, `k * d` is a reasonable initial proxy + budget. Keep it at least as large as `max_inference_duration_s`. +3. Exercise the longest allowed singleton, then increase the budget gradually + while observing GPU memory, throughput, latency, and failures. +4. Benchmark representative short, medium, long, exact-boundary, and remainder + inputs. Average duration hides padding and tail effects. +5. Compare enabled and disabled modes using identical candidate windows, + model settings, software, input order, and hardware. Establish output + parity before comparing performance. +6. Tune backend `batch_size` separately. A small window offers few regrouping + choices; a very large window increases waveform-preparation latency and + host memory before inference begins. + +A duration budget cannot represent fixed per-item memory. A finite candidate +window containing many very short or zero-duration items can therefore +produce a large item-count call. If a model has a hard native item-count +limit, enforce or expose that adapter-specific constraint at the model +boundary rather than silently redefining `batch_size`. + +### Complexity, guarantees, and non-goals + +For `N` prepared segments: + +- bucketing disabled: `O(N)` greedy planning; +- bucketing enabled: `O(N log N)` stable sorting plus `O(N^2)` dynamic + programming, for `O(N^2)` total time and `O(N)` auxiliary storage. + +Given the same prepared segment sequence and configuration, planning is +deterministic. Enabled mode guarantees the minimum number of calls and then +minimum padded seconds within its sorted contiguous-span search space. It +does not guarantee lower end-to-end latency or higher throughput for every +model and duration distribution. The model-time or resource savings from +reduced padding must exceed the added sorting and dynamic-programming cost for +the optimization to pay off. + +The current design intentionally does not provide: + +- global, partition-wide, cross-worker, or cross-`process_batch()` bucketing; +- queues, timers, flush hooks, or end-of-stream state; +- static duration buckets or per-bucket configuration; +- a separate adapter batch size or per-bucket item-count map; +- feature-weighted cost estimation; +- arbitrary noncontiguous grouping after the stable duration sort; +- adaptive OOM retries or automatic budget tuning; +- a GPU-memory-safety proof from the duration proxy; +- a throughput win for every workload; +- boundary-safe ASR overlap and transcript reconciliation. + +### Tests a reusable implementation needs + +At minimum, tests should cover: + +- missing, Boolean, nonnumeric, zero, negative, `NaN`, and infinite budgets; +- a model limit larger than the actor budget and a limit shorter than one + sample; +- empty input, zero-duration audio, exact fills, and floating-point boundary + equality; +- original-order packing, stable duration ordering, and equal-duration + stability; +- the `[1, 2, 2]` under budget `4` case that distinguishes the dynamic program + from sorted greedy filling; +- segments from multiple parent rows sharing one call; +- strict isolation between separate `process_batch()` calls; +- exact model-duration boundaries, long parents, and every final remainder; +- exact-once submission, wrong adapter result counts, scatter, stitching, + skip, and error behavior; +- proof that backend `batch_size` is not treated as an adapter item cap; +- parity of user-visible outputs, using exact equality or the stage's + documented tolerance, between enabled and disabled modes before comparing + padding efficiency, throughput, latency, RAM, or VRAM. diff --git a/nemo_curator/stages/audio/inference/README.md b/nemo_curator/stages/audio/inference/README.md new file mode 100644 index 0000000000..a32aba9c4c --- /dev/null +++ b/nemo_curator/stages/audio/inference/README.md @@ -0,0 +1,428 @@ +# Local Duration Bucketing for Audio GPU Inference + +This document describes the local duration-batching contract implemented by +[`ASRStage`](asr/stage.py) and how to apply the same pattern to another audio +GPU inference stage. + +The contract has three controls: + +| Field | Default | Purpose | +|---|---:|---| +| `max_audio_sec_per_actor` | required | Maximum padded audio seconds planned for one adapter call | +| `max_inference_duration_s` | `2400` | Model-specific maximum duration of one segment | +| `local_bucketing` | `false` | Stable-sort by duration and optimize call boundaries inside the current planning window | + +There are no bucket edges, per-bucket limits, item-count caps, timers, queues, +flush operations, generic cost features, or separate adapter batch-size +setting. `batch_size` remains a backend candidate-window setting; it does not +cap the number of segments in an adapter call. + +## Mental model + +Audio tensors in one GPU call are commonly padded to the longest item. For a +nonempty call with durations `d[0] ... d[n-1]`, define: + +```text +useful audio seconds = sum(d) +padded audio seconds = n * max(d) +padding efficiency = sum(d) / (n * max(d)), when max(d) > 0 +``` + +For an all-zero-duration call, padded cost is zero and padding efficiency is +undefined. + +`ASRStage` uses `padded audio seconds` as a simple capacity proxy. Every +planned adapter call must satisfy: + +```text +number of items in the call * longest item duration + <= max_audio_sec_per_actor +``` + +For an adapter that forms a jointly padded batch, this proxy captures the main +cost of padding more directly than summing raw durations. A serial, ragged, or +packed adapter may not realize the expected benefit. The proxy is also not a +proof of GPU-memory safety: model architecture, precision, decoder state, +framework workspaces, and fixed per-item overhead matter. The value must be +measured and tuned for each adapter, model, and hardware configuration. + +Despite its name, `max_audio_sec_per_actor` limits each adapter invocation +planned by an actor. It is not a lifetime quota and does not accumulate across +`process_batch()` calls. + +## Scope and data flow + +The complete planning horizon is exactly one finite `process_batch(tasks)` +call: + +```text +backend candidate rows (`batch_size`) + -> eligibility checks and waveform preparation + -> model-safe segmentation of every eligible parent + -> one flat list of segments from all eligible parents + -> optional stable duration ordering + -> padded-seconds packing + -> adapter calls + -> scatter segment results to original positions + -> stitch segments for each parent + -> return parent rows in original order +``` + +The distinctions in that flow are important: + +1. A **parent row** is one `AudioTask` received from the backend. +2. A **segment** is one model-safe waveform derived from a parent. One parent + may create several segments. +3. An **adapter call** is one capacity-bounded group of segments. It may + contain segments from several parents, including a short remainder from a + long parent. + +Skipped languages, reused outputs, and waveforms that fail preparation do not +enter the segment plan. All other segments produced from all input rows in the +same `process_batch()` call share one local planning window. Nothing is held +for a later call, another actor, or another worker. + +## Segmentation always happens first + +`local_bucketing` controls reordering only. It never controls segmentation. + +After decode, downmix, and resampling, `ASRStage` always calls +[`plan_audio_segments`](../model_input_segmentation.py) with +`max_inference_duration_s`. For sample rate `s` and model limit `D`, the +maximum samples in one segment are: + +```text +max_samples = int(D * s) +``` + +The intervals are contiguous, nonoverlapping, and cover the full waveform. +An exact multiple creates no empty tail, and every nonempty remainder becomes +another segment. A zero-sample waveform remains representable as one +zero-duration segment. If `max_samples` is less than one, segmentation rejects +the configured limit at processing time rather than clamping it. + +The stage derives `audio_seconds` from the actual prepared segment, not from a +possibly stale manifest duration: + +```python +audio_seconds = number_of_segment_samples / target_sample_rate +``` + +`max_inference_duration_s` must be less than or equal to +`max_audio_sec_per_actor`. That construction-time validation guarantees that +one model-safe segment can fit in one planned call. Choose the model limit for +model correctness first, then choose an actor budget large enough to contain +it. + +## Packing algorithm + +Planning is deterministic, but the two modes use different boundary planners. +Both enforce the padded-audio budget on every adapter call. Durations need no +second validation here: they are produced internally from model-safe segments. + +With `local_bucketing=false`, the planner preserves input order and scans it +greedily. It extends the current call while the candidate padded cost fits, +emits that call when the next segment would exceed the budget, and then starts +the next call with that segment. + +With `local_bucketing=true`, the planner: + +1. Stable-sorts all segments by ascending duration. Equal-duration segments + retain their original relative order. +2. Treats every adapter call as one contiguous span of that sorted sequence. + For a span `[start, stop)`, its padded cost is + `(stop - start) * duration[stop - 1]`. +3. Uses dynamic programming from the end of the sequence to score every + feasible next boundary. +4. Minimizes the score lexicographically: first the total number of adapter + calls, then the sum of padded seconds across those calls. +5. Reconstructs the selected spans, executes them in sorted-plan order, and + scatters each result through its saved pre-planning segment index. + +The enabled planner is therefore the exact optimum for its stable +duration-sorted, contiguous-span planning space. It is not a static bucket +scheme or a scheduler across process calls, actors, or workers. With bucketing +off, the same budget still applies; only reordering and dynamic-programming +boundary optimization are disabled. + +### Worked example + +Suppose the segment durations in original order are: + +```text +indices: [0, 1, 2, 3] +durations: [8, 2, 7, 3] +budget: 16 padded audio seconds +``` + +Without local bucketing, original-order packing produces: + +```text +[8, 2] -> 2 * 8 = 16 padded seconds +[7, 3] -> 2 * 7 = 14 padded seconds +``` + +With local bucketing, stable ascending order is indices `[1, 3, 2, 0]`: + +```text +[2, 3] -> 2 * 3 = 6 padded seconds +[7, 8] -> 2 * 8 = 16 padded seconds +``` + +The same 20 useful audio seconds require 30 proxy seconds without reordering +and 22 with reordering. The adapter results are then scattered to indices +`[0, 1, 2, 3]`, so adapter-call order cannot reorder output rows. + +### Why the enabled planner uses dynamic programming + +Sorted greedy packing does not always minimize padding when several plans use +the same number of calls. For sorted durations `[1, 2, 2]` and a budget of `4`, +a greedy scan would produce: + +```text +[1, 2] -> 2 * 2 = 4 padded seconds +[2] -> 1 * 2 = 2 padded seconds +score -> 2 calls, 6 padded seconds +``` + +The dynamic-programming planner instead selects: + +```text +[1] -> 1 * 1 = 1 padded second +[2, 2] -> 2 * 2 = 4 padded seconds +score -> 2 calls, 5 padded seconds +``` + +Three singleton calls would also total 5 padded seconds, but they lose on the +primary call-count objective. Thus the selected `[1]`, `[2, 2]` plan is the +lexicographic optimum. Python tuple ordering implements the calls-first, +padded-seconds-second comparison directly. Exact score ties keep the first +boundary encountered. Budget feasibility alone uses a small floating-point +tolerance (`rel_tol=1e-12`, `abs_tol=1e-9`), so decimal boundaries such as +three 0.1-second items under a 0.3-second budget do not split spuriously. + +## Configuration + +### Python + +```python +from nemo_curator.stages.audio.inference.asr.stage import ASRStage + +asr = ASRStage( + adapter_target="nemo_curator.models.asr.nemo_asr.NeMoASRAdapter", + model_id="nvidia/stt_en_fastconformer_ctc_large", + audio_filepath_key="resampled_audio_filepath", + max_audio_sec_per_actor=240, + max_inference_duration_s=120, + local_bucketing=True, + batch_size=32, +) +``` + +### Hydra YAML + +```yaml +- _target_: nemo_curator.stages.audio.inference.asr.stage.ASRStage + adapter_target: nemo_curator.models.asr.nemo_asr.NeMoASRAdapter + model_id: nvidia/stt_en_fastconformer_ctc_large + audio_filepath_key: resampled_audio_filepath + max_audio_sec_per_actor: 240 + max_inference_duration_s: 120 + local_bucketing: true + batch_size: 32 +``` + +Set `local_bucketing: false` for the original-order control. Do not remove +`max_audio_sec_per_actor`: it is required and bounds calls in both modes. + +The example budget permits either two 120-second segments, four 60-second +segments, or twenty-four 10-second segments when those are the longest items +in their respective calls. Actual group membership depends on all durations +in the current candidate window. + +## What each control does + +| Control | Boundary | Effect | +|---|---|---| +| `batch_size` | Backend to stage | Maximum candidate parent rows normally supplied to one `process_batch()` call; a larger window offers more regrouping choices but uses more host memory | +| `max_inference_duration_s` | Parent waveform to segment | Always splits prepared audio at the model-specific single-input limit | +| `local_bucketing` | Segment planning order | `true` uses stable ascending duration order; `false` keeps original order | +| `max_audio_sec_per_actor` | Planned adapter call | Caps the padded-seconds proxy in both ordering modes | +| Adapter-native settings | Inside adapter/model | Preserve any model-specific decoding or implementation controls | + +In particular, `batch_size` is not reused as an adapter item cap. A backend +window of three parent rows may yield one adapter call, several adapter calls, +or more than three segments after long-audio splitting. + +## Tuning + +There is no universal budget. Tune with the exact model, precision, GPU, +decoder settings, and post-segmentation duration distribution. + +1. Establish the largest semantically safe single input and set + `max_inference_duration_s`. Segmentation and stitching behavior is part of + the model contract, not a memory-tuning trick. +2. Start `max_audio_sec_per_actor` from a known-safe uniform call. If `k` + clips of duration `d` are safe, `k * d` is a reasonable initial proxy + budget. Keep it at least as large as `max_inference_duration_s`. +3. Exercise the longest allowed segment by itself. Then increase the budget + gradually while measuring peak GPU memory, throughput, and failures. +4. Use representative short, medium, long, exact-boundary, and remainder + inputs. Average duration alone hides padding and tail behavior. +5. Compare `local_bucketing=false` and `true` with the same candidate windows, + inputs, model settings, and hardware. Validate outputs before comparing + speed. +6. Tune `batch_size` separately. Too small gives the local optimizer few + alternatives; too large increases decode/preparation latency and host + memory because all candidate waveforms are prepared before model execution. + +A duration budget cannot express fixed per-item memory. Very short or +zero-duration items may therefore produce a large call from a finite candidate +window. If an adapter has a hard native count limit, it must enforce or expose +that model-specific constraint at its actual model boundary; do not silently +reinterpret `batch_size` as that limit. + +## Correctness invariants + +An implementation is correct only when all of these properties hold: + +1. Model-safe segmentation runs whether local bucketing is on or off. +2. Every eligible segment has one duration computed after final waveform + preparation. +3. Every eligible segment is submitted exactly once. +4. Every nonempty adapter call satisfies + `len(call) * max(duration) <= max_audio_sec_per_actor`. +5. Bucketing off greedily packs the original segment order. Bucketing on uses + stable ascending duration order and chooses the contiguous partition with + the fewest calls and then the fewest total padded seconds. +6. Each adapter call returns exactly one result per submitted segment. +7. Results are scattered to original segment positions before parent + assembly. +8. Each parent's segment results are stitched in temporal order. +9. For successfully returned ordered adapter results, scatter preserves + parent association and parent output order. A stateful adapter's outputs or + exceptions may still depend on call composition or execution order. +10. No pending audio, planner state, timer, or flush obligation survives the + current `process_batch()` call. + +The persistent worker-local adapter and model are unrelated to planner state; +they may live across calls as usual. + +## Applying the pattern to another audio GPU stage + +The planner is currently implemented inside `ASRStage`; it is not a generic +base-class hook. Another stage can adopt the approach when independent audio +inputs can be grouped safely, results can be mapped back unambiguously, and +`item_count * longest_duration` meaningfully approximates the adapter's +jointly padded work. + +Use this integration sequence: + +1. Add the same required padded-seconds budget and local-bucketing Boolean to + the stage. Keep `batch_size` as its backend window. +2. In one `process_batch()`, validate and prepare all eligible parents before + planning adapter calls. +3. Apply any model-specific segmentation unconditionally. Do not copy ASR's + split-and-text-stitch semantics unless they are valid for that model. +4. Flatten prepared model inputs and save, for each item, its original item + index, parent index, segment ordinal, and any timestamp or side-effect + metadata needed later. +5. Compute duration from the final resampled or sliced waveform. +6. Preserve input order when local bucketing is off; stable-sort by ascending + duration when it is on. +7. With bucketing off, greedily pack the input-order sequence. With bucketing + on, use dynamic programming over sorted contiguous spans to minimize call + count and then total padded seconds, using + `span_count * span_max_duration` as each span's cost. +8. Call the adapter once per planned group, require a complete ordered result + mapping, scatter by saved item index, and only then reassemble or write + parent outputs. Native model-call behavior remains adapter-specific. +9. Test enabled and disabled modes against the same correctness oracle. + +Do not introduce buffering across `process_batch()` calls to improve the +duration distribution. That changes latency, failure, ownership, and end-of- +stream behavior and would require a separate queue/flush design. + +### Model-specific cautions + +- **ASR:** the current contract concatenates independently decoded segment + text in temporal order. +- **SED:** preserve frame arrays, valid-frame counts, timestamps, and sidecar + paths when scattering. Splitting also requires frame-offset reconstruction. +- **VAD:** independent splits can change onset/offset decisions and merged + speech regions. +- **Diarization:** speaker clustering and identity may depend on the whole + recording; arbitrary segmentation is not equivalent. +- **Alignment:** flattened segment metadata must remain paired with the right + hypothesis after reordered calls. +- **File-producing stages:** derive paths from saved parent metadata, never + from reordered call position. + +Reordering independent whole inputs is often safe. Segmentation and stitching +always require a model-specific correctness design. + +## Testing and parity + +Unit tests should use a recording adapter stub and cover: + +- missing, Boolean, negative, `NaN`, and infinite budgets, plus model limits + shorter than one sample; +- an empty input and zero-duration audio; +- exact budget fills and a candidate that begins the next call; +- original-order and stable duration-order call membership; +- equal-duration stability; +- the `[1, 2, 2]` budget-`4` case that distinguishes the optimal boundary + plan from sorted greedy packing; +- planning across segments from multiple rows in one call, but never across + two `process_batch()` calls; +- exact model-duration boundaries, multi-segment parents, and every final + remainder; +- result scatter, parent stitching, skip/error paths, and wrong result counts; +- proof that `batch_size` does not impose an adapter item-count cap. + +For a local parity run, use a fixed cohort with short, medium, boundary-equal, +long, and segmented-remainder audio. Hold model settings, software, hardware, +input order, and candidate windows constant. First verify row IDs, row counts, +skip classifications, transcript equality or the documented tolerance, and +exact-once segment coverage. Only then compare padding efficiency, throughput, +latency, RAM, and VRAM. Local bucketing guarantees the planning contract, not +a performance win for every workload. + +## Explicit non-goals + +This design intentionally does not provide: + +- global, partition-wide, or cross-worker bucketing; +- planning across multiple `process_batch()` calls; +- queues, timers, flush hooks, or end-of-stream state; +- static bucket edges or per-bucket configuration; +- separate adapter or per-duration-group item-count controls; +- feature-weighted cost estimation; +- reordering beyond the stable duration sort or optimization across + noncontiguous subsets; +- adaptive or OOM-retry scheduling; +- automatic budget tuning; +- a guarantee of higher throughput for every model and distribution. + +For `N` segments, planning is `O(N)` with bucketing off. With bucketing on, +stable sorting costs `O(N log N)` and the exact boundary dynamic program costs +`O(N^2)`, for `O(N^2)` total time and `O(N)` auxiliary storage. Given the same +ordered prepared segments and configuration, the plan is deterministic. + +## Source and test map + +- [`asr/stage.py`](asr/stage.py): waveform preparation, unconditional + segmentation, local planning, adapter execution, scatter, and stitching. +- [`model_input_segmentation.py`](../model_input_segmentation.py): contiguous + model-safe segment planning and validation. +- [`inference/base.py`](base.py): shared adapter lifecycle and input handling; + deliberately not a generic batching planner. +- [`models/asr/base.py`](../../../models/asr/base.py): ordered ASR adapter input + and result contract. +- [`test_asr_stage.py`](../../../../tests/stages/audio/inference/test_asr_stage.py): + actor budgets, mode behavior, local scope, segmentation, ordering, and error + coverage. +- [`test_model_input_segmentation.py`](../../../../tests/stages/audio/test_model_input_segmentation.py): + exact boundaries, remainders, one-sample limits, and zero-length inputs. +- [FastConformer tutorial](../../../../tutorials/audio/nemo_fastconformer/README.md): + runnable ASR configuration and usage. diff --git a/nemo_curator/stages/audio/inference/asr/stage.py b/nemo_curator/stages/audio/inference/asr/stage.py index f7342ecbb0..2970a46334 100644 --- a/nemo_curator/stages/audio/inference/asr/stage.py +++ b/nemo_curator/stages/audio/inference/asr/stage.py @@ -21,7 +21,9 @@ from __future__ import annotations +import math from dataclasses import dataclass, field +from numbers import Real from typing import TYPE_CHECKING, Any, cast import numpy as np @@ -31,6 +33,10 @@ from nemo_curator.models.asr.base import ASRAdapter, ASRResult from nemo_curator.stages.audio.inference.base import AdapterInferenceStage +from nemo_curator.stages.audio.model_input_segmentation import ( + plan_audio_segments, + resolve_max_model_input_duration, +) from nemo_curator.stages.resources import Resources if TYPE_CHECKING: @@ -95,6 +101,8 @@ _NOTES_KEY = "additional_notes" _MONO_DIMENSIONS = 1 _CHANNEL_FIRST_DIMENSIONS = 2 +_PADDED_SECONDS_REL_TOL = 1e-12 +_PADDED_SECONDS_ABS_TOL = 1e-9 def _set_note(task_data: dict[str, Any], stage_name: str, value: str) -> None: @@ -112,11 +120,19 @@ class ASRStage(AdapterInferenceStage[ASRAdapter]): The stage writes ``pred_text_key`` and optional control columns ``_skipme`` and ``additional_notes``. When ``extras_key`` is configured, it also writes non-empty adapter metadata as one nested dictionary under that key. + + Audio longer than ``max_inference_duration_s`` is always split into + model-safe segments and stitched back to one result per parent row. Every + segment prepared by one backend-provided ``process_batch`` call is packed + into adapter calls bounded by ``max_audio_sec_per_actor``. Enabling + ``local_bucketing`` first orders those segments by duration to reduce GPU + padding; disabling it preserves their input order. """ # Adapter selection. adapter_target: str model_id: str + max_audio_sec_per_actor: float name: str = "ASR_inference" # Task I/O keys. @@ -140,6 +156,8 @@ class ASRStage(AdapterInferenceStage[ASRAdapter]): resources: Resources = field(default_factory=lambda: Resources(gpus=1.0)) batch_size: int = 32 + max_inference_duration_s: float = 2400.0 + local_bucketing: bool = False def __post_init__(self) -> None: super().__post_init__() @@ -163,10 +181,35 @@ def __post_init__(self) -> None: if int(self.target_sample_rate) <= 0: msg = f"ASRStage.target_sample_rate must be > 0, got {self.target_sample_rate}" raise ValueError(msg) + self.max_inference_duration_s = resolve_max_model_input_duration( + max_duration_s=self.max_inference_duration_s, + owner="ASRStage", + ) + self.max_audio_sec_per_actor = self._validate_max_audio_sec_per_actor(self.max_audio_sec_per_actor) + if self.max_inference_duration_s > self.max_audio_sec_per_actor: + msg = ( + "ASRStage.max_inference_duration_s must be <= max_audio_sec_per_actor; " + f"got {self.max_inference_duration_s} > {self.max_audio_sec_per_actor}" + ) + raise ValueError(msg) + if not isinstance(self.local_bucketing, bool): + msg = f"ASRStage.local_bucketing must be a bool, got {type(self.local_bucketing).__name__}" + raise TypeError(msg) self.batch_size = int(self.batch_size) self.target_sample_rate = int(self.target_sample_rate) self._supported_language_codes = self._normalise_supported_language_codes(self.supported_language_codes) + @staticmethod + def _validate_max_audio_sec_per_actor(value: object) -> float: + if isinstance(value, bool) or not isinstance(value, Real): + msg = f"ASRStage.max_audio_sec_per_actor must be numeric, got {type(value).__name__}" + raise TypeError(msg) + maximum = float(value) + if not math.isfinite(maximum) or maximum <= 0: + msg = f"ASRStage.max_audio_sec_per_actor must be finite and > 0, got {value}" + raise ValueError(msg) + return maximum + @staticmethod def _normalise_supported_language_codes(value: object) -> set[str] | None: """Normalize an optional adapter-specific supported-language allowlist.""" @@ -318,7 +361,7 @@ def run_inference(self, items: list[dict[str, Any]]) -> list[ASRResult]: """Transcribe one stage batch via the adapter.""" supported_indices = [index for index, item in enumerate(items) if self._is_language_supported(item)] by_index: dict[int, ASRResult] = {} - adapter_indices: list[int] = [] + adapter_parent_indices: list[int] = [] adapter_items: list[dict[str, Any]] = [] for index in supported_indices: item = items[index] @@ -344,26 +387,35 @@ def run_inference(self, items: list[dict[str, Any]]) -> list[ASRResult]: ) by_index[index] = ASRResult(text="", skipped=True, skip_reason="audio_load_error") continue - adapter_indices.append(index) - adapter_items.append( - { - "waveform": waveform, - "sample_rate": self.target_sample_rate, - "language": item["language"], - "language_code": item["language_code"], - "task_id": item["task_id"], - } + segments = plan_audio_segments( + num_samples=int(waveform.shape[0]), + sample_rate=self.target_sample_rate, + max_duration_s=self.max_inference_duration_s, + owner="ASRStage", ) + for segment in segments: + adapter_parent_indices.append(index) + adapter_items.append( + { + "waveform": np.ascontiguousarray( + waveform[segment.start_sample : segment.stop_sample], + dtype=np.float32, + ), + "sample_rate": self.target_sample_rate, + "audio_seconds": segment.duration_s, + "language": item["language"], + "language_code": item["language_code"], + "task_id": item["task_id"], + } + ) if adapter_items: - adapter_results = self._adapter.transcribe_batch(adapter_items) - if len(adapter_results) != len(adapter_items): - msg = ( - f"Adapter returned {len(adapter_results)} results for " - f"{len(adapter_items)} supported items (must match 1:1)" - ) - raise RuntimeError(msg) - by_index.update(zip(adapter_indices, adapter_results, strict=True)) + adapter_results = self._run_adapter_batches(adapter_items) + per_parent: dict[int, list[ASRResult]] = {} + for parent_index, result in zip(adapter_parent_indices, adapter_results, strict=True): + per_parent.setdefault(parent_index, []).append(result) + for parent_index, chunk_results in per_parent.items(): + by_index[parent_index] = self._stitch_chunk_results(chunk_results) return [ by_index.get( index, @@ -381,6 +433,167 @@ def run_inference(self, items: list[dict[str, Any]]) -> list[ASRResult]: for index, item in enumerate(items) ] + @staticmethod + def _stitch_chunk_results(results: list[ASRResult]) -> ASRResult: + """Join ordered chunk outputs into one parent-row result.""" + if not results: + return ASRResult(text="", skipped=True, skip_reason="empty_audio") + if len(results) == 1: + return results[0] + + texts = [text for result in results if (text := (result.text or "").strip())] + any_skipped = any(result.skipped for result in results) + skip_reason = next((result.skip_reason for result in results if result.skip_reason), None) + unsupported_language = next( + (result.unsupported_language for result in results if result.unsupported_language), + None, + ) + extras: dict[str, Any] = {} + for result in results: + extras.update(result.extras) + return ASRResult( + text=" ".join(texts), + skipped=any_skipped, + skip_reason=skip_reason if any_skipped else None, + unsupported_language=unsupported_language, + extras=extras, + ) + + def _run_adapter_batches(self, items: list[dict[str, Any]]) -> list[ASRResult]: + """Run capacity-bounded adapter calls and restore segment order.""" + if self._adapter is None: + msg = "Adapter not initialized - setup() was not called" + raise RuntimeError(msg) + + sub_batches = self._plan_adapter_batches(items) + + aligned: list[ASRResult | None] = [None] * len(items) + for indices, sub_items in sub_batches: + sub_results = self._adapter.transcribe_batch(sub_items) + if len(sub_results) != len(sub_items): + msg = ( + f"Adapter returned {len(sub_results)} results for " + f"{len(sub_items)} supported items (must match 1:1)" + ) + raise RuntimeError(msg) + for index, result in zip(indices, sub_results, strict=True): + aligned[index] = result + + if any(result is None for result in aligned): + msg = "Local batch planning did not produce a result for every supported item" + raise RuntimeError(msg) + return [result for result in aligned if result is not None] + + def _plan_adapter_batches( + self, + items: list[dict[str, Any]], + ) -> list[tuple[list[int], list[dict[str, Any]]]]: + """Optimally pack one finite segment list under the padded-audio budget. + + The proxy cost of an adapter call is its longest audio duration times + its item count. This models the padded tensor work more closely than a + sum of unpadded durations. With local bucketing enabled, dynamic + programming over the stable duration order first minimizes adapter-call + count and then total padded seconds. The budget is enforced in both + modes. + """ + indexed_items = [(index, item, item["audio_seconds"]) for index, item in enumerate(items)] + + if not self.local_bucketing: + return self._pack_in_order(indexed_items) + + indexed_items.sort(key=lambda indexed_item: indexed_item[2]) + return self._pack_duration_sorted(indexed_items) + + def _pack_duration_sorted( + self, + indexed_items: list[tuple[int, dict[str, Any], float]], + ) -> list[tuple[list[int], list[dict[str, Any]]]]: + """Find the exact lexicographic optimum over sorted contiguous spans.""" + item_count = len(indexed_items) + if item_count == 0: + return [] + + # best_score[start] is the exact optimum for the suffix beginning at + # start. A batch is always one contiguous span in stable duration + # order, so the recurrence considers every possible next boundary. + best_score: list[tuple[int, float] | None] = [None] * (item_count + 1) + next_boundary = [item_count] * item_count + best_score[item_count] = (0, 0.0) + + for start in range(item_count - 1, -1, -1): + for stop in range(start + 1, item_count + 1): + padded_seconds = indexed_items[stop - 1][2] * (stop - start) + if not self._fits_audio_budget(padded_seconds): + break + + suffix_score = best_score[stop] + if suffix_score is None: # pragma: no cover - every singleton is feasible + continue + candidate_score = (suffix_score[0] + 1, suffix_score[1] + padded_seconds) + current_score = best_score[start] + if current_score is None or candidate_score < current_score: + best_score[start] = candidate_score + next_boundary[start] = stop + + planned: list[tuple[list[int], list[dict[str, Any]]]] = [] + start = 0 + while start < item_count: + stop = next_boundary[start] + batch = indexed_items[start:stop] + planned.append( + ( + [index for index, _item, _audio_seconds in batch], + [item for _index, item, _audio_seconds in batch], + ), + ) + start = stop + return planned + + def _pack_in_order( + self, + indexed_items: list[tuple[int, dict[str, Any], float]], + ) -> list[tuple[list[int], list[dict[str, Any]]]]: + """Greedily preserve input order while enforcing padded capacity.""" + planned: list[tuple[list[int], list[dict[str, Any]]]] = [] + current: list[tuple[int, dict[str, Any], float]] = [] + current_max_duration = 0.0 + + for indexed_item in indexed_items: + audio_seconds = indexed_item[2] + candidate_max_duration = max(current_max_duration, audio_seconds) + candidate_padded_seconds = candidate_max_duration * (len(current) + 1) + if current and not self._fits_audio_budget(candidate_padded_seconds): + planned.append( + ( + [index for index, _item, _audio_seconds in current], + [item for _index, item, _audio_seconds in current], + ), + ) + current = [] + current_max_duration = 0.0 + + current.append(indexed_item) + current_max_duration = max(current_max_duration, audio_seconds) + + if current: + planned.append( + ( + [index for index, _item, _audio_seconds in current], + [item for _index, item, _audio_seconds in current], + ), + ) + return planned + + def _fits_audio_budget(self, padded_seconds: float) -> bool: + """Treat representation-level equality as within the configured budget.""" + return padded_seconds <= self.max_audio_sec_per_actor or math.isclose( + padded_seconds, + self.max_audio_sec_per_actor, + rel_tol=_PADDED_SECONDS_REL_TOL, + abs_tol=_PADDED_SECONDS_ABS_TOL, + ) + def assemble( self, tasks: list[AudioTask], diff --git a/nemo_curator/stages/audio/model_input_segmentation.py b/nemo_curator/stages/audio/model_input_segmentation.py new file mode 100644 index 0000000000..79cc2c171b --- /dev/null +++ b/nemo_curator/stages/audio/model_input_segmentation.py @@ -0,0 +1,91 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared audio model-input segmentation helpers. + +Segmentation creates model-safe work units. Duration-aware bucketing is a +separate packing step that consumes these bounded units. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + + +@dataclass(frozen=True) +class AudioSegment: + """One contiguous audio model-input segment in sample coordinates.""" + + start_sample: int + stop_sample: int + duration_s: float + + +def resolve_max_model_input_duration( + *, + max_duration_s: float, + owner: str, +) -> float: + """Validate and normalize the model-input duration ceiling.""" + + maximum = float(max_duration_s) + if not math.isfinite(maximum) or maximum <= 0: + msg = f"{owner}.max_inference_duration_s must be finite and > 0 s, got {max_duration_s}" + raise ValueError(msg) + return maximum + + +def plan_audio_segments( + *, + num_samples: int, + sample_rate: int, + max_duration_s: float, + owner: str, +) -> tuple[AudioSegment, ...]: + """Create bounded contiguous segment specs for one audio input.""" + + maximum = resolve_max_model_input_duration( + max_duration_s=max_duration_s, + owner=owner, + ) + if sample_rate <= 0: + msg = f"{owner}.sample_rate must be > 0, got {sample_rate}" + raise ValueError(msg) + max_samples = int(maximum * float(sample_rate)) + if max_samples < 1: + msg = f"{owner}.max_inference_duration_s must cover at least one sample at sample_rate={sample_rate}" + raise ValueError(msg) + if num_samples <= 0: + return ( + AudioSegment( + start_sample=0, + stop_sample=0, + duration_s=0.0, + ), + ) + + starts = list(range(0, int(num_samples), max_samples)) + segments: list[AudioSegment] = [] + for start in starts: + stop = min(start + max_samples, int(num_samples)) + duration_s = float(stop - start) / float(sample_rate) + segments.append( + AudioSegment( + start_sample=start, + stop_sample=stop, + duration_s=duration_s, + ) + ) + return tuple(segments) diff --git a/tests/config/test_run.py b/tests/config/test_run.py index 7af6e897c7..fb238bc184 100644 --- a/tests/config/test_run.py +++ b/tests/config/test_run.py @@ -18,7 +18,7 @@ from unittest.mock import MagicMock, patch import pytest -from hydra import compose, initialize_config_dir +from hydra import compose, initialize, initialize_config_dir from omegaconf import OmegaConf from nemo_curator.config.run import create_executor_from_yaml, create_pipeline_from_yaml, main @@ -465,6 +465,9 @@ def test_qwen_tutorial_yaml_matches_reference_runner_config(): "ur", ] assert stage.batch_size == 32 + assert stage.max_audio_sec_per_actor == 2400.0 + assert stage.max_inference_duration_s == 2400.0 + assert stage.local_bucketing is True assert stage.resources.gpus == 2 assert dict(stage.adapter_kwargs) == { "revision": "abc123", @@ -566,6 +569,9 @@ def test_qwen_asr_tutorial_yaml_uses_generic_adapter_contract(): "ro", ] assert stage.batch_size == 128 + assert stage.max_audio_sec_per_actor == 240.0 + assert stage.max_inference_duration_s == 120.0 + assert stage.local_bucketing is True assert stage.resources.gpus == 1 assert dict(stage.adapter_kwargs) == { "revision": "abc123", @@ -731,6 +737,9 @@ def test_faster_whisper_tutorial_yaml_matches_reference_contract(): ["custom_prediction", "_skipme", "additional_notes", "asr_extras"], ) assert stage.batch_size == 128 + assert stage.max_audio_sec_per_actor == 2400.0 + assert stage.max_inference_duration_s == 2400.0 + assert stage.local_bucketing is True assert stage.resources.gpus == 1 assert dict(stage.adapter_kwargs) == { "revision": "abc123", @@ -778,6 +787,9 @@ def test_nemo_fastconformer_tutorial_yaml_uses_shared_adapter_contract(): assert stage.target_sample_rate == 16000 assert stage.pred_text_key == "custom_prediction" assert stage.batch_size == 16 + assert stage.max_audio_sec_per_actor == 240.0 + assert stage.max_inference_duration_s == 120.0 + assert stage.local_bucketing is True assert stage.resources.gpus == 1 assert dict(stage.adapter_kwargs) == { "num_workers": 0, @@ -789,6 +801,25 @@ def test_nemo_fastconformer_tutorial_yaml_uses_shared_adapter_contract(): assert executor.config == {} +def test_nemo_fastconformer_tutorial_accepts_local_bucketing_config() -> None: + with initialize(config_path="../../tutorials/audio/nemo_fastconformer", version_base=None): + cfg = compose( + config_name="pipeline", + overrides=[ + "manifest_path=tests/fixtures/audio/tagging/sample_input.jsonl", + "max_audio_sec_per_actor=360.0", + "max_inference_duration_s=90.0", + "local_bucketing=false", + ], + ) + + _, _, stage, _ = create_pipeline_from_yaml(cfg, log_config=False).stages + + assert stage.max_audio_sec_per_actor == 360.0 + assert stage.max_inference_duration_s == 90.0 + assert stage.local_bucketing is False + + def test_run_cli_defaults_to_pipeline_config_for_fastconformer() -> None: repo_root = Path(__file__).parents[2] result = subprocess.run( # noqa: S603 diff --git a/tests/models/asr/test_nemo_asr.py b/tests/models/asr/test_nemo_asr.py index ea6a98182d..c24dda7d01 100644 --- a/tests/models/asr/test_nemo_asr.py +++ b/tests/models/asr/test_nemo_asr.py @@ -29,6 +29,8 @@ from nemo_curator.models.asr import nemo_asr from nemo_curator.models.asr.base import ASRAdapter from nemo_curator.models.asr.nemo_asr import NeMoASRAdapter +from nemo_curator.stages.audio.inference.asr.stage import ASRStage +from nemo_curator.tasks import AudioTask _MODEL_ID = "nvidia/stt_en_fastconformer_ctc_large" _SAMPLE_RATE = 16_000 @@ -139,6 +141,40 @@ def test_transcribe_batch_uses_one_exact_nemo_batch() -> None: assert len(kwargs["audio"]) == 2 +def test_asr_stage_drives_nemo_adapter_with_exact_local_batches() -> None: + model = _mock_model([]) + model.transcribe.side_effect = [ + ["short-a", "short-b"], + ["long"], + ] + adapter = NeMoASRAdapter() + adapter._model = model + stage = ASRStage( + adapter_target="nemo_curator.models.asr.nemo_asr.NeMoASRAdapter", + model_id=_MODEL_ID, + max_audio_sec_per_actor=4.0, + max_inference_duration_s=4.0, + local_bucketing=True, + waveform_key="waveform", + sample_rate_key="sampling_rate", + keep_waveform=True, + ) + stage._adapter = adapter + tasks = [ + AudioTask(data={"waveform": np.zeros(_SAMPLE_RATE, dtype=np.float32), "sampling_rate": _SAMPLE_RATE}), + AudioTask(data={"waveform": np.zeros(4 * _SAMPLE_RATE, dtype=np.float32), "sampling_rate": _SAMPLE_RATE}), + AudioTask( + data={"waveform": np.zeros(int(1.5 * _SAMPLE_RATE), dtype=np.float32), "sampling_rate": _SAMPLE_RATE} + ), + ] + + results = stage.process_batch(tasks) + + assert [call.kwargs["batch_size"] for call in model.transcribe.call_args_list] == [2, 1] + assert [len(call.kwargs["audio"]) for call in model.transcribe.call_args_list] == [2, 1] + assert [task.data["pred_text"] for task in results] == ["short-a", "long", "short-b"] + + def test_transcribe_batch_preserves_empty_positions() -> None: model = _mock_model(["valid"]) adapter = NeMoASRAdapter() diff --git a/tests/models/asr/test_qwen_asr.py b/tests/models/asr/test_qwen_asr.py index 7805cfcdc1..549d16a7e1 100644 --- a/tests/models/asr/test_qwen_asr.py +++ b/tests/models/asr/test_qwen_asr.py @@ -127,6 +127,7 @@ def test_asr_stage_prefetches_qwen_adapter_with_adapter_owned_revision() -> None stage = ASRStage( adapter_target="nemo_curator.models.asr.qwen_asr.QwenASRAdapter", model_id="Qwen/Qwen3-ASR-0.6B", + max_audio_sec_per_actor=2400.0, adapter_kwargs={"revision": "abc123"}, ) with ( @@ -341,6 +342,7 @@ def test_asr_stage_drives_qwen_adapter_end_to_end() -> None: stage = ASRStage( adapter_target="nemo_curator.models.asr.qwen_asr.QwenASRAdapter", model_id="Qwen/Qwen3-ASR-0.6B", + max_audio_sec_per_actor=2400.0, batch_size=2, waveform_key="waveform", sample_rate_key="sample_rate", diff --git a/tests/stages/audio/inference/test_asr_stage.py b/tests/stages/audio/inference/test_asr_stage.py index b276a47a39..47d884f6ff 100644 --- a/tests/stages/audio/inference/test_asr_stage.py +++ b/tests/stages/audio/inference/test_asr_stage.py @@ -38,6 +38,10 @@ def _make_stage( # noqa: PLR0913 *, default_language: str | None = None, batch_size: int = 32, + max_audio_sec_per_actor: float = 2400.0, + local_bucketing: bool = False, + target_sample_rate: int = _SR, + max_inference_duration_s: float = 2400.0, supported_language_codes: list[str] | None = None, skip_if_output_exists: bool = False, waveform_key: str | None = None, @@ -52,6 +56,10 @@ def _make_stage( # noqa: PLR0913 pred_text_key="pred_text", default_language=default_language, batch_size=batch_size, + max_audio_sec_per_actor=max_audio_sec_per_actor, + local_bucketing=local_bucketing, + target_sample_rate=target_sample_rate, + max_inference_duration_s=max_inference_duration_s, supported_language_codes=supported_language_codes, skip_if_output_exists=skip_if_output_exists, waveform_key=waveform_key, @@ -118,16 +126,22 @@ def test_basic_inference() -> None: assert set(inferred_item) == { "waveform", "sample_rate", + "audio_seconds", "language", "language_code", "task_id", } assert inferred_item["waveform"].shape == (_SR,) assert inferred_item["sample_rate"] == _SR + assert inferred_item["audio_seconds"] == 1.0 def test_adapter_not_initialized_raises() -> None: - stage = ASRStage(adapter_target=_QWEN_ADAPTER_TARGET, model_id="mock/model") + stage = ASRStage( + adapter_target=_QWEN_ADAPTER_TARGET, + model_id="mock/model", + max_audio_sec_per_actor=2400.0, + ) with pytest.raises(RuntimeError, match="setup"): stage.process_batch([_make_task()]) @@ -144,6 +158,310 @@ def test_multi_task_batch_preserves_order() -> None: assert results[1].data["pred_text"] == "text2" +def test_local_bucketing_groups_all_rows_and_restores_task_order() -> None: + stage = _make_stage( + waveform_key="waveform", + keep_waveform=True, + max_audio_sec_per_actor=50.0, + max_inference_duration_s=50.0, + local_bucketing=True, + ) + tasks = [ + _make_waveform_task(waveform=np.zeros(5 * _SR, dtype=np.float32)), + _make_waveform_task(waveform=np.zeros(40 * _SR, dtype=np.float32)), + _make_waveform_task(waveform=np.zeros(15 * _SR, dtype=np.float32)), + ] + for index, task in enumerate(tasks): + task.task_id = f"task-{index}" + stage._adapter.transcribe_batch.side_effect = lambda items: [ + ASRResult(text=str(item["task_id"])) for item in items + ] + + results = stage.process_batch(tasks) + + durations_by_call = [ + [item["audio_seconds"] for item in call.args[0]] for call in stage._adapter.transcribe_batch.call_args_list + ] + assert durations_by_call == [[5.0, 15.0], [40.0]] + assert [task.data["pred_text"] for task in results] == ["task-0", "task-1", "task-2"] + + +@pytest.mark.parametrize( + ("local_bucketing", "expected_call_durations", "expected_call_task_ids"), + [ + (False, [[1.0, 4.0], [1.0]], [["task-0", "task-1"], ["task-2"]]), + (True, [[1.0, 1.0], [4.0]], [["task-0", "task-2"], ["task-1"]]), + ], + ids=["input-order", "duration-order"], +) +def test_actor_audio_budget_is_enforced_with_bucketing_on_or_off( + local_bucketing: bool, + expected_call_durations: list[list[float]], + expected_call_task_ids: list[list[str]], +) -> None: + stage = _make_stage( + waveform_key="waveform", + keep_waveform=True, + max_audio_sec_per_actor=8.0, + max_inference_duration_s=8.0, + local_bucketing=local_bucketing, + ) + durations = [1.0, 4.0, 1.0] + tasks = [_make_waveform_task(waveform=np.zeros(int(duration * _SR), dtype=np.float32)) for duration in durations] + for index, task in enumerate(tasks): + task.task_id = f"task-{index}" + stage._adapter.transcribe_batch.side_effect = lambda items: [ + ASRResult(text=str(item["task_id"])) for item in items + ] + + results = stage.process_batch(tasks) + + call_durations = [ + [item["audio_seconds"] for item in call.args[0]] for call in stage._adapter.transcribe_batch.call_args_list + ] + call_task_ids = [ + [item["task_id"] for item in call.args[0]] for call in stage._adapter.transcribe_batch.call_args_list + ] + assert call_durations == expected_call_durations + assert call_task_ids == expected_call_task_ids + assert [task.data["pred_text"] for task in results] == ["task-0", "task-1", "task-2"] + assert all(len(call) * max(call) <= 8.0 for call in call_durations) + + +def test_local_bucketing_minimizes_padded_seconds_after_call_count() -> None: + stage = _make_stage( + waveform_key="waveform", + keep_waveform=True, + max_audio_sec_per_actor=4.0, + max_inference_duration_s=4.0, + local_bucketing=True, + ) + durations = [1.0, 2.0, 2.0] + tasks = [_make_waveform_task(waveform=np.zeros(int(duration * _SR), dtype=np.float32)) for duration in durations] + for index, task in enumerate(tasks): + task.task_id = f"task-{index}" + stage._adapter.transcribe_batch.side_effect = lambda items: [ + ASRResult(text=str(item["task_id"])) for item in items + ] + + results = stage.process_batch(tasks) + + durations_by_call = [ + [item["audio_seconds"] for item in call.args[0]] for call in stage._adapter.transcribe_batch.call_args_list + ] + assert durations_by_call == [[1.0], [2.0, 2.0]] + assert sum(len(call) * max(call) for call in durations_by_call) == 5.0 + assert [task.data["pred_text"] for task in results] == ["task-0", "task-1", "task-2"] + + +@pytest.mark.parametrize("local_bucketing", [False, True]) +def test_actor_budget_accepts_decimal_roundoff_at_exact_boundary(local_bucketing: bool) -> None: + stage = _make_stage( + max_audio_sec_per_actor=0.3, + max_inference_duration_s=0.3, + local_bucketing=local_bucketing, + ) + items = [{"audio_seconds": 0.1, "name": name} for name in ["a", "b", "c"]] + + plan = stage._plan_adapter_batches(items) + + assert [indices for indices, _items in plan] == [[0, 1, 2]] + assert [[item["name"] for item in batch] for _indices, batch in plan] == [["a", "b", "c"]] + + +def test_local_bucketing_is_scoped_to_each_process_batch_call() -> None: + stage = _make_stage(waveform_key="waveform", keep_waveform=True, local_bucketing=True) + first = _make_waveform_task() + second = _make_waveform_task() + first.task_id = "first-window" + second.task_id = "second-window" + stage._adapter.transcribe_batch.side_effect = lambda items: [ + ASRResult(text=str(item["task_id"])) for item in items + ] + + first_result = stage.process_batch([first]) + second_result = stage.process_batch([second]) + + task_ids_by_call = [ + [item["task_id"] for item in call.args[0]] for call in stage._adapter.transcribe_batch.call_args_list + ] + assert task_ids_by_call == [["first-window"], ["second-window"]] + assert first_result[0].data["pred_text"] == "first-window" + assert second_result[0].data["pred_text"] == "second-window" + + +def test_batch_size_does_not_cap_adapter_calls() -> None: + stage = _make_stage( + waveform_key="waveform", + keep_waveform=True, + batch_size=2, + max_audio_sec_per_actor=3.0, + max_inference_duration_s=3.0, + ) + tasks = [_make_waveform_task() for _ in range(3)] + stage._adapter.transcribe_batch.return_value = [ + ASRResult(text="a"), + ASRResult(text="b"), + ASRResult(text="c"), + ] + + results = stage.process_batch(tasks) + + assert [len(call.args[0]) for call in stage._adapter.transcribe_batch.call_args_list] == [3] + assert [task.data["pred_text"] for task in results] == ["a", "b", "c"] + + +@pytest.mark.parametrize("local_bucketing", [False, True]) +def test_model_safe_segmentation_preserves_samples_and_stitches_parent_order(local_bucketing: bool) -> None: + sample_rate = 10 + stage = _make_stage( + waveform_key="waveform", + keep_waveform=True, + target_sample_rate=sample_rate, + max_audio_sec_per_actor=9.0, + max_inference_duration_s=3.0, + local_bucketing=local_bucketing, + ) + long_waveform = np.arange(5 * sample_rate, dtype=np.float32) + short_waveform = np.arange(2 * sample_rate, dtype=np.float32) + 100 + tasks = [ + _make_waveform_task(waveform=long_waveform, sample_rate=sample_rate), + _make_waveform_task(waveform=short_waveform, sample_rate=sample_rate), + ] + + def transcribe(items: list[dict[str, object]]) -> list[ASRResult]: + results: list[ASRResult] = [] + for item in items: + item_waveform = np.asarray(item["waveform"]) + if item_waveform[0] == 0: + text = "first" + elif item_waveform[0] == 30: + text = "tail" + else: + text = "second" + results.append(ASRResult(text=text)) + return results + + stage._adapter.transcribe_batch.side_effect = transcribe + + results = stage.process_batch(tasks) + + inferred = [item for call in stage._adapter.transcribe_batch.call_args_list for item in call.args[0]] + assert sorted(item["audio_seconds"] for item in inferred) == [2.0, 2.0, 3.0] + inferred_by_start = {float(item["waveform"][0]): item for item in inferred} + np.testing.assert_array_equal(inferred_by_start[0.0]["waveform"], long_waveform[: 3 * sample_rate]) + np.testing.assert_array_equal(inferred_by_start[30.0]["waveform"], long_waveform[3 * sample_rate :]) + np.testing.assert_array_equal(inferred_by_start[100.0]["waveform"], short_waveform) + np.testing.assert_array_equal( + np.concatenate([inferred_by_start[0.0]["waveform"], inferred_by_start[30.0]["waveform"]]), + long_waveform, + ) + assert all(item["waveform"].dtype == np.float32 for item in inferred) + assert all(item["waveform"].flags.c_contiguous for item in inferred) + assert [task.data["pred_text"] for task in results] == ["first tail", "second"] + + +def test_segmented_parent_obeys_actor_audio_budget_before_stitching() -> None: + sample_rate = 10 + waveform = np.arange(10 * sample_rate, dtype=np.float32) + stage = _make_stage( + waveform_key="waveform", + keep_waveform=True, + target_sample_rate=sample_rate, + max_audio_sec_per_actor=6.0, + max_inference_duration_s=3.0, + local_bucketing=True, + ) + stage._adapter.transcribe_batch.side_effect = lambda items: [ + ASRResult(text=f"chunk-{int(np.asarray(item['waveform'])[0]) // 30}") for item in items + ] + + result = stage.process_batch([_make_waveform_task(waveform=waveform, sample_rate=sample_rate)])[0] + + assert [len(call.args[0]) for call in stage._adapter.transcribe_batch.call_args_list] == [2, 2] + inferred = [item for call in stage._adapter.transcribe_batch.call_args_list for item in call.args[0]] + assert [ + [item["audio_seconds"] for item in call.args[0]] for call in stage._adapter.transcribe_batch.call_args_list + ] == [ + [1.0, 3.0], + [3.0, 3.0], + ] + ordered = sorted(inferred, key=lambda item: float(item["waveform"][0])) + np.testing.assert_array_equal(np.concatenate([item["waveform"] for item in ordered]), waveform) + assert result.data["pred_text"] == "chunk-0 chunk-1 chunk-2 chunk-3" + + +def test_long_row_tail_can_co_bucket_after_model_safe_segmentation() -> None: + sample_rate = 10 + stage = _make_stage( + waveform_key="waveform", + keep_waveform=True, + target_sample_rate=sample_rate, + max_audio_sec_per_actor=240.0, + max_inference_duration_s=120.0, + local_bucketing=True, + ) + tasks = [ + _make_waveform_task(waveform=np.zeros(250 * sample_rate, dtype=np.float32), sample_rate=sample_rate), + _make_waveform_task(waveform=np.zeros(10 * sample_rate, dtype=np.float32), sample_rate=sample_rate), + _make_waveform_task(waveform=np.zeros(sample_rate, dtype=np.float32), sample_rate=sample_rate), + ] + stage._adapter.transcribe_batch.side_effect = [ + [ASRResult(text="tiny"), ASRResult(text="tail"), ASRResult(text="ten")], + [ASRResult(text="long-0"), ASRResult(text="long-1")], + ] + + results = stage.process_batch(tasks) + + durations_by_call = [ + [item["audio_seconds"] for item in call.args[0]] for call in stage._adapter.transcribe_batch.call_args_list + ] + assert durations_by_call == [[1.0, 10.0, 10.0], [120.0, 120.0]] + assert [task.data["pred_text"] for task in results] == ["long-0 long-1 tail", "ten", "tiny"] + + +def test_segmented_parent_marks_partial_chunk_failure() -> None: + sample_rate = 10 + stage = _make_stage( + waveform_key="waveform", + keep_waveform=True, + target_sample_rate=sample_rate, + max_inference_duration_s=3.0, + ) + task = _make_waveform_task(waveform=np.zeros(5 * sample_rate, dtype=np.float32), sample_rate=sample_rate) + stage._adapter.transcribe_batch.return_value = [ + ASRResult(text="", skipped=True, skip_reason="empty_audio"), + ASRResult(text="recovered"), + ] + + result = stage.process_batch([task])[0] + + assert result.data["pred_text"] == "recovered" + assert result.data["_skipme"] == "empty_audio" + + +def test_segmented_parent_preserves_skip_reason_and_flat_adapter_extras() -> None: + sample_rate = 10 + stage = _make_stage( + waveform_key="waveform", + keep_waveform=True, + extras_key="asr_extras", + target_sample_rate=sample_rate, + max_inference_duration_s=3.0, + ) + task = _make_waveform_task(waveform=np.zeros(5 * sample_rate, dtype=np.float32), sample_rate=sample_rate) + stage._adapter.transcribe_batch.return_value = [ + ASRResult(text="", skipped=True, skip_reason="decode_failed", extras={"first_chunk": 0}), + ASRResult(text="", skipped=True, skip_reason="empty_audio", extras={"last_chunk": 1}), + ] + + result = stage.process_batch([task])[0] + + assert result.data["pred_text"] == "" + assert result.data["_skipme"] == "decode_failed" + assert result.data["asr_extras"] == {"first_chunk": 0, "last_chunk": 1} + + def test_audio_load_failure_skips_only_failed_item_and_preserves_order() -> None: stage = _make_stage() waveform = np.zeros(_SR, dtype=np.float32) @@ -306,6 +624,7 @@ def test_inputs_and_exact_output_contract() -> None: stage = ASRStage( adapter_target=_QWEN_ADAPTER_TARGET, model_id="mock/model", + max_audio_sec_per_actor=2400.0, pred_text_key="custom_prediction", ) _required, required_inputs = stage.inputs() @@ -355,6 +674,7 @@ def test_in_memory_input_contract_requires_waveform_and_sample_rate() -> None: stage = ASRStage( adapter_target=_QWEN_ADAPTER_TARGET, model_id="mock/model", + max_audio_sec_per_actor=2400.0, waveform_key="waveform", sample_rate_key="sampling_rate", ) @@ -413,6 +733,7 @@ def test_faster_whisper_empty_8khz_audio_preserves_reference_output(waveform: np stage = ASRStage( adapter_target=_FASTER_WHISPER_ADAPTER_TARGET, model_id="large-v3", + max_audio_sec_per_actor=2400.0, waveform_key="waveform", sample_rate_key="sampling_rate", supported_language_codes=["fil"], @@ -453,10 +774,65 @@ def test_invalid_target_sample_rate_is_rejected() -> None: ASRStage( adapter_target=_QWEN_ADAPTER_TARGET, model_id="mock/model", + max_audio_sec_per_actor=2400.0, target_sample_rate=0, ) +@pytest.mark.parametrize( + ("max_audio_sec_per_actor", "expected_exception", "match"), + [ + (0, ValueError, "max_audio_sec_per_actor must be finite and > 0"), + (-1, ValueError, "max_audio_sec_per_actor must be finite and > 0"), + (float("inf"), ValueError, "max_audio_sec_per_actor must be finite and > 0"), + (float("nan"), ValueError, "max_audio_sec_per_actor must be finite and > 0"), + ("2400", TypeError, "max_audio_sec_per_actor must be numeric"), + (True, TypeError, "max_audio_sec_per_actor must be numeric"), + ], +) +def test_invalid_max_audio_sec_per_actor_is_rejected( + max_audio_sec_per_actor: object, + expected_exception: type[Exception], + match: str, +) -> None: + with pytest.raises(expected_exception, match=match): + ASRStage( + adapter_target=_QWEN_ADAPTER_TARGET, + model_id="mock/model", + max_audio_sec_per_actor=max_audio_sec_per_actor, # type: ignore[arg-type] + ) + + +def test_invalid_local_bucketing_is_rejected() -> None: + with pytest.raises(TypeError, match="local_bucketing must be a bool"): + ASRStage( + adapter_target=_QWEN_ADAPTER_TARGET, + model_id="mock/model", + max_audio_sec_per_actor=2400.0, + local_bucketing=1, # type: ignore[arg-type] + ) + + +def test_actor_budget_must_fit_one_model_safe_segment() -> None: + with pytest.raises(ValueError, match="max_inference_duration_s must be <= max_audio_sec_per_actor"): + ASRStage( + adapter_target=_QWEN_ADAPTER_TARGET, + model_id="mock/model", + max_audio_sec_per_actor=10.0, + max_inference_duration_s=11.0, + ) + + +def test_invalid_max_inference_duration_is_rejected() -> None: + with pytest.raises(ValueError, match="max_inference_duration_s must be finite and > 0"): + ASRStage( + adapter_target=_QWEN_ADAPTER_TARGET, + model_id="mock/model", + max_audio_sec_per_actor=2400.0, + max_inference_duration_s=0, + ) + + def test_stage_requires_resampled_path_and_does_not_fallback_to_original_audio() -> None: stage = _make_stage() task = AudioTask(data={"audio_filepath": "/data/original.wav", "source_lang": "en"}) @@ -473,6 +849,7 @@ def test_empty_prediction_key_is_rejected() -> None: ASRStage( adapter_target=_QWEN_ADAPTER_TARGET, model_id="mock/model", + max_audio_sec_per_actor=2400.0, pred_text_key="", ) @@ -482,6 +859,7 @@ def test_empty_extras_key_is_rejected() -> None: ASRStage( adapter_target=_QWEN_ADAPTER_TARGET, model_id="mock/model", + max_audio_sec_per_actor=2400.0, extras_key=" ", ) @@ -492,6 +870,7 @@ def test_extras_key_cannot_collide_with_another_output(extras_key: str) -> None: ASRStage( adapter_target=_QWEN_ADAPTER_TARGET, model_id="mock/model", + max_audio_sec_per_actor=2400.0, extras_key=extras_key, ) @@ -502,6 +881,7 @@ def test_control_columns_cannot_be_used_as_prediction_key(pred_text_key: str) -> ASRStage( adapter_target=_QWEN_ADAPTER_TARGET, model_id="mock/model", + max_audio_sec_per_actor=2400.0, pred_text_key=pred_text_key, ) @@ -526,6 +906,7 @@ def test_setup_on_node_downloads_weights(mock_download: MagicMock) -> None: stage = ASRStage( adapter_target=_QWEN_ADAPTER_TARGET, model_id="mock/model", + max_audio_sec_per_actor=2400.0, adapter_kwargs={"revision": "abc123"}, ) stage.setup_on_node() @@ -537,6 +918,7 @@ def test_setup_on_node_downloads_faster_whisper_weights(mock_download: MagicMock stage = ASRStage( adapter_target=_FASTER_WHISPER_ADAPTER_TARGET, model_id="large-v3", + max_audio_sec_per_actor=2400.0, adapter_kwargs={"revision": "abc123"}, ) stage.setup_on_node() @@ -548,7 +930,11 @@ def test_setup_on_node_downloads_faster_whisper_weights(mock_download: MagicMock side_effect=RuntimeError("missing auth"), ) def test_setup_on_node_raises_by_default(mock_download: MagicMock) -> None: - stage = ASRStage(adapter_target=_QWEN_ADAPTER_TARGET, model_id="mock/model") + stage = ASRStage( + adapter_target=_QWEN_ADAPTER_TARGET, + model_id="mock/model", + max_audio_sec_per_actor=2400.0, + ) with pytest.raises(RuntimeError, match="download_weights_on_node failed"): stage.setup_on_node() mock_download.assert_called_once_with("mock/model") @@ -562,6 +948,7 @@ def test_setup_on_node_can_warn_and_retry_later(mock_download: MagicMock) -> Non stage = ASRStage( adapter_target=_QWEN_ADAPTER_TARGET, model_id="mock/model", + max_audio_sec_per_actor=2400.0, prefetch_fail_on_error=False, ) stage.setup_on_node() @@ -570,12 +957,17 @@ def test_setup_on_node_can_warn_and_retry_later(mock_download: MagicMock) -> Non def test_adapter_target_required() -> None: with pytest.raises(TypeError): - ASRStage(model_id="mock/model") + ASRStage(model_id="mock/model", max_audio_sec_per_actor=2400.0) def test_model_id_required() -> None: with pytest.raises(TypeError): - ASRStage(adapter_target=_QWEN_ADAPTER_TARGET) + ASRStage(adapter_target=_QWEN_ADAPTER_TARGET, max_audio_sec_per_actor=2400.0) + + +def test_max_audio_sec_per_actor_required() -> None: + with pytest.raises(TypeError): + ASRStage(adapter_target=_QWEN_ADAPTER_TARGET, model_id="mock/model") def test_stage_rejects_model_specific_revision_field() -> None: @@ -583,6 +975,7 @@ def test_stage_rejects_model_specific_revision_field() -> None: ASRStage( adapter_target=_QWEN_ADAPTER_TARGET, model_id="mock/model", + max_audio_sec_per_actor=2400.0, revision="abc123", # type: ignore[call-arg] ) @@ -593,6 +986,7 @@ def test_setup_uses_adapter_target_and_kwargs() -> None: stage = ASRStage( adapter_target=_QWEN_ADAPTER_TARGET, model_id="mock/model", + max_audio_sec_per_actor=2400.0, adapter_kwargs={ "revision": "abc123", "vllm_kwargs": { @@ -632,6 +1026,7 @@ def test_setup_derives_adapter_gpu_count_from_stage_resources( stage = ASRStage( adapter_target="tests.fake.Adapter", model_id="mock/model", + max_audio_sec_per_actor=2400.0, resources=Resources(gpus=requested_gpus), ) fake_adapter = MagicMock() @@ -647,6 +1042,7 @@ def test_setup_rejects_invalid_stage_gpu_resource(requested_gpus: float) -> None stage = ASRStage( adapter_target="tests.fake.Adapter", model_id="mock/model", + max_audio_sec_per_actor=2400.0, resources=Resources(gpus=requested_gpus), ) fake_adapter = MagicMock() @@ -661,7 +1057,11 @@ def test_setup_rejects_invalid_stage_gpu_resource(requested_gpus: float) -> None def test_setup_failure_cleans_partial_adapter_and_allows_retry() -> None: - stage = ASRStage(adapter_target=_QWEN_ADAPTER_TARGET, model_id="mock/model") + stage = ASRStage( + adapter_target=_QWEN_ADAPTER_TARGET, + model_id="mock/model", + max_audio_sec_per_actor=2400.0, + ) failed_adapter = MagicMock() failed_adapter.load_model.side_effect = RuntimeError("engine init failed") working_adapter = MagicMock() diff --git a/tests/stages/audio/inference/test_base.py b/tests/stages/audio/inference/test_base.py index 73a8a07ce4..956a7a9587 100644 --- a/tests/stages/audio/inference/test_base.py +++ b/tests/stages/audio/inference/test_base.py @@ -39,7 +39,11 @@ def test_common_adapter_infrastructure_is_not_reimplemented() -> None: def test_worker_sizing_uses_the_processing_stage_override() -> None: sed = SEDInferenceStage(adapter_target="package.Adapter", checkpoint_path="/checkpoint.pth") - asr = ASRStage(adapter_target="package.Adapter", model_id="model") + asr = ASRStage( + adapter_target="package.Adapter", + model_id="model", + max_audio_sec_per_actor=2400.0, + ) assert "num_workers_override" not in SEDInferenceStage.__dataclass_fields__ assert sed.num_workers() is None diff --git a/tests/stages/audio/test_model_input_segmentation.py b/tests/stages/audio/test_model_input_segmentation.py new file mode 100644 index 0000000000..929c3b330e --- /dev/null +++ b/tests/stages/audio/test_model_input_segmentation.py @@ -0,0 +1,99 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from nemo_curator.stages.audio.model_input_segmentation import ( + plan_audio_segments, + resolve_max_model_input_duration, +) + + +@pytest.mark.parametrize("max_duration_s", [0, -1, float("nan"), float("inf")]) +def test_resolve_max_model_input_duration_rejects_invalid_values(max_duration_s: float) -> None: + with pytest.raises(ValueError, match="max_inference_duration_s must be finite and > 0"): + resolve_max_model_input_duration(max_duration_s=max_duration_s, owner="test") + + +def test_plan_audio_segments_rejects_invalid_sample_rate() -> None: + with pytest.raises(ValueError, match="sample_rate must be > 0"): + plan_audio_segments(num_samples=100, sample_rate=0, max_duration_s=10.0, owner="test") + + +def test_plan_audio_segments_rejects_sub_sample_max_duration() -> None: + with pytest.raises(ValueError, match="must cover at least one sample"): + plan_audio_segments(num_samples=100, sample_rate=16000, max_duration_s=1e-6, owner="test") + + +def test_plan_audio_segments_keeps_zero_sample_inputs_representable() -> None: + segments = plan_audio_segments(num_samples=0, sample_rate=16000, max_duration_s=30.0, owner="test") + + assert len(segments) == 1 + assert segments[0].start_sample == 0 + assert segments[0].stop_sample == 0 + assert segments[0].duration_s == 0.0 + + +def test_plan_audio_segments_handles_non_divisible_final_segment() -> None: + segments = plan_audio_segments(num_samples=95, sample_rate=10, max_duration_s=3.0, owner="test") + + assert [(segment.start_sample, segment.stop_sample) for segment in segments] == [ + (0, 30), + (30, 60), + (60, 90), + (90, 95), + ] + assert [segment.duration_s for segment in segments] == [3.0, 3.0, 3.0, 0.5] + assert sum(segment.stop_sample - segment.start_sample for segment in segments) == 95 + + +def test_plan_audio_segments_exact_boundary_has_no_empty_tail() -> None: + segments = plan_audio_segments(num_samples=60, sample_rate=10, max_duration_s=3.0, owner="test") + + assert [(segment.start_sample, segment.stop_sample) for segment in segments] == [ + (0, 30), + (30, 60), + ] + assert [segment.duration_s for segment in segments] == [3.0, 3.0] + + +def test_plan_audio_segments_default_2400s_boundary_at_16khz() -> None: + sample_rate = 16000 + max_duration_s = 2400.0 + boundary_samples = int(sample_rate * max_duration_s) + + exact = plan_audio_segments( + num_samples=boundary_samples, + sample_rate=sample_rate, + max_duration_s=max_duration_s, + owner="ASRStage", + ) + just_over = plan_audio_segments( + num_samples=boundary_samples + 1, + sample_rate=sample_rate, + max_duration_s=max_duration_s, + owner="ASRStage", + ) + + assert [(segment.start_sample, segment.stop_sample) for segment in exact] == [ + (0, boundary_samples), + ] + assert exact[0].duration_s == max_duration_s + assert [(segment.start_sample, segment.stop_sample) for segment in just_over] == [ + (0, boundary_samples), + (boundary_samples, boundary_samples + 1), + ] + assert just_over[0].duration_s == max_duration_s + assert just_over[1].duration_s == 1.0 / sample_rate + assert sum(segment.stop_sample - segment.start_sample for segment in just_over) == boundary_samples + 1 diff --git a/tutorials/audio/faster_whisper/README.md b/tutorials/audio/faster_whisper/README.md index a0f2d750b0..48a0272e5c 100644 --- a/tutorials/audio/faster_whisper/README.md +++ b/tutorials/audio/faster_whisper/README.md @@ -59,6 +59,9 @@ revision for node prefetch and worker model loading. |---|---:| | Model | `large-v3` | | ASR stage batch size | `128` | +| Maximum padded audio per adapter call | `2400` seconds | +| Maximum duration per stage segment | `2400` seconds | +| Local duration bucketing | Enabled | | GPUs per ASR actor | `1` | | GPU compute type | `float16` | | CPU compute type | `int8` | @@ -68,12 +71,18 @@ revision for node prefetch and worker model loading. | Prediction field | `asr_prediction` | | Adapter extras field | `asr_extras` | -The stage batch size controls how many tasks Curator sends to one -`transcribe_batch()` call. `FasterWhisperASR` then calls -`WhisperModel.transcribe()` once per eligible audio, in order. This is -sequential per-audio inference; it does not use Faster-Whisper's -`BatchedInferencePipeline` and does not turn 128 files into one native model -batch. +The stage batch size controls the candidate-row window supplied to one +`process_batch()` call. Inside that window, `ASRStage` always segments audio at +`max_inference_duration_s`, then plans one or more `transcribe_batch()` calls +whose padded-audio cost does not exceed `max_audio_sec_per_actor`. With +`local_bucketing=true`, segments are ordered by duration while those call +boundaries are optimized; output is restored to parent-row order afterward. + +`FasterWhisperASR` still calls `WhisperModel.transcribe()` once per eligible +segment. This adapter is sequential per audio; it does not use +Faster-Whisper's `BatchedInferencePipeline`. The shared batching settings make +the stage contract explicit and bound each adapter call, but do not turn the +128-row candidate window into one native Faster-Whisper model batch. ## Select GPU or CPU execution @@ -107,9 +116,12 @@ language codes. It also accepts the input aliases normalized by the adapter: `nb` to `no`. Missing or unsupported languages are filtered by `ASRStage` before the adapter is called. -`ASRStage` writes transcription text to `asr_prediction`. It writes the forced, -normalized language code under `asr_extras.language_code`; this value is the -requested inference language, not a detected language. +`ASRStage` writes transcription text to `asr_prediction`. For a single stage +segment, it writes the forced, normalized language code under +`asr_extras.language_code`; this value is the requested inference language, +not a detected language. For a segmented input, adapter extras are merged in +temporal order without changing that flat schema; every segment uses the same +forced language code. The adapter intentionally discards Faster-Whisper's `TranscriptionInfo` and does not emit detected-language confidence, duration metadata, segment diff --git a/tutorials/audio/faster_whisper/pipeline.yaml b/tutorials/audio/faster_whisper/pipeline.yaml index 03238c02e2..387ae40f5b 100644 --- a/tutorials/audio/faster_whisper/pipeline.yaml +++ b/tutorials/audio/faster_whisper/pipeline.yaml @@ -136,6 +136,9 @@ supported_language_codes: pred_text_key: asr_prediction extras_key: asr_extras gpus_per_actor: 1 +max_audio_sec_per_actor: 2400.0 +max_inference_duration_s: 2400.0 +local_bucketing: true backend: ray_data execution_mode: streaming @@ -160,6 +163,9 @@ stages: pred_text_key: ${pred_text_key} extras_key: ${extras_key} batch_size: 128 + max_audio_sec_per_actor: ${max_audio_sec_per_actor} + max_inference_duration_s: ${max_inference_duration_s} + local_bucketing: ${local_bucketing} resources: _target_: nemo_curator.stages.resources.Resources gpus: ${gpus_per_actor} diff --git a/tutorials/audio/fleurs/README.md b/tutorials/audio/fleurs/README.md index 9219c0d04a..7a937d58f7 100644 --- a/tutorials/audio/fleurs/README.md +++ b/tutorials/audio/fleurs/README.md @@ -214,6 +214,9 @@ pipeline = Pipeline( ASRStage( adapter_target="nemo_curator.models.asr.nemo_asr.NeMoASRAdapter", model_id="nvidia/parakeet-tdt-0.6b-v2", + max_audio_sec_per_actor=240, + max_inference_duration_s=120, + local_bucketing=True, audio_filepath_key="audio_filepath", ), GetPairwiseWerStage(text_key="text", pred_text_key="pred_text", wer_key="wer_pct"), @@ -233,7 +236,7 @@ finally: | Problem | Cause | Fix | |---|---|---| | Output directory already exists | Previous run left `${raw_data_dir}/result/${lang}/` | Remove the directory before re-running | -| OOM during ASR inference | GPU VRAM too small for model + batch | Reduce `stages.1.batch_size` or use a smaller model | +| OOM during ASR inference | GPU VRAM too small for padded adapter work | Reduce `max_audio_sec_per_actor` (and, if needed, `max_inference_duration_s`) or use a smaller model | | `CUDA error: invalid argument` in RNNT label-loop decoding | NeMo CUDA-graph decoder is unsupported by the local CUDA runtime/driver combination | Set `stages.1.adapter_kwargs.use_cuda_graph_decoder=false` (the supplied FLEURS config already does this) | | CPU inference very slow | CPU is 10–50x slower than GPU | Set `stages.1.resources.gpus=1`; CPU is only for testing | | Empty output JSONL | `wer_threshold` too strict for the model+language pair | Increase `wer_threshold` or use a better-matching ASR model | diff --git a/tutorials/audio/fleurs/fleurs_tutorial.ipynb b/tutorials/audio/fleurs/fleurs_tutorial.ipynb index 20939158d6..1de3951fc7 100644 --- a/tutorials/audio/fleurs/fleurs_tutorial.ipynb +++ b/tutorials/audio/fleurs/fleurs_tutorial.ipynb @@ -13,7 +13,66 @@ }, "outputs": [], "source": [ - "# Silence noisy logs BEFORE importing Curator / Ray / NeMo.\nimport contextlib\nimport logging\nimport os\nimport sys\nimport warnings\nfrom collections.abc import Iterator\n\n# Ensure the active venv's `ray` CLI is on PATH (needed for RayClient.start()).\n_venv_bin = os.path.dirname(sys.executable)\nif os.path.isdir(_venv_bin):\n os.environ[\"PATH\"] = _venv_bin + os.pathsep + os.environ.get(\"PATH\", \"\")\n\n# Curator uses loguru; NeMo uses Python logging; Ray forwards worker stdout to the driver.\nos.environ[\"LOGURU_LEVEL\"] = \"ERROR\"\nos.environ[\"NEMO_LOG_LEVEL\"] = \"ERROR\"\nos.environ.setdefault(\"RAY_BACKEND_LOG_LEVEL\", \"error\")\nos.environ.setdefault(\"RAY_DEDUP_LOGS\", \"1\")\n\nwarnings.filterwarnings(\"ignore\", message=r\".*concurrency.*deprecated.*\")\nwarnings.filterwarnings(\"ignore\", message=r\".*num_cpus and num_gpus.*experimental.*\")\n\n\ndef _silence_third_party_loggers() -> None:\n logging.getLogger(\"nemo_logger\").setLevel(logging.ERROR)\n for name in (\"ray\", \"ray.data\", \"ray.worker\", \"ray.util\", \"urllib3\", \"matplotlib\", \"filelock\"):\n logging.getLogger(name).setLevel(logging.ERROR)\n\n\ndef configure_loguru() -> None:\n from loguru import logger\n\n logger.remove()\n logger.add(sys.stderr, level=\"ERROR\")\n\n\ndef configure_quiet_ray() -> None:\n \"\"\"Disable Ray Data progress bars without pre-initing Ray on the driver.\n\n Do not call ray.init() here: a driver-side init (especially with\n log_to_driver=False) breaks cosmos_xenna GPU resource discovery even when\n the cluster exposes GPUs. XennaExecutor / RayDataExecutor connect themselves.\n \"\"\"\n from ray.data import DataContext\n\n ctx = DataContext.get_current()\n ctx.enable_progress_bars = False\n ctx.enable_operator_progress_bars = False\n ctx.use_ray_tqdm = False\n\n _silence_third_party_loggers()\n\n\n@contextlib.contextmanager\ndef quiet_pipeline_run() -> Iterator[None]:\n \"\"\"Suppress Ray / NeMo / Xenna driver noise during pipeline.run().\"\"\"\n configure_loguru()\n _silence_third_party_loggers()\n with open(os.devnull, \"w\") as devnull, contextlib.redirect_stdout(devnull), contextlib.redirect_stderr(devnull):\n yield\n\n" + "# Silence noisy logs BEFORE importing Curator / Ray / NeMo.\n", + "import contextlib\n", + "import logging\n", + "import os\n", + "import sys\n", + "import warnings\n", + "from collections.abc import Iterator\n", + "\n", + "# Ensure the active venv's `ray` CLI is on PATH (needed for RayClient.start()).\n", + "_venv_bin = os.path.dirname(sys.executable)\n", + "if os.path.isdir(_venv_bin):\n", + " os.environ[\"PATH\"] = _venv_bin + os.pathsep + os.environ.get(\"PATH\", \"\")\n", + "\n", + "# Curator uses loguru; NeMo uses Python logging; Ray forwards worker stdout to the driver.\n", + "os.environ[\"LOGURU_LEVEL\"] = \"ERROR\"\n", + "os.environ[\"NEMO_LOG_LEVEL\"] = \"ERROR\"\n", + "os.environ.setdefault(\"RAY_BACKEND_LOG_LEVEL\", \"error\")\n", + "os.environ.setdefault(\"RAY_DEDUP_LOGS\", \"1\")\n", + "\n", + "warnings.filterwarnings(\"ignore\", message=r\".*concurrency.*deprecated.*\")\n", + "warnings.filterwarnings(\"ignore\", message=r\".*num_cpus and num_gpus.*experimental.*\")\n", + "\n", + "\n", + "def _silence_third_party_loggers() -> None:\n", + " logging.getLogger(\"nemo_logger\").setLevel(logging.ERROR)\n", + " for name in (\"ray\", \"ray.data\", \"ray.worker\", \"ray.util\", \"urllib3\", \"matplotlib\", \"filelock\"):\n", + " logging.getLogger(name).setLevel(logging.ERROR)\n", + "\n", + "\n", + "def configure_loguru() -> None:\n", + " from loguru import logger\n", + "\n", + " logger.remove()\n", + " logger.add(sys.stderr, level=\"ERROR\")\n", + "\n", + "\n", + "def configure_quiet_ray() -> None:\n", + " \"\"\"Disable Ray Data progress bars without pre-initing Ray on the driver.\n", + "\n", + " Do not call ray.init() here: a driver-side init (especially with\n", + " log_to_driver=False) breaks cosmos_xenna GPU resource discovery even when\n", + " the cluster exposes GPUs. XennaExecutor / RayDataExecutor connect themselves.\n", + " \"\"\"\n", + " from ray.data import DataContext\n", + "\n", + " ctx = DataContext.get_current()\n", + " ctx.enable_progress_bars = False\n", + " ctx.enable_operator_progress_bars = False\n", + " ctx.use_ray_tqdm = False\n", + "\n", + " _silence_third_party_loggers()\n", + "\n", + "\n", + "@contextlib.contextmanager\n", + "def quiet_pipeline_run() -> Iterator[None]:\n", + " \"\"\"Suppress Ray / NeMo / Xenna driver noise during pipeline.run().\"\"\"\n", + " configure_loguru()\n", + " _silence_third_party_loggers()\n", + " with open(os.devnull, \"w\") as devnull, contextlib.redirect_stdout(devnull), contextlib.redirect_stderr(devnull):\n", + " yield" ] }, { @@ -65,7 +124,23 @@ } ], "source": [ - "import json\nimport os\nimport shutil\n\nfrom nemo_curator.backends.xenna import XennaExecutor\nfrom nemo_curator.core.client import RayClient\nfrom nemo_curator.pipeline import Pipeline\nfrom nemo_curator.stages.audio.common import GetAudioDurationStage, PreserveByValueStage\nfrom nemo_curator.stages.audio.datasets.fleurs.create_initial_manifest import CreateInitialManifestFleursStage\nfrom nemo_curator.stages.audio.inference.asr.stage import ASRStage\nfrom nemo_curator.stages.audio.io.convert import AudioToDocumentStage\nfrom nemo_curator.stages.audio.metrics.wer import GetPairwiseWerStage\nfrom nemo_curator.stages.resources import Resources\nfrom nemo_curator.stages.text.io.writer import JsonlWriter\n\nconfigure_loguru()\n_silence_third_party_loggers()" + "import json\n", + "import os\n", + "import shutil\n", + "\n", + "from nemo_curator.backends.xenna import XennaExecutor\n", + "from nemo_curator.core.client import RayClient\n", + "from nemo_curator.pipeline import Pipeline\n", + "from nemo_curator.stages.audio.common import GetAudioDurationStage, PreserveByValueStage\n", + "from nemo_curator.stages.audio.datasets.fleurs.create_initial_manifest import CreateInitialManifestFleursStage\n", + "from nemo_curator.stages.audio.inference.asr.stage import ASRStage\n", + "from nemo_curator.stages.audio.io.convert import AudioToDocumentStage\n", + "from nemo_curator.stages.audio.metrics.wer import GetPairwiseWerStage\n", + "from nemo_curator.stages.resources import Resources\n", + "from nemo_curator.stages.text.io.writer import JsonlWriter\n", + "\n", + "configure_loguru()\n", + "_silence_third_party_loggers()" ] }, { @@ -88,14 +163,26 @@ }, "outputs": [], "source": [ - "RAW_DATA_DIR = os.path.abspath(\"./example_audio/fleurs\")\nLANG = \"hy_am\"\nSPLIT = \"dev\" # matches audio_fleurs_benchmark.py default (nightly CI uses train)\nMODEL_NAME = \"nvidia/stt_hy_fastconformer_hybrid_large_pc\"\n# Keep GPU inference enabled while avoiding NeMo's RNNT label-loop CUDA graph,\n# which is unsupported by some CUDA runtime/driver combinations.\nUSE_CUDA_GRAPH_DECODER = False\nWER_THRESHOLD = 5.5\nGPUS = 1.0\n\nRESULT_DIR = os.path.join(RAW_DATA_DIR, \"result\", LANG)\nif os.path.isdir(RESULT_DIR):\n shutil.rmtree(RESULT_DIR)" + "RAW_DATA_DIR = os.path.abspath(\"./example_audio/fleurs\")\n", + "LANG = \"hy_am\"\n", + "SPLIT = \"dev\" # matches audio_fleurs_benchmark.py default (nightly CI uses train)\n", + "MODEL_NAME = \"nvidia/stt_hy_fastconformer_hybrid_large_pc\"\n", + "# Keep GPU inference enabled while avoiding NeMo's RNNT label-loop CUDA graph,\n", + "# which is unsupported by some CUDA runtime/driver combinations.\n", + "USE_CUDA_GRAPH_DECODER = False\n", + "WER_THRESHOLD = 5.5\n", + "GPUS = 1.0\n", + "\n", + "RESULT_DIR = os.path.join(RAW_DATA_DIR, \"result\", LANG)\n", + "if os.path.isdir(RESULT_DIR):\n", + " shutil.rmtree(RESULT_DIR)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "## Step 1: Build the pipeline\n\nThe pipeline has 7 stages: download \u2192 ASR \u2192 WER \u2192 duration \u2192 filter \u2192 convert \u2192 write.\nWe define a function so we can rebuild the pipeline for each backend run." + "## Step 1: Build the pipeline\n\nThe pipeline has 7 stages: download → ASR → WER → duration → filter → convert → write.\nWe define a function so we can rebuild the pipeline for each backend run." ] }, { @@ -119,7 +206,35 @@ } ], "source": [ - "def build_pipeline(result_dir: str) -> Pipeline:\n \"\"\"Create a fresh pipeline writing to *result_dir*.\"\"\"\n if os.path.isdir(result_dir):\n shutil.rmtree(result_dir)\n p = Pipeline(name=\"fleurs_tutorial\", description=\"Download FLEURS, run ASR, filter by WER\")\n p.add_stage(\n CreateInitialManifestFleursStage(lang=LANG, split=SPLIT, raw_data_dir=RAW_DATA_DIR).with_(batch_size=4)\n )\n p.add_stage(\n ASRStage(\n adapter_target=\"nemo_curator.models.asr.nemo_asr.NeMoASRAdapter\",\n model_id=MODEL_NAME,\n audio_filepath_key=\"audio_filepath\",\n adapter_kwargs={\"use_cuda_graph_decoder\": USE_CUDA_GRAPH_DECODER},\n batch_size=16,\n ).with_(resources=Resources(gpus=GPUS))\n )\n p.add_stage(GetPairwiseWerStage(text_key=\"text\", pred_text_key=\"pred_text\", wer_key=\"wer_pct\"))\n p.add_stage(GetAudioDurationStage(audio_filepath_key=\"audio_filepath\", duration_key=\"duration\"))\n p.add_stage(PreserveByValueStage(input_value_key=\"wer_pct\", target_value=WER_THRESHOLD, operator=\"le\"))\n p.add_stage(AudioToDocumentStage().with_(batch_size=1))\n p.add_stage(JsonlWriter(path=result_dir, write_kwargs={\"force_ascii\": False}))\n return p\n\n\nprint(build_pipeline(RESULT_DIR).describe())\n" + "def build_pipeline(result_dir: str) -> Pipeline:\n", + " \"\"\"Create a fresh pipeline writing to *result_dir*.\"\"\"\n", + " if os.path.isdir(result_dir):\n", + " shutil.rmtree(result_dir)\n", + " p = Pipeline(name=\"fleurs_tutorial\", description=\"Download FLEURS, run ASR, filter by WER\")\n", + " p.add_stage(\n", + " CreateInitialManifestFleursStage(lang=LANG, split=SPLIT, raw_data_dir=RAW_DATA_DIR).with_(batch_size=4)\n", + " )\n", + " p.add_stage(\n", + " ASRStage(\n", + " adapter_target=\"nemo_curator.models.asr.nemo_asr.NeMoASRAdapter\",\n", + " model_id=MODEL_NAME,\n", + " max_audio_sec_per_actor=240.0,\n", + " max_inference_duration_s=120.0,\n", + " local_bucketing=True,\n", + " audio_filepath_key=\"audio_filepath\",\n", + " adapter_kwargs={\"use_cuda_graph_decoder\": USE_CUDA_GRAPH_DECODER},\n", + " batch_size=16,\n", + " ).with_(resources=Resources(gpus=GPUS))\n", + " )\n", + " p.add_stage(GetPairwiseWerStage(text_key=\"text\", pred_text_key=\"pred_text\", wer_key=\"wer_pct\"))\n", + " p.add_stage(GetAudioDurationStage(audio_filepath_key=\"audio_filepath\", duration_key=\"duration\"))\n", + " p.add_stage(PreserveByValueStage(input_value_key=\"wer_pct\", target_value=WER_THRESHOLD, operator=\"le\"))\n", + " p.add_stage(AudioToDocumentStage().with_(batch_size=1))\n", + " p.add_stage(JsonlWriter(path=result_dir, write_kwargs={\"force_ascii\": False}))\n", + " return p\n", + "\n", + "\n", + "print(build_pipeline(RESULT_DIR).describe())" ] }, { @@ -222,7 +337,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "[xenna] 44.20s \u2014 50 samples, mean WER 2.6%\n" + "[xenna] 44.20s — 50 samples, mean WER 2.6%\n" ] }, { @@ -236,12 +351,61 @@ "name": "stdout", "output_type": "stream", "text": [ - "[ray_data] 38.56s \u2014 50 samples, mean WER 2.6%\n" + "[ray_data] 38.56s — 50 samples, mean WER 2.6%\n" ] } ], "source": [ - "import time\n\nfrom nemo_curator.backends.ray_data import RayDataExecutor\n\n# Avoid attaching to a stale Ray cluster left over from a prior notebook run.\nos.environ.pop(\"RAY_ADDRESS\", None)\n\nray_client = RayClient(num_gpus=int(GPUS))\nray_client.start()\nconfigure_quiet_ray()\n\n\ndef load_results(result_dir: str) -> list[dict]:\n \"\"\"Read all JSONL files from a result directory.\"\"\"\n data = []\n for fname in os.listdir(result_dir):\n if fname.endswith(\".jsonl\"):\n with open(os.path.join(result_dir, fname)) as f:\n data.extend(json.loads(line) for line in f if line.strip())\n return data\n\n\nbackends = {\n \"xenna\": XennaExecutor,\n \"ray_data\": RayDataExecutor,\n}\n\nrun_results = {}\n\nfor name, executor_cls in backends.items():\n result_dir = os.path.join(RAW_DATA_DIR, f\"result_{name}\")\n pipeline = build_pipeline(result_dir)\n executor = executor_cls()\n\n t0 = time.time()\n with quiet_pipeline_run():\n pipeline.run(executor)\n elapsed = time.time() - t0\n\n data = load_results(result_dir)\n wers = [r.get(\"wer_pct\", 0) for r in data]\n\n run_results[name] = {\n \"time\": elapsed,\n \"samples\": len(data),\n \"mean_wer\": sum(wers) / len(wers) if wers else 0,\n \"total_dur\": sum(r.get(\"duration\", 0) for r in data),\n \"data\": data,\n }\n print(f\"[{name}] {elapsed:.2f}s \u2014 {len(data)} samples, mean WER {run_results[name]['mean_wer']:.1f}%\")" + "import time\n", + "\n", + "from nemo_curator.backends.ray_data import RayDataExecutor\n", + "\n", + "# Avoid attaching to a stale Ray cluster left over from a prior notebook run.\n", + "os.environ.pop(\"RAY_ADDRESS\", None)\n", + "\n", + "ray_client = RayClient(num_gpus=int(GPUS))\n", + "ray_client.start()\n", + "configure_quiet_ray()\n", + "\n", + "\n", + "def load_results(result_dir: str) -> list[dict]:\n", + " \"\"\"Read all JSONL files from a result directory.\"\"\"\n", + " data = []\n", + " for fname in os.listdir(result_dir):\n", + " if fname.endswith(\".jsonl\"):\n", + " with open(os.path.join(result_dir, fname)) as f:\n", + " data.extend(json.loads(line) for line in f if line.strip())\n", + " return data\n", + "\n", + "\n", + "backends = {\n", + " \"xenna\": XennaExecutor,\n", + " \"ray_data\": RayDataExecutor,\n", + "}\n", + "\n", + "run_results = {}\n", + "\n", + "for name, executor_cls in backends.items():\n", + " result_dir = os.path.join(RAW_DATA_DIR, f\"result_{name}\")\n", + " pipeline = build_pipeline(result_dir)\n", + " executor = executor_cls()\n", + "\n", + " t0 = time.time()\n", + " with quiet_pipeline_run():\n", + " pipeline.run(executor)\n", + " elapsed = time.time() - t0\n", + "\n", + " data = load_results(result_dir)\n", + " wers = [r.get(\"wer_pct\", 0) for r in data]\n", + "\n", + " run_results[name] = {\n", + " \"time\": elapsed,\n", + " \"samples\": len(data),\n", + " \"mean_wer\": sum(wers) / len(wers) if wers else 0,\n", + " \"total_dur\": sum(r.get(\"duration\", 0) for r in data),\n", + " \"data\": data,\n", + " }\n", + " print(f\"[{name}] {elapsed:.2f}s — {len(data)} samples, mean WER {run_results[name]['mean_wer']:.1f}%\")" ] }, { @@ -260,12 +424,35 @@ "name": "stdout", "output_type": "stream", "text": [ - "\n============================================================\nBackend Comparison\n============================================================\n Xenna Ray Data Match\n Time (s) 44.20 38.56 \n Samples 50 50 \u2713\n Mean WER 2.6 2.6 \u2713\n Audio (s) 554.0 554.0 \u2713\n\n\u2192 ray_data was 1.1x faster on this dataset\n" + "\n============================================================\nBackend Comparison\n============================================================\n Xenna Ray Data Match\n Time (s) 44.20 38.56 \n Samples 50 50 ✓\n Mean WER 2.6 2.6 ✓\n Audio (s) 554.0 554.0 ✓\n\n→ ray_data was 1.1x faster on this dataset\n" ] } ], "source": [ - "MATCH_TOL = 0.1\n\nprint(\"\\n\" + \"=\" * 60)\nprint(\"Backend Comparison\")\nprint(\"=\" * 60)\nprint(f\"{'':>12s} {'Xenna':>10s} {'Ray Data':>10s} {'Match':>6s}\")\nprint(f\"{'Time (s)':>12s} {run_results['xenna']['time']:10.2f} {run_results['ray_data']['time']:10.2f} {'':>6s}\")\nprint(\n f\"{'Samples':>12s} {run_results['xenna']['samples']:10d} {run_results['ray_data']['samples']:10d}\"\n f\" {'\u2713' if run_results['xenna']['samples'] == run_results['ray_data']['samples'] else '\u2717':>6s}\"\n)\nprint(\n f\"{'Mean WER':>12s} {run_results['xenna']['mean_wer']:10.1f} {run_results['ray_data']['mean_wer']:10.1f}\"\n f\" {'\u2713' if abs(run_results['xenna']['mean_wer'] - run_results['ray_data']['mean_wer']) < MATCH_TOL else '\u2717':>6s}\"\n)\nprint(\n f\"{'Audio (s)':>12s} {run_results['xenna']['total_dur']:10.1f} {run_results['ray_data']['total_dur']:10.1f}\"\n f\" {'\u2713' if abs(run_results['xenna']['total_dur'] - run_results['ray_data']['total_dur']) < MATCH_TOL else '\u2717':>6s}\"\n)\nspeedup = run_results[\"ray_data\"][\"time\"] / run_results[\"xenna\"][\"time\"]\nfaster = \"xenna\" if speedup > 1 else \"ray_data\"\nprint(f\"\\n\u2192 {faster} was {max(speedup, 1 / speedup):.1f}x faster on this dataset\")\n\nresults = run_results[\"xenna\"][\"data\"]" + "MATCH_TOL = 0.1\n", + "\n", + "print(\"\\n\" + \"=\" * 60)\n", + "print(\"Backend Comparison\")\n", + "print(\"=\" * 60)\n", + "print(f\"{'':>12s} {'Xenna':>10s} {'Ray Data':>10s} {'Match':>6s}\")\n", + "print(f\"{'Time (s)':>12s} {run_results['xenna']['time']:10.2f} {run_results['ray_data']['time']:10.2f} {'':>6s}\")\n", + "print(\n", + " f\"{'Samples':>12s} {run_results['xenna']['samples']:10d} {run_results['ray_data']['samples']:10d}\"\n", + " f\" {'✓' if run_results['xenna']['samples'] == run_results['ray_data']['samples'] else '✗':>6s}\"\n", + ")\n", + "print(\n", + " f\"{'Mean WER':>12s} {run_results['xenna']['mean_wer']:10.1f} {run_results['ray_data']['mean_wer']:10.1f}\"\n", + " f\" {'✓' if abs(run_results['xenna']['mean_wer'] - run_results['ray_data']['mean_wer']) < MATCH_TOL else '✗':>6s}\"\n", + ")\n", + "print(\n", + " f\"{'Audio (s)':>12s} {run_results['xenna']['total_dur']:10.1f} {run_results['ray_data']['total_dur']:10.1f}\"\n", + " f\" {'✓' if abs(run_results['xenna']['total_dur'] - run_results['ray_data']['total_dur']) < MATCH_TOL else '✗':>6s}\"\n", + ")\n", + "speedup = run_results[\"ray_data\"][\"time\"] / run_results[\"xenna\"][\"time\"]\n", + "faster = \"xenna\" if speedup > 1 else \"ray_data\"\n", + "print(f\"\\n→ {faster} was {max(speedup, 1 / speedup):.1f}x faster on this dataset\")\n", + "\n", + "results = run_results[\"xenna\"][\"data\"]" ] }, { @@ -291,12 +478,14 @@ "name": "stdout", "output_type": "stream", "text": [ - "Total samples after filtering: 50\n\nSample entry:\n{\n \"audio_filepath\": \"/home/aaftabv/prs-into-curator/CuratorPR2254FastConformer/tutorials/audio/fleurs/example_audio/fleurs/hy_am/dev/11348083370274933042.wav\",\n \"text\": \"\u053b\u057d\u0580\u0561\u0575\u0565\u056c\u0568 \u057a\u0561\u0570\u0561\u0576\u057b\u0578\u0582\u0574 \u0567 \u0577\u0561\u0580\u0578\u0582\u0576\u0561\u056f\u0561\u056f\u0561\u0576 \u057c\u0561\u0566\u0574\u0561\u056f\u0561\u0576 \u0576\u0565\u0580\u056f\u0561\u0575\u0578\u0582\u0569\u0575\u0578\u0582\u0576 \u0570\u0578\u057e\u057f\u0578\u0582\u0574 \u057f\u0561\u057d\u0568 \u057f\u0561\u0580\u057e\u0561 \u0568\u0576\u0569\u0561\u0581\u0584\u0578\u0582\u0574 \u057a\u0561\u0575\u0574\u0561\u0576\u0561\u0563\u056b\u0580\u0568 \u056f\u0576\u0584\u0565\u056c\u0578\u0582\u0581 \u0570\u0565\u057f\u0578, \u0574\u056b\u0576\u0579\u0564\u0565\u057c \u054a\u0561\u0572\u0565\u057d\u057f\u056b\u0576\u056b \u056b\u0580\u0561\u057e\u0561\u057d\u0578\u0582 \u0574\u0561\u0580\u0574\u056b\u0576\u0576\u0565\u0580\u0568 \u0570\u0561\u0574\u0561\u0571\u0561\u0575\u0576\u057e\u0578\u0582\u0574 \u0565\u0576 \u0569\u0578\u0572\u0576\u0565\u056c \u0561\u0575\u0564\u057a\u056b\u057d\u056b \u0576\u0565\u0580\u056f\u0561\u0575\u0578\u0582\u0569\u0575\u0561\u0576\u0568 \u0574\u056b\u0561\u0575\u0576 \u0570\u056b\u0576\u0563 \u057f\u0561\u0580\u0578\u057e:\",\n \"pred_text\": \"\u053b\u057d\u0580\u0561\u0575\u0565\u056c\u0568 \u057a\u0561\u0570\u0561\u0576\u057b\u0578\u0582\u0574 \u0567 \u0577\u0561\u0580\u0578\u0582\u0576\u0561\u056f\u0561\u056f\u0561\u0576 \u057c\u0561\u0566\u0574\u0561\u056f\u0561\u0576 \u0576\u0565\u0580\u056f\u0561\u0575\u0578\u0582\u0569\u0575\u0578\u0582\u0576 \u0570\u0578\u057e\u057f\u0578\u0582\u0574 \u057f\u0561\u057d\u0568 \u057f\u0561\u0580\u057e\u0561 \u0568\u0576\u0569\u0561\u0581\u0584\u0578\u0582\u0574 \u057a\u0561\u0575\u0574\u0561\u0576\u0561\u0563\u056b\u0580\u0568 \u056f\u0576\u0584\u0565\u056c\u0578\u0582\u0581 \u0570\u0565\u057f\u0578, \u0574\u056b\u0576\u0579\u0564\u0565\u057c \u054a\u0561\u0572\u0565\u057d\u057f\u056b\u0576\u056b \u056b\u0580\u0561\u057e\u0561\u057d\u0578\u0582 \u0574\u0561\u0580\u0574\u056b\u0576\u0576\u0565\u0580\u0568 \u0570\u0561\u0574\u0561\u0571\u0561\u0575\u0576\u057e\u0578\u0582\u0574 \u0565\u0576 \u0569\u0578\u0572\u0576\u0565\u056c \u0561\u0575\u0564\u057a\u056b\u057d\u056b \u0576\u0565\u0580\u056f\u0561\u0575\u0578\u0582\u0569\u0575\u0561\u0576\u0568 \u0574\u056b\u0561\u0575\u0576 \u0570\u056b\u0576\u0563 \u057f\u0561\u0580\u0578\u057e\u0589\",\n \"wer_pct\": 4.0,\n \"duration\": 19.2\n}\n" + "Total samples after filtering: 50\n\nSample entry:\n{\n \"audio_filepath\": \"/home/aaftabv/prs-into-curator/CuratorPR2254FastConformer/tutorials/audio/fleurs/example_audio/fleurs/hy_am/dev/11348083370274933042.wav\",\n \"text\": \"Իսրայելը պահանջում է շարունակական ռազմական ներկայություն հովտում տասը տարվա ընթացքում պայմանագիրը կնքելուց հետո, մինչդեռ Պաղեստինի իրավասու մարմինները համաձայնվում են թողնել այդպիսի ներկայությանը միայն հինգ տարով:\",\n \"pred_text\": \"Իսրայելը պահանջում է շարունակական ռազմական ներկայություն հովտում տասը տարվա ընթացքում պայմանագիրը կնքելուց հետո, մինչդեռ Պաղեստինի իրավասու մարմինները համաձայնվում են թողնել այդպիսի ներկայությանը միայն հինգ տարով։\",\n \"wer_pct\": 4.0,\n \"duration\": 19.2\n}\n" ] } ], "source": [ - "print(f\"Total samples after filtering: {len(results)}\")\nprint(\"\\nSample entry:\")\nprint(json.dumps(results[0], indent=2, ensure_ascii=False) if results else \"No results\")" + "print(f\"Total samples after filtering: {len(results)}\")\n", + "print(\"\\nSample entry:\")\n", + "print(json.dumps(results[0], indent=2, ensure_ascii=False) if results else \"No results\")" ] }, { @@ -332,12 +521,67 @@ "name": "stdout", "output_type": "stream", "text": [ - "\nWER \u2014 min: 0.0%, max: 5.3%, mean: 2.6%, median: 3.7%\nDuration \u2014 min: 3.66s, max: 19.20s, total: 554.0s\n" + "\nWER — min: 0.0%, max: 5.3%, mean: 2.6%, median: 3.7%\nDuration — min: 3.66s, max: 19.20s, total: 554.0s\n" ] } ], "source": [ - "import matplotlib.pyplot as plt\nimport numpy as np\n\nwers = [r.get(\"wer_pct\", 0) for r in results]\ndurations = [r.get(\"duration\", 0) for r in results]\n\nif not wers:\n print(\"No results to visualize. Try relaxing WER_THRESHOLD or re-running the pipeline.\")\nelse:\n fig, axes = plt.subplots(2, 2, figsize=(14, 10))\n\n # 1. WER histogram with threshold line\n ax = axes[0, 0]\n ax.hist(wers, bins=30, color=\"#4C72B0\", edgecolor=\"white\", alpha=0.85)\n ax.axvline(WER_THRESHOLD, color=\"#C44E52\", linestyle=\"--\", linewidth=2, label=f\"Threshold ({WER_THRESHOLD}%)\")\n ax.set_xlabel(\"WER (%)\")\n ax.set_ylabel(\"Count\")\n ax.set_title(\"WER Distribution\")\n ax.legend()\n\n # 2. Duration distribution\n ax = axes[0, 1]\n ax.hist(durations, bins=30, color=\"#55A868\", edgecolor=\"white\", alpha=0.85)\n ax.set_xlabel(\"Duration (seconds)\")\n ax.set_ylabel(\"Count\")\n ax.set_title(\"Audio Duration Distribution\")\n\n # 3. WER vs Duration scatter\n ax = axes[1, 0]\n scatter = ax.scatter(durations, wers, c=wers, cmap=\"RdYlGn_r\", alpha=0.6, s=20, edgecolors=\"none\")\n ax.axhline(WER_THRESHOLD, color=\"#C44E52\", linestyle=\"--\", linewidth=1.5, alpha=0.7)\n ax.set_xlabel(\"Duration (seconds)\")\n ax.set_ylabel(\"WER (%)\")\n ax.set_title(\"WER vs Duration\")\n plt.colorbar(scatter, ax=ax, label=\"WER %\")\n\n # 4. Pass rate at multiple thresholds\n ax = axes[1, 1]\n thresholds = [5, 10, 25, 50, 75, 100]\n pass_rates = [sum(1 for w in wers if w <= t) / len(wers) * 100 for t in thresholds]\n bars = ax.bar([str(t) for t in thresholds], pass_rates, color=\"#8172B2\", edgecolor=\"white\")\n ax.set_xlabel(\"WER Threshold (%)\")\n ax.set_ylabel(\"Samples Passing (%)\")\n ax.set_title(\"Dataset Yield by Threshold\")\n for bar, rate in zip(bars, pass_rates, strict=True):\n ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 1, f\"{rate:.0f}%\", ha=\"center\", fontsize=9)\n ax.set_ylim(0, 110)\n\n fig.suptitle(f\"FLEURS {LANG} / {SPLIT} \u2014 {len(results)} samples (WER \u2264 {WER_THRESHOLD}%)\", fontsize=13, y=1.01)\n fig.tight_layout()\n plt.show()\n\n print(\n f\"\\nWER \u2014 min: {min(wers):.1f}%, max: {max(wers):.1f}%, mean: {np.mean(wers):.1f}%, median: {np.median(wers):.1f}%\"\n )\n print(f\"Duration \u2014 min: {min(durations):.2f}s, max: {max(durations):.2f}s, total: {sum(durations):.1f}s\")" + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "wers = [r.get(\"wer_pct\", 0) for r in results]\n", + "durations = [r.get(\"duration\", 0) for r in results]\n", + "\n", + "if not wers:\n", + " print(\"No results to visualize. Try relaxing WER_THRESHOLD or re-running the pipeline.\")\n", + "else:\n", + " fig, axes = plt.subplots(2, 2, figsize=(14, 10))\n", + "\n", + " # 1. WER histogram with threshold line\n", + " ax = axes[0, 0]\n", + " ax.hist(wers, bins=30, color=\"#4C72B0\", edgecolor=\"white\", alpha=0.85)\n", + " ax.axvline(WER_THRESHOLD, color=\"#C44E52\", linestyle=\"--\", linewidth=2, label=f\"Threshold ({WER_THRESHOLD}%)\")\n", + " ax.set_xlabel(\"WER (%)\")\n", + " ax.set_ylabel(\"Count\")\n", + " ax.set_title(\"WER Distribution\")\n", + " ax.legend()\n", + "\n", + " # 2. Duration distribution\n", + " ax = axes[0, 1]\n", + " ax.hist(durations, bins=30, color=\"#55A868\", edgecolor=\"white\", alpha=0.85)\n", + " ax.set_xlabel(\"Duration (seconds)\")\n", + " ax.set_ylabel(\"Count\")\n", + " ax.set_title(\"Audio Duration Distribution\")\n", + "\n", + " # 3. WER vs Duration scatter\n", + " ax = axes[1, 0]\n", + " scatter = ax.scatter(durations, wers, c=wers, cmap=\"RdYlGn_r\", alpha=0.6, s=20, edgecolors=\"none\")\n", + " ax.axhline(WER_THRESHOLD, color=\"#C44E52\", linestyle=\"--\", linewidth=1.5, alpha=0.7)\n", + " ax.set_xlabel(\"Duration (seconds)\")\n", + " ax.set_ylabel(\"WER (%)\")\n", + " ax.set_title(\"WER vs Duration\")\n", + " plt.colorbar(scatter, ax=ax, label=\"WER %\")\n", + "\n", + " # 4. Pass rate at multiple thresholds\n", + " ax = axes[1, 1]\n", + " thresholds = [5, 10, 25, 50, 75, 100]\n", + " pass_rates = [sum(1 for w in wers if w <= t) / len(wers) * 100 for t in thresholds]\n", + " bars = ax.bar([str(t) for t in thresholds], pass_rates, color=\"#8172B2\", edgecolor=\"white\")\n", + " ax.set_xlabel(\"WER Threshold (%)\")\n", + " ax.set_ylabel(\"Samples Passing (%)\")\n", + " ax.set_title(\"Dataset Yield by Threshold\")\n", + " for bar, rate in zip(bars, pass_rates, strict=True):\n", + " ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 1, f\"{rate:.0f}%\", ha=\"center\", fontsize=9)\n", + " ax.set_ylim(0, 110)\n", + "\n", + " fig.suptitle(f\"FLEURS {LANG} / {SPLIT} — {len(results)} samples (WER ≤ {WER_THRESHOLD}%)\", fontsize=13, y=1.01)\n", + " fig.tight_layout()\n", + " plt.show()\n", + "\n", + " print(\n", + " f\"\\nWER — min: {min(wers):.1f}%, max: {max(wers):.1f}%, mean: {np.mean(wers):.1f}%, median: {np.median(wers):.1f}%\"\n", + " )\n", + " print(f\"Duration — min: {min(durations):.2f}s, max: {max(durations):.2f}s, total: {sum(durations):.1f}s\")" ] }, { @@ -363,12 +607,16 @@ "name": "stdout", "output_type": "stream", "text": [ - " WER \u2264 5.0%: 37 samples (74%)\n WER \u2264 5.5%: 50 samples (100%)\n WER \u2264 10.0%: 50 samples (100%)\n WER \u2264 25.0%: 50 samples (100%)\n WER \u2264 50.0%: 50 samples (100%)\n WER \u2264 75.0%: 50 samples (100%)\n" + " WER ≤ 5.0%: 37 samples (74%)\n WER ≤ 5.5%: 50 samples (100%)\n WER ≤ 10.0%: 50 samples (100%)\n WER ≤ 25.0%: 50 samples (100%)\n WER ≤ 50.0%: 50 samples (100%)\n WER ≤ 75.0%: 50 samples (100%)\n" ] } ], "source": [ - "thresholds = [5, 5.5, 10, 25, 50, 75]\nfor t in thresholds:\n passing = [r for r in results if r.get(\"wer_pct\", 100) <= t]\n pct = len(passing) / len(results) * 100 if results else 0\n print(f\" WER \u2264 {t:4.1f}%: {len(passing):4d} samples ({pct:.0f}%)\")" + "thresholds = [5, 5.5, 10, 25, 50, 75]\n", + "for t in thresholds:\n", + " passing = [r for r in results if r.get(\"wer_pct\", 100) <= t]\n", + " pct = len(passing) / len(results) * 100 if results else 0\n", + " print(f\" WER ≤ {t:4.1f}%: {len(passing):4d} samples ({pct:.0f}%)\")" ] }, { diff --git a/tutorials/audio/fleurs/pipeline.yaml b/tutorials/audio/fleurs/pipeline.yaml index 2f09915bfa..fb8d9b0483 100644 --- a/tutorials/audio/fleurs/pipeline.yaml +++ b/tutorials/audio/fleurs/pipeline.yaml @@ -41,6 +41,9 @@ lang: hy_am data_split: dev wer_threshold: 5.5 output_dir: ${raw_data_dir}/result/${lang} +max_audio_sec_per_actor: 240.0 +max_inference_duration_s: 120.0 +local_bucketing: true backend: xenna stages: @@ -56,6 +59,9 @@ stages: audio_filepath_key: audio_filepath target_sample_rate: 16000 batch_size: 16 + max_audio_sec_per_actor: ${max_audio_sec_per_actor} + max_inference_duration_s: ${max_inference_duration_s} + local_bucketing: ${local_bucketing} adapter_kwargs: # Avoid NeMo RNNT label-loop CUDA graphs, which are not supported by all # CUDA runtime/driver combinations. This does not disable GPU inference. diff --git a/tutorials/audio/nemo_fastconformer/README.md b/tutorials/audio/nemo_fastconformer/README.md index 7a63b6b509..807c0a52f1 100644 --- a/tutorials/audio/nemo_fastconformer/README.md +++ b/tutorials/audio/nemo_fastconformer/README.md @@ -73,7 +73,10 @@ uses those fields to retain and explain rows that could not be transcribed. | `model_id` | `nvidia/stt_en_fastconformer_ctc_large` | Any compatible pretrained NeMo ASR checkpoint | | `pred_text_key` | `pred_text` | Output transcript column | | `gpus_per_actor` | `1` | GPUs scheduled for each ASR worker; set `0` for CPU | -| `stages.2.batch_size` | `16` | Number of waveforms per NeMo transcription call | +| `stages.2.batch_size` | `16` | Backend candidate-row window supplied to one `process_batch` call | +| `max_audio_sec_per_actor` | `240` | Maximum padded audio seconds in one NeMo adapter call | +| `max_inference_duration_s` | `120` | Model-input ceiling; longer rows are always segmented and stitched afterward | +| `local_bucketing` | `true` | Sort segments in the current process batch by duration before packing | | `stages.2.adapter_kwargs.num_workers` | `0` | NeMo transcription data-loader workers | | `stages.2.adapter_kwargs.enable_local_attention` | `false` | Convert a compatible FastConformer checkpoint to local attention | @@ -84,6 +87,28 @@ stages.2.adapter_kwargs.enable_local_attention=true \ 'stages.2.adapter_kwargs.local_attention_context_size=[128,128]' ``` +The supplied config enables local duration bucketing with three direct stage +settings: + +```yaml +batch_size: 16 +max_audio_sec_per_actor: 240 +max_inference_duration_s: 120 +local_bucketing: true +``` + +`max_audio_sec_per_actor` bounds the padded work of each adapter call as +`longest segment seconds × item count`. The stage always splits audio at +`max_inference_duration_s` first. It then considers all resulting segments +from the current `process_batch` together, packs duration-near segments when +local bucketing is on, and restores segment and parent-row order afterward. +Set `local_bucketing: false` to preserve input order while retaining the same +capacity bound. No grouping crosses a backend `process_batch` boundary. + +See [Local Duration Bucketing for Audio GPU Inference](../../../nemo_curator/stages/audio/inference/README.md) +for the full algorithm, tuning model, correctness invariants, and adoption +guidance for other audio inference stages. + ## Use the adapter in Python ```python @@ -92,6 +117,9 @@ from nemo_curator.stages.audio.inference.asr.stage import ASRStage asr = ASRStage( adapter_target="nemo_curator.models.asr.nemo_asr.NeMoASRAdapter", model_id="nvidia/stt_en_fastconformer_ctc_large", + max_audio_sec_per_actor=240, + max_inference_duration_s=120, + local_bucketing=True, audio_filepath_key="audio_filepath", batch_size=16, ) @@ -107,7 +135,7 @@ audio to the configured `target_sample_rate` before calling the adapter. | Symptom | Action | |---|---| | `ffmpeg` is not found | Install `ffmpeg` and ensure it is on `PATH` | -| CUDA out of memory | Reduce `stages.2.batch_size` or select a smaller checkpoint | +| CUDA out of memory | Reduce `max_audio_sec_per_actor`; keep it at least as large as `max_inference_duration_s` | | Model import fails | Install `audio_cuda12` or `audio_cpu` for your platform | | First run appears idle | Wait for the NeMo checkpoint download and inspect the Ray logs | | Local-attention conversion fails | Disable it or use a FastConformer checkpoint exposing the required conversion APIs | diff --git a/tutorials/audio/nemo_fastconformer/pipeline.yaml b/tutorials/audio/nemo_fastconformer/pipeline.yaml index e35bd67bae..3fedbbd7e7 100644 --- a/tutorials/audio/nemo_fastconformer/pipeline.yaml +++ b/tutorials/audio/nemo_fastconformer/pipeline.yaml @@ -22,6 +22,9 @@ resampled_audio_dir: ${workspace_dir}/audio_resampled model_id: nvidia/stt_en_fastconformer_ctc_large pred_text_key: pred_text gpus_per_actor: 1 +max_audio_sec_per_actor: 240.0 +max_inference_duration_s: 120.0 +local_bucketing: true backend: ray_data execution_mode: streaming @@ -42,6 +45,9 @@ stages: model_id: ${model_id} pred_text_key: ${pred_text_key} batch_size: 16 + max_audio_sec_per_actor: ${max_audio_sec_per_actor} + max_inference_duration_s: ${max_inference_duration_s} + local_bucketing: ${local_bucketing} resources: _target_: nemo_curator.stages.resources.Resources gpus: ${gpus_per_actor} diff --git a/tutorials/audio/qwen_asr/README.md b/tutorials/audio/qwen_asr/README.md index 4a53bf3b6d..1e4e73a033 100644 --- a/tutorials/audio/qwen_asr/README.md +++ b/tutorials/audio/qwen_asr/README.md @@ -67,6 +67,9 @@ input can run on a 12 GB GPU. This does not change the adapter's default. | Executor | Ray Data | | ASR stage batch size | `128` | | GPUs per ASR actor | `1` | +| Maximum padded audio per adapter call | `240` seconds | +| Maximum single model input | `120` seconds | +| Local duration bucketing | enabled | | Hugging Face model revision | `null` (Hugging Face default) | | GPU memory limit | `0.7` of device memory | | Maximum vLLM model length | `8192` tokens | @@ -75,8 +78,12 @@ input can run on a 12 GB GPU. This does not change the adapter's default. | Prediction field | `pred_text` | | Adapter extras field | `asr_extras` | -The stage batch is passed to one adapter `transcribe_batch()` call. The adapter -forwards `max_inference_batch_size` to Qwen3-ASR as its internal cap. +The stage always segments audio at `max_inference_duration_s`. It then packs +all resulting segments from the current process batch under +`max_audio_sec_per_actor`, using `item count × longest segment duration` as +the padded-work estimate. Local bucketing sorts by duration before packing; +disable it with `local_bucketing=false` to preserve input order. The adapter +also forwards `max_inference_batch_size` to Qwen3-ASR as its internal cap. ## Select the executor diff --git a/tutorials/audio/qwen_asr/pipeline.yaml b/tutorials/audio/qwen_asr/pipeline.yaml index 0cc3cffb55..d528fe8f4c 100644 --- a/tutorials/audio/qwen_asr/pipeline.yaml +++ b/tutorials/audio/qwen_asr/pipeline.yaml @@ -27,6 +27,9 @@ supported_language_codes: [zh, en, yue, ar, de, fr, es, pt, id, it, ko, ru, th, pred_text_key: pred_text extras_key: asr_extras gpus_per_actor: 1 +max_audio_sec_per_actor: 240.0 +max_inference_duration_s: 120.0 +local_bucketing: true backend: ray_data execution_mode: streaming @@ -50,6 +53,9 @@ stages: pred_text_key: ${pred_text_key} extras_key: ${extras_key} batch_size: 128 + max_audio_sec_per_actor: ${max_audio_sec_per_actor} + max_inference_duration_s: ${max_inference_duration_s} + local_bucketing: ${local_bucketing} resources: _target_: nemo_curator.stages.resources.Resources gpus: ${gpus_per_actor} diff --git a/tutorials/audio/qwen_omni_inprocess/README.md b/tutorials/audio/qwen_omni_inprocess/README.md index 0705ac11e6..57c057e797 100644 --- a/tutorials/audio/qwen_omni_inprocess/README.md +++ b/tutorials/audio/qwen_omni_inprocess/README.md @@ -63,6 +63,9 @@ explicit: | Executor | Ray Data | | ASR stage batch size | `32` | | GPUs per ASR actor / derived vLLM tensor parallelism | `2` / `2`, from `gpus_per_actor` | +| Maximum padded audio per adapter call | `2400` seconds | +| Maximum single model input | `2400` seconds | +| Local duration bucketing | enabled | | Prompt | `Transcribe the audio.` | | Prompt content order | text, then audio | | Concurrent vLLM sequences | `16` | @@ -180,6 +183,8 @@ annotated with `language_missing` unless `default_language` is explicitly set. This is a functional, local manifest-to-transcript example. It does not provide recovery ASR, hallucination filtering, WER calculation, -duration-aware bucketing, sharded resumability, or benchmark reporting. +sharded resumability, or benchmark reporting. Local duration bucketing is +limited to the rows and model-safe segments in one `process_batch` call; it +does not globally reorder the manifest. Validate output quality and row accounting on representative audio before larger runs. diff --git a/tutorials/audio/qwen_omni_inprocess/pipeline.yaml b/tutorials/audio/qwen_omni_inprocess/pipeline.yaml index d3f20a0801..e61844da8b 100644 --- a/tutorials/audio/qwen_omni_inprocess/pipeline.yaml +++ b/tutorials/audio/qwen_omni_inprocess/pipeline.yaml @@ -28,6 +28,9 @@ prompt_text: Transcribe the audio. prompt_file: null pred_text_key: pred_text gpus_per_actor: 2 +max_audio_sec_per_actor: 2400.0 +max_inference_duration_s: 2400.0 +local_bucketing: true backend: ray_data execution_mode: streaming @@ -50,6 +53,9 @@ stages: supported_language_codes: ${supported_language_codes} pred_text_key: ${pred_text_key} batch_size: 32 + max_audio_sec_per_actor: ${max_audio_sec_per_actor} + max_inference_duration_s: ${max_inference_duration_s} + local_bucketing: ${local_bucketing} resources: _target_: nemo_curator.stages.resources.Resources gpus: ${gpus_per_actor}