fix(audio): benchmark tagging on pre-staged AMI meetings

[mohammadaaftabv]

Summary

Fix the two audio-tagging nightly benchmarks so they run representative AMI meeting audio and cannot pass after doing no useful work.

The effective PR is now 8 files (+1136/-116). It keeps the benchmark focused on tagging, removes the earlier unrelated FLEURS/tutorial/notebook/runtime-download scope, and makes the AMI setup reproducible instead of relying on hand-built manifests or manual model copies. Test-file changes have been removed from the net PR diff.

Workload

The entries use three AMI single-distant-microphone meetings with long, multi-speaker audio and overlap. The repeat entry processes each meeting twice with deterministic repeated IDs.

Entry	Rows	Source audio
audio_tagging_tts_xenna	3	1.563 h
audio_tagging_tts_xenna_repeat	6	3.126 h

FLEURS was not retained for tagging because its short, mostly single-speaker clips do not exercise the long-audio splitting, multi-speaker diarization, or overlap behavior that dominates this pipeline.

Reproducible Setup

Audio tagging now follows the same two-layer setup pattern as FLEURS:

1. Run a one-time data-prep script on the benchmark host.
2. Run nightly with --raw-data-dir and --no-auto-download.

One-time setup:

python benchmarking/data_prep/prepare_audio_tagging_data.py \
  --output-path {datasets_path}/audio_tagging_ami_sdm \
  --model-output-path {model_weights_path}/audio_tagging/pyannote-speaker-diarization-community-1 \
  --hf-repo-id <token-free-audio-tagging-assets-repo>

If the PyAnnote snapshot already exists locally, pass --model-source-path to copy it instead of downloading it from the token-free assets repo.

The prep script stages:

• manifest.jsonl
• audio/EN2002b.Array1-01.wav
• audio/ES2004c.Array1-01.wav
• audio/TS3003a.Array1-01.wav
• pyannote-speaker-diarization-community-1/

It also verifies that the manifest has exactly the expected AMI rows, expected audio basenames, and unique nonempty audio_item_id values. That makes missing or mis-staged data fail during setup instead of surfacing only after an expensive nightly run.

Nightly Behavior

Nightly consumes only pre-staged local paths:

--scratch-output-path={session_entry_dir}/scratch
--raw-data-dir={dataset:audio_tagging_ami_sdm,files}
--no-auto-download
--diarization-model-path={model_weights_path}/audio_tagging/pyannote-speaker-diarization-community-1

Nightly does not need HF_TOKEN, does not download AMI or PyAnnote assets at runtime, and fails loudly if the staged dataset/model is missing or malformed.

Standalone debugging can still auto-download into <scratch-output-path>/audio_tagging_ami_sdm when --raw-data-dir is omitted and --hf-repo-id or CURATOR_AUDIO_TAGGING_HF_REPO_ID is set. The standalone cache uses CURATOR_AUDIO_TAGGING_CACHE_DIR or /tmp/curator/audio_tagging_cache, so deleting scratch does not force blob refetches.

Correctness Contract

Every downstream stage preserves outer AudioTask rows, so success requires:

input manifest rows * repeat factor
  == returned pipeline tasks
  == written output-manifest rows

Success also requires:

• complete second-pass ASR and finite WER output
• nonzero work from all 12 downstream stages, including ASRAlignment2, ComputeWER, and ManifestWriter
• positive source-audio duration
• nightly floors for complete segments, segment-output coverage, tagged duration, and throughput
• a nonzero exit code with persisted error metrics on validation failure

Repeated inputs receive deterministic IDs so resampled audio and RTTM files cannot collide.

Runtime Fixes

• Allow PyAnnote to load a local pipeline without a token and keep tokens out of stage representations.
• Declare the correct NeMoASRAlignerStage schema for second-pass segment mode.
• Reject empty and zero-duration prepared segments.
• Keep local-GPU batch and CUDA-graph controls while leaving nightly defaults unchanged.
• Fix RepeatEntriesStage so it accepts valid falsy IDs like integer 0 while rejecting only None and the empty string.

Validation

Current checks after the setup refactor:

• ruff check benchmarking/data_prep/prepare_audio_tagging_data.py benchmarking/scripts/audio_tagging_benchmark.py
• ruff format --check benchmarking/data_prep/prepare_audio_tagging_data.py benchmarking/scripts/audio_tagging_benchmark.py
• python3 -m py_compile benchmarking/data_prep/prepare_audio_tagging_data.py benchmarking/scripts/audio_tagging_benchmark.py
• python3 benchmarking/data_prep/prepare_audio_tagging_data.py --help
• YAML assertion that both tagging nightly entries pass --raw-data-dir={dataset:audio_tagging_ami_sdm,files}, pass --no-auto-download, and no longer pass --input-manifest
• git diff --check

Earlier full one-GPU PR validation, before the data-prep-only refactor:

• 3 input rows = 3 returned tasks = 3 written rows
• all 3 meetings produced complete output
• 182 complete segments from 250 emitted
• all 12 required stages processed data
• 1.563 source-audio hours and 0.279 tagged-audio hours
• 337.6 seconds pipeline time, 346.9 seconds including subprocess overhead
• 16.67 audio-hours per wall-hour
• all configured benchmark requirements passed

The full GPU benchmark was not rerun after this data-prep refactor in the current local environment.

Checklist

• I am familiar with the Contributing Guide.
• Test-file changes are intentionally absent from this PR; validation uses existing coverage plus static/config checks to match the FLEURS benchmark-script convention.
• Documentation reflects the actual workload, setup flow, standalone flow, and pass/fail contract.

[copy-pr-bot]

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

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR fixes two audio-tagging nightly benchmarks by replacing their previous reliance on hand-built manifests and an HF_SECRET_KEY token with a reproducible two-layer setup: a one-time prepare_audio_tagging_data.py script stages three AMI SDM meetings and a local PyAnnote snapshot, and the nightly entries then consume those pre-staged paths via --raw-data-dir / --no-auto-download.

• Core runtime fixes: hf_token on PyAnnoteDiarizationStage is made optional (str | None, repr=False) so local model paths work without a token; NeMoASRAlignerStage gains the correct inputs()/outputs() schema for infer_segment_only mode and a new use_cuda_graphs field; PrepareModuleSegmentsStage.is_valid_segment now rejects empty-word and zero-duration segments.
• Correctness contract: _validate_outputs enforces input-row == output-row == manifest-row counts, nonzero work from all 12 required stages, and finite WER; nightly YAML requirements add explicit floors for segment counts, coverage ratios, audio duration, and throughput.
• RepeatEntriesStage fix: The previous not source_id guard is replaced with source_id is None or source_id == "", correctly preserving valid falsy IDs like integer 0.

Confidence Score: 5/5

Safe to merge; all core logic changes are correct and the correctness-contract validation adds meaningful regression protection.

The runtime fixes (optional hf_token, correct stage schemas, segment rejection, RepeatEntriesStage ID guard) are all narrow and correct. The new benchmark correctness contract is well-designed — it raises on row loss, skipped stages, and zero tagged output, so the benchmark can no longer pass silently. The only findings are a documentation mismatch where CURATOR_AUDIO_TAGGING_HF_REPO_ID cannot actually override the dataset repo in the prep script, and the row-index-based AMI meeting selection which has no meeting-ID cross-check.

benchmarking/data_prep/prepare_audio_tagging_data.py — env-var override is unreachable and meeting selection has no ID validation against actual dataset rows.

Important Files Changed

Filename	Overview
benchmarking/data_prep/prepare_audio_tagging_data.py	New one-time data-prep script for AMI SDM meetings + PyAnnote snapshot; row-index-based meeting selection is fragile, and the env-var override documented in help text is unreachable.
benchmarking/scripts/audio_tagging_benchmark.py	Substantial rewrite: replaces HF-token + hand-authored manifest flow with reproducible AMI staging + correctness-contract validation; logic is sound and the new metrics are wired up correctly.
benchmarking/scripts/utils.py	RepeatEntriesStage: falsy-but-valid ID fix (is None or == '' instead of not source_id) and deterministic suffix assignment are both correct.
nemo_curator/stages/audio/inference/speaker_diarization/pyannote.py	hf_token made optional (str
nemo_curator/stages/audio/tagging/inference/nemo_asr_align.py	Correct inputs()/outputs() schema for infer_segment_only mode and new use_cuda_graphs field with proper CTC/RNNT config paths.
nemo_curator/stages/audio/tagging/prepare_module_segments.py	is_valid_segment now rejects empty word lists and zero/negative-duration segments before the single-word over-length check; safe change.
benchmarking/nightly-benchmark.yaml	Both tagging entries updated to use --raw-data-dir / --no-auto-download, removing HF_SECRET_KEY; new correctness-contract requirements (row counts, coverage, duration floors) added.
benchmarking/README.md	Documentation additions for audio benchmark data setup and the audio tagging benchmark flow are accurate and match the code.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Prep as prepare_audio_tagging_data.py
    participant HF as HuggingFace
    participant Disk as Staged Disk
    participant Bench as audio_tagging_benchmark.py
    participant Pipeline as NeMo Pipeline

    Note over Prep,Disk: One-time setup (run once on benchmark host)
    Prep->>HF: load_dataset(diarizers-community/ami, sdm, test)
    HF-->>Prep: streaming rows at indices 1,6,12
    Prep->>Disk: audio/EN2002b.Array1-01.wav + ES2004c + TS3003a
    Prep->>Disk: manifest.jsonl (3 rows)
    Prep->>HF: snapshot_download(pyannote-community/speaker-diarization-community-1)
    HF-->>Prep: model snapshot
    Prep->>Disk: pyannote-speaker-diarization-community-1/

    Note over Bench,Pipeline: Nightly run (--raw-data-dir + --no-auto-download)
    Bench->>Disk: _locate_prestaged_data(raw_data_dir)
    Disk-->>Bench: manifest.jsonl + audio/
    Bench->>Bench: "_write_staged_manifest -> run manifest"
    Bench->>Pipeline: "ManifestReader -> [RepeatEntries] -> ResampleAudio"
    Pipeline->>Pipeline: PyAnnoteDiarization (local model, no token)
    Pipeline->>Pipeline: "SplitLongAudio -> ASRAlignment -> JoinSplitMetadata"
    Pipeline->>Pipeline: "MergeAlignmentDiar -> BandwidthEstimation -> SquimMetrics"
    Pipeline->>Pipeline: "PrepareModuleSegments -> ASRAlignment2 -> ComputeWER"
    Pipeline->>Pipeline: ManifestWriter
    Pipeline-->>Bench: tasks[]
    Bench->>Bench: _validate_outputs (row counts, stage coverage, duration floors)
    Bench->>Disk: benchmark_results.json

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Prep as prepare_audio_tagging_data.py
    participant HF as HuggingFace
    participant Disk as Staged Disk
    participant Bench as audio_tagging_benchmark.py
    participant Pipeline as NeMo Pipeline

    Note over Prep,Disk: One-time setup (run once on benchmark host)
    Prep->>HF: load_dataset(diarizers-community/ami, sdm, test)
    HF-->>Prep: streaming rows at indices 1,6,12
    Prep->>Disk: audio/EN2002b.Array1-01.wav + ES2004c + TS3003a
    Prep->>Disk: manifest.jsonl (3 rows)
    Prep->>HF: snapshot_download(pyannote-community/speaker-diarization-community-1)
    HF-->>Prep: model snapshot
    Prep->>Disk: pyannote-speaker-diarization-community-1/

    Note over Bench,Pipeline: Nightly run (--raw-data-dir + --no-auto-download)
    Bench->>Disk: _locate_prestaged_data(raw_data_dir)
    Disk-->>Bench: manifest.jsonl + audio/
    Bench->>Bench: "_write_staged_manifest -> run manifest"
    Bench->>Pipeline: "ManifestReader -> [RepeatEntries] -> ResampleAudio"
    Pipeline->>Pipeline: PyAnnoteDiarization (local model, no token)
    Pipeline->>Pipeline: "SplitLongAudio -> ASRAlignment -> JoinSplitMetadata"
    Pipeline->>Pipeline: "MergeAlignmentDiar -> BandwidthEstimation -> SquimMetrics"
    Pipeline->>Pipeline: "PrepareModuleSegments -> ASRAlignment2 -> ComputeWER"
    Pipeline->>Pipeline: ManifestWriter
    Pipeline-->>Bench: tasks[]
    Bench->>Bench: _validate_outputs (row counts, stage coverage, duration floors)
    Bench->>Disk: benchmark_results.json

Loading

_{Reviews (7): Last reviewed commit: "Make audio tagging setup download reprod..." | Re-trigger Greptile}

[greptile-apps]

Tutorial now fails hard on empty-segment pipelines

validate_tagging_outputs raises RuntimeError when the pipeline produces zero tagged segments. Because there is no try/except around this call in the tutorial's main, any small or synthetic dataset that passes diarization but produces no PrepareModuleSegments output will crash with an unhandled exception — the PIPELINE COMPLETE block is never reached. Previously the tutorial logged num_tasks = 0 and completed normally, so this is a behavioural break for developers iterating against toy fixtures. Consider documenting in the tutorial that a non-trivial audio dataset producing at least one segment is required, or catching the validation error and printing a user-friendly hint before re-raising.

[greptile-apps]

Use an explicit None/empty-string check so that a legitimate ID value of 0 (or any other falsy non-empty value) does not incorrectly raise.

Suggested change

	source_id = data.get(self._unique_id_key)
	if not source_id:
	msg = f"Cannot repeat entry without '{self._unique_id_key}'"
	raise ValueError(msg)
	source_id = data.get(self._unique_id_key)
	if source_id is None or source_id == "":
	msg = f"Cannot repeat entry without '{self._unique_id_key}'"
	raise ValueError(msg)

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[github-actions]

🌿 Preview your docs: https://nvidia-preview-aaftabv-fix-audio-tagging-benchmarks.docs.buildwithfern.com/nemo/curator

Here are the markdown pages you've updated:

• nemo/curator/curate-audio/load-data/fleurs-dataset
• nemo/curator/curate-audio/tutorials/beginner
• nemo/curator/get-started/audio