nemotron-asr reproducers (single/batch/streaming RNNT)

_hf_capture.py

"""Capture the HF NemotronAsrForRNNT outputs to bake into the integration tests.

Not a reproducer (those are NeMo reference). This runs the *HF* model on the same audio the tests use:
  - offline single   (librispeech_asr_dummy sample 0, default att_context)
  - offline batched   (librispeech_asr_dummy samples 0..4)
  - streaming         (obama.mp3, att_context_size=[70, 6], mel-frame chunks 49 then 56)
and prints JSON so the exact HF strings can be pasted into EXPECTED_* in the test.
"""

import json
from threading import Thread

import torch
from datasets import Audio, load_dataset

from transformers import AutoModelForRNNT, AutoProcessor, TextIteratorStreamer
from transformers.audio_utils import load_audio


MODEL_ID = "/raid/eustache/nemotron-speech-streaming-en-0.6b-hf"
device = "cuda" if torch.cuda.is_available() else "cpu"

processor = AutoProcessor.from_pretrained(MODEL_ID)
sr = processor.feature_extractor.sampling_rate
model = AutoModelForRNNT.from_pretrained(MODEL_ID, dtype=torch.float32, device_map=device).eval()

# --- offline single + batched on librispeech_asr_dummy ---------------------------------------------
ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
ds = ds.cast_column("audio", Audio(sampling_rate=sr))
samples = [x["array"] for x in ds.sort("id")[:5]["audio"]]


def transcribe_offline(audio_samples):
    inputs = processor(audio_samples, sampling_rate=sr)
    inputs.to(model.device, dtype=model.dtype)
    out = model.generate(**inputs, return_dict_in_generate=True)
    return processor.batch_decode(out.sequences, skip_special_tokens=True)


offline_single = transcribe_offline(samples[:1])
offline_batched = transcribe_offline(samples[:5])

# --- streaming on obama.mp3 ------------------------------------------------------------------------
audio = load_audio(
    "https://huggingface.co/datasets/hf-internal-testing/dummy-audio-samples/resolve/main/obama.mp3",
    sampling_rate=sr,
)
stream_inputs = processor(audio, sampling_rate=sr)
stream_inputs.to(model.device, dtype=model.dtype)


def input_features_generator():
    start_idx, first_chunk_size, chunk_size = 0, 49, 56
    chunk = first_chunk_size
    input_length = stream_inputs.input_features.shape[1]
    while start_idx < input_length:
        end_idx = min(start_idx + chunk, input_length)
        yield stream_inputs.input_features[:, start_idx:end_idx, :]
        start_idx = end_idx
        chunk = chunk_size


streamer = TextIteratorStreamer(processor.tokenizer, skip_special_tokens=True, clean_up_tokenization_spaces=True)
generate_kwargs = {
    "input_features": input_features_generator(),
    "att_context_size": [70, 6],
    "streamer": streamer,
}
thread = Thread(target=model.generate, kwargs=generate_kwargs)
thread.start()
streaming = "".join(streamer)
thread.join()

print("HF_CAPTURE_JSON_START")
print(
    json.dumps(
        {
            "offline_single": offline_single,
            "offline_batched": offline_batched,
            "streaming": streaming,
        },
        indent=2,
        ensure_ascii=False,
    )
)
print("HF_CAPTURE_JSON_END")

pyproject.toml

[project]
name = "nemotron-asr-reproducers"
version = "0.0.0"
description = "NeMo reference reproducers for the HF NemotronAsrForRNNT integration-test expected values."
# NeMo's deps lag on the newest CPython; 3.10–3.12 resolve cleanly.
requires-python = ">=3.10,<3.13"
dependencies = [
    # The cache-aware streaming Nemotron checkpoint (nvidia/nemotron-speech-streaming-en-0.6b) was validated
    # against a NeMo 2.8.0rc0 dev build (the same build used elsewhere in this project). That pre-release is
    # not on PyPI, so `uv run` resolves the newest published `nemo-toolkit` instead. If a published release
    # cannot load the checkpoint, run the reproducers from the project's NeMo dev environment instead (see
    # run_reproducers.sh, which honours $NEMO_PYTHON).
    "nemo_toolkit[asr]>=2.7.3",
    "datasets",
    "soundfile",
    "librosa",
]

reproducer_batch_rnnt.py

"""Reference values for the batched NemotronAsrForRNNT integration test, straight from NeMo.

Greedy RNN-T transcription of the first 5 librispeech_asr_dummy samples (batched) with the original NeMo
cache-aware streaming model `nvidia/nemotron-speech-streaming-en-0.6b`, run *offline* (full context) at the
widest attention context `[70, 13]` — the setting the HF model uses by default for non-streaming `generate`.

Deps come from this folder's pyproject.toml. Run with:
    uv run reproducer_batch_rnnt.py [output.json]
"""

import io
import json
import os
import sys
import tempfile

import nemo.collections.asr as nemo_asr
import soundfile as sf
from datasets import Audio, load_dataset


MODEL_NAME = "nvidia/nemotron-speech-streaming-en-0.6b"
ATT_CONTEXT_SIZE = [70, 13]  # widest = best WER; matches the HF offline default (first config entry)
SAMPLING_RATE = 16000
NUM_SAMPLES = 5

model = nemo_asr.models.ASRModel.from_pretrained(model_name=MODEL_NAME).eval()
model.encoder.set_default_att_context_size(att_context_size=ATT_CONTEXT_SIZE)

ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
ds = ds.cast_column("audio", Audio(decode=False)).sort("id")

with tempfile.TemporaryDirectory() as tmp:
    paths = []
    for i in range(NUM_SAMPLES):
        a = ds[i]["audio"]
        arr, sr = sf.read(io.BytesIO(a["bytes"]) if a.get("bytes") else a["path"])
        assert sr == SAMPLING_RATE
        p = os.path.join(tmp, f"sample_{i}.wav")
        sf.write(p, arr, SAMPLING_RATE)
        paths.append(p)
    hyps = model.transcribe(paths, batch_size=NUM_SAMPLES)

transcriptions = [h.text if hasattr(h, "text") else h for h in hyps]
payload = json.dumps({"att_context_size": ATT_CONTEXT_SIZE, "transcriptions": transcriptions}, indent=2)

if len(sys.argv) > 1:
    with open(sys.argv[1], "w") as f:
        f.write(payload + "\n")
else:
    print(payload)

reproducer_single_rnnt.py

"""Reference value for the single-sample NemotronAsrForRNNT integration test, straight from NeMo.

Greedy RNN-T transcription of the first librispeech_asr_dummy sample with the original NeMo cache-aware
streaming model `nvidia/nemotron-speech-streaming-en-0.6b`, run *offline* (full context) at the widest
attention context `[70, 13]` — the setting the HF model uses by default for non-streaming `generate`.

Deps come from this folder's pyproject.toml. Run with:
    uv run reproducer_single_rnnt.py [output.json]
With no argument the JSON is printed; with a path it is written there (used by run_reproducers.sh).
"""

import io
import json
import os
import sys
import tempfile

import nemo.collections.asr as nemo_asr
import soundfile as sf
from datasets import Audio, load_dataset


MODEL_NAME = "nvidia/nemotron-speech-streaming-en-0.6b"
ATT_CONTEXT_SIZE = [70, 13]  # widest = best WER; matches the HF offline default (first config entry)
SAMPLING_RATE = 16000
NUM_SAMPLES = 1

model = nemo_asr.models.ASRModel.from_pretrained(model_name=MODEL_NAME).eval()
model.encoder.set_default_att_context_size(att_context_size=ATT_CONTEXT_SIZE)

# This NeMo env's `datasets` may lack a torch codec to auto-decode audio, so decode the raw bytes ourselves with
# soundfile (librispeech_asr_dummy is already 16 kHz mono) and hand NeMo plain wav paths.
ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
ds = ds.cast_column("audio", Audio(decode=False)).sort("id")

with tempfile.TemporaryDirectory() as tmp:
    paths = []
    for i in range(NUM_SAMPLES):
        a = ds[i]["audio"]
        arr, sr = sf.read(io.BytesIO(a["bytes"]) if a.get("bytes") else a["path"])
        assert sr == SAMPLING_RATE
        p = os.path.join(tmp, f"sample_{i}.wav")
        sf.write(p, arr, SAMPLING_RATE)
        paths.append(p)
    hyps = model.transcribe(paths, batch_size=NUM_SAMPLES)

transcriptions = [h.text if hasattr(h, "text") else h for h in hyps]
payload = json.dumps({"att_context_size": ATT_CONTEXT_SIZE, "transcriptions": transcriptions}, indent=2)

if len(sys.argv) > 1:
    with open(sys.argv[1], "w") as f:
        f.write(payload + "\n")
else:
    print(payload)

reproducer_streaming_rnnt.py

"""Reference value for the streaming NemotronAsrForRNNT integration test, straight from NeMo.

Simulated cache-aware streaming RNN-T transcription of the `obama.mp3` dummy sample with the original NeMo
`nvidia/nemotron-speech-streaming-en-0.6b`, at attention context `[70, 6]` (the latency the HF streaming test
uses). This mirrors NeMo's reference
`examples/asr/asr_cache_aware_streaming/speech_to_text_cache_aware_streaming_infer.py`, trimmed to a single
file: the audio is fed to the FastConformer cache-aware encoder chunk by chunk while the decoder state is
threaded across chunks.

Unlike the multilingual `nemotron-3.5-asr-streaming-0.6b`, the English-only checkpoint here is NOT
prompt-conditioned, so the language-prompt / lang-tag handling of the generic reference is dropped.

The HF model is a re-implementation and chunks the *mel frames* (49 then 56) rather than driving NeMo's
`CacheAwareStreamingAudioBuffer`, so the HF streaming transcript may differ from this NeMo reference by a few
sub-word emissions (~1e-3 numerical drift over the conformer stack flipping borderline greedy emissions). This
JSON is the NeMo reference used to sanity-check / annotate the HF expected value baked into the test.

Deps come from this folder's pyproject.toml. Run with:
    uv run reproducer_streaming_rnnt.py [output.json]
"""

import io
import json
import sys
import urllib.request

import librosa
import torch
from omegaconf import open_dict

import nemo.collections.asr as nemo_asr
from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
from nemo.collections.asr.parts.utils.streaming_utils import CacheAwareStreamingAudioBuffer


MODEL_NAME = "nvidia/nemotron-speech-streaming-en-0.6b"
AUDIO_URL = "https://huggingface.co/datasets/hf-internal-testing/dummy-audio-samples/resolve/main/obama.mp3"
ATT_CONTEXT_SIZE = [70, 6]
SAMPLING_RATE = 16000
MAX_SYMBOLS = 10


def text_of(hyps):
    """conformer_stream_step returns either Hypothesis objects or plain strings."""
    if hyps and isinstance(hyps[0], Hypothesis):
        return hyps[0].text
    return hyps[0] if hyps else ""


torch.set_grad_enabled(False)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = nemo_asr.models.ASRModel.from_pretrained(model_name=MODEL_NAME, map_location=device)
# Cache-aware models only run correctly in float32 (some layers force-cast).
model = model.to(device=device, dtype=torch.float32).eval()

# Use torch SDPA (read live in each attention forward) to match the HF encoder.
for module in model.modules():
    if hasattr(module, "use_pytorch_sdpa"):
        module.use_pytorch_sdpa = True

model.encoder.set_default_att_context_size(att_context_size=ATT_CONTEXT_SIZE)

# RNN-T greedy streaming decoding; fused_batch_size=-1 threads the decoder state across chunks.
decoding_cfg = model.cfg.decoding
with open_dict(decoding_cfg):
    decoding_cfg.strategy = "greedy_batch"
    decoding_cfg.fused_batch_size = -1
    if "greedy" in decoding_cfg:
        decoding_cfg.greedy.max_symbols = MAX_SYMBOLS
model.change_decoding_strategy(decoding_cfg)

# Load obama.mp3 as mono 16 kHz (handles mp3 decode + downmix), like the HF streaming test input.
with urllib.request.urlopen(AUDIO_URL) as resp:
    samples, _ = librosa.load(io.BytesIO(resp.read()), sr=SAMPLING_RATE, mono=True)

buffer = CacheAwareStreamingAudioBuffer(model=model, online_normalization=False)
buffer.append_audio(samples)

cache_last_channel, cache_last_time, cache_last_channel_len = model.encoder.get_initial_cache_state(batch_size=1)

previous_hypotheses = None
pred_out_stream = None
transcript = ""
for step, (chunk_audio, chunk_lengths) in enumerate(buffer):
    with torch.inference_mode():
        (
            pred_out_stream,
            transcribed_texts,
            cache_last_channel,
            cache_last_time,
            cache_last_channel_len,
            previous_hypotheses,
        ) = model.conformer_stream_step(
            processed_signal=chunk_audio.to(torch.float32),
            processed_signal_length=chunk_lengths,
            cache_last_channel=cache_last_channel,
            cache_last_time=cache_last_time,
            cache_last_channel_len=cache_last_channel_len,
            keep_all_outputs=buffer.is_buffer_empty(),
            previous_hypotheses=previous_hypotheses,
            previous_pred_out=pred_out_stream,
            drop_extra_pre_encoded=(0 if step == 0 else model.encoder.streaming_cfg.drop_extra_pre_encoded),
            return_transcription=True,
        )
    transcript = text_of(transcribed_texts)

payload = json.dumps(
    {"att_context_size": ATT_CONTEXT_SIZE, "audio": "obama.mp3", "transcription": transcript}, indent=2
)

if len(sys.argv) > 1:
    with open(sys.argv[1], "w") as f:
        f.write(payload + "\n")
else:
    print(payload)

run_reproducers.sh

#!/usr/bin/env bash
#
# Regenerate the NemotronAsrForRNNT integration-test reference values from the ORIGINAL NeMo model
# (nvidia/nemotron-speech-streaming-en-0.6b). Each reproducer writes its JSON straight into the test
# fixtures dir (transformers/tests/fixtures/nemotron_asr/), where `NemotronAsrForRNNTIntegrationTest`
# (tests/models/nemotron_asr/test_modeling_nemotron_asr.py) loads it as the expected values — the same
# fixture pattern as the Parakeet CTC/TDT integration tests.
#
# These are the NeMo *reference* outputs. Offline single/batch match the HF re-implementation exactly. The
# streaming reference can differ from the HF output by a sub-word (~1e-3 numerical drift in the re-implemented
# FastConformer encoder flipping a borderline greedy emission); the streaming test asserts against this NeMo
# reference anyway, so it is expected to fail on that one sub-word until the drift is addressed.
#
# Override the destination with FIXTURES_DIR=/some/path ./run_reproducers.sh
#
# Usage:
#   ./run_reproducers.sh
#
# By default each reproducer is run with `uv run` (deps resolved from this folder's pyproject.toml). The
# cache-aware streaming checkpoint needs a NeMo 2.8.0rc0 dev build that is not on PyPI, so if `uv run` cannot
# load it, point the runner at a Python that already has that NeMo build:
#   NEMO_PYTHON=/path/to/nemo-venv/bin/python ./run_reproducers.sh
#
# NOTE: these are the NeMo *reference* values. The HF FastConformer encoder is a re-implementation, so a few
# greedy emissions can drift (especially in streaming, which also chunks differently). Review against the HF
# output before baking values into the test; the streaming test annotates any drift inline.

set -uo pipefail

here="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
out="${FIXTURES_DIR:-$here/../transformers/tests/fixtures/nemotron_asr}"
mkdir -p "$out"

if [ -n "${NEMO_PYTHON:-}" ]; then
    runner=("$NEMO_PYTHON")
else
    runner=(uv run)
fi

status=0
run() {  # <reproducer.py> <output.json>
    local script="$1" dst="$out/$2"
    echo "==> ${runner[*]} $script  ->  $dst"
    if ( cd "$here" && "${runner[@]}" "$script" "$dst" ); then
        echo "    OK"
    else
        echo "    FAILED" >&2
        status=1
    fi
}

run reproducer_single_rnnt.py    expected_results_single.json
run reproducer_batch_rnnt.py     expected_results_batch.json
run reproducer_streaming_rnnt.py expected_results_streaming.json

exit $status