feat(stt): Add a Nemotron 3.5 ASR Finetuning Recipe

examples/nemotron-asr-nemo/.python-version

@@ -0,0 +1 @@
+3.12.0

examples/nemotron-asr-nemo/README.md

@@ -0,0 +1,79 @@
+# Nemotron 3.5 ASR Fine-Tuning with NeMo
+
+This example fine-tunes NVIDIA's [Nemotron 3.5 ASR streaming](https://huggingface.co/nvidia/nemotron-3.5-asr-streaming-0.6b) speech recognition model using the NVIDIA NeMo framework on Baseten.
+
+Nemotron 3.5 ASR is a 600M-parameter, multilingual (40 language-locales), real-time streaming model built on a Cache-Aware FastConformer-RNNT architecture with prompt-based language conditioning. This recipe does a full fine-tune from the base `.nemo` checkpoint using NeMo's `speech_to_text_finetune.py`, following NVIDIA's [official fine-tuning recipe](https://huggingface.co/blog/nvidia/fine-tuning-nemotron-35-asr). It uses the small public AN4 corpus as a stand-in so it runs end-to-end quickly.
+
+The job runs on the `nvcr.io/nvidia/nemo:26.02` container — the last consolidated NeMo Framework image that bundles the ASR/TTS collection with a pre-tested CUDA/numba/lightning/kaldialign stack (the newer `26.06` image is Megatron/LLM-only, and the collection now lives in [`NVIDIA-NeMo/Speech`](https://github.com/NVIDIA-NeMo/Speech)). `run.sh` uses the container's bundled NeMo when it has the streaming-prompt recipe, and otherwise clones it and prepends it to `PYTHONPATH` so the recipe code runs against the container's already-installed, correctly CUDA-linked deps. The base checkpoint is delivered read-only via the [Baseten Delivery Network (BDN)](https://docs.baseten.co/training/concepts/storage) — declared as a `WeightsSource` in `config.py` — so it's on local disk before training starts and never costs billed GPU time to download.
+
+**Resources:** 1 node, 1x H100 GPU
+
+## Prerequisites
+
+1. [Create a Baseten account](https://baseten.co/signup) if you don't already have one.
+2. Add a Hugging Face access token as the Baseten secret `hf_access_token` (used to download the gated base checkpoint). See [secrets](https://app.baseten.co/settings/secrets).
+3. Install the Truss CLI:
+   ```bash
+   # pip
+   pip install -U truss
+   # or uv
+   uv add truss
+   ```
+
+## Getting Started
+
+Initialize the example, navigate into the directory, and push the training job:
+
+```bash
+truss train init --examples nemotron-asr-nemo
+cd nemotron-asr-nemo
+truss train push config.py
+```
+
+## Using Your Own Data
+
+The example trains on AN4 purely as a smoke test. To fine-tune on your own language, domain, or accent, replace `prepare_data.py` with your own manifest builder that emits NeMo JSON-lines manifests:
+
+```json
+{"audio_filepath": "/abs/path/clip.wav", "duration": 4.27, "text": "Reference transcript.", "lang": "en-US", "target_lang": "en-US"}
+```
+
+Two details matter most:
+
+- **Every clip needs a `target_lang` tag** matching a locale the model recognizes (e.g. `en-US`, `es-ES`, `el-GR`, `bg-BG`). This drives the model's prompt-based language conditioning and is unforgiving of mismatched labels.
+- **Match the base model's text style** — punctuated, properly-cased transcripts, since that's what the model produces.
+
+For larger datasets, point the trainer at tarred NeMo/Lhotse shards and drive training with a fixed step budget (`trainer.max_steps`) rather than epochs. When specializing on a few languages in the multilingual model, blend in a slice of the other languages ("replay") to avoid eroding them.
+
+For real, **frozen** datasets (e.g. tarred shards on Hugging Face, S3, or GCS), mount them through BDN the same way as the base checkpoint rather than downloading them in `run.sh` — add another `WeightsSource` in `config.py` and read from the mount path:
+
+```python
+weights=[
+    WeightsSource(source=f"hf://{INIT_MODEL}", mount_location=INIT_MODEL_MOUNT, auth_secret_name="hf_access_token"),
+    WeightsSource(source="s3://my-bucket/asr-shards", mount_location="/app/data"),
+]
+```
+
+The AN4 smoke test stays as an in-container download because it's tiny and needs on-the-fly preprocessing (`.sph`→`.wav` + manifest building), which a read-only BDN mount can't do in place.
+
+## Checkpoints
+
+NeMo saves a `.nemo` tarball rather than the HF `config.json` + safetensors layout Baseten's checkpoint detector looks for, so `run.sh` moves the final `.nemo` into a `checkpoint-<N>/` directory under `$BT_CHECKPOINT_DIR` and writes a minimal `config.json` marker. This registers the artifact so it's backed up and appears in `truss train checkpoints list`. It's a NeMo checkpoint — download it and load it with NeMo (see below), not the standard vLLM/HF deploy path.
+
+## Evaluate
+
+Evaluate at your deployment latency on held-out data using NeMo's streaming inference script. The lowest-latency setting (`att_context_size=[56,0]`, 80ms chunk, 0ms look-ahead) is the most demanding, honest condition:
+
+```bash
+python ${NEMO_DIR}/examples/asr/asr_cache_aware_streaming/speech_to_text_cache_aware_streaming_infer.py \
+    model_path=<path-to-finetuned.nemo> \
+    dataset_manifest=<path-to-test_manifest.json> \
+    target_lang=auto \
+    att_context_size="[56,3]" \
+    decoder_type=rnnt \
+    pad_and_drop_preencoded=true \
+    batch_size=8 \
+    strip_lang_tags=false
+```
+
+The same checkpoint covers the whole latency/accuracy spectrum — pick the operating point at inference time via `att_context_size`, no retraining required.

examples/nemotron-asr-nemo/config.py

@@ -0,0 +1,59 @@
+from truss_train import definitions, WeightsSource
+from truss.base import truss_config
+
+# Last consolidated NeMo Framework container that bundles the ASR/TTS collection
+# with a pre-tested CUDA/numba/lightning/kaldialign stack (the newer 26.06 image
+# is Megatron/LLM-only). NVIDIA-NeMo/Speech recommends 26.02 as the current
+# stable runtime for speech models, so it sidesteps the source-build + numba-JIT
+# issues you hit on a plain PyTorch image.
+BASE_IMAGE = "nvcr.io/nvidia/nemo:26.02"
+
+# Base checkpoint. Delivered read-only via BDN before start commands run, so we
+# never pay for the ~2.4GB download on billed GPU time (BDN mirrors + caches it
+# across jobs). See https://docs.baseten.co/training/concepts/storage
+INIT_MODEL = "nvidia/nemotron-3.5-asr-streaming-0.6b"
+INIT_MODEL_MOUNT = f"/app/models/{INIT_MODEL}"
+
+training_runtime = definitions.Runtime(
+    start_commands=["/bin/sh -c 'chmod +x ./run.sh && ./run.sh'"],
+    environment_variables={
+        # Gated on Hugging Face - used by BDN auth and any in-container downloads.
+        "HF_TOKEN": definitions.SecretReference(name="hf_access_token"),
+        "HF_HUB_ENABLE_HF_TRANSFER": "true",
+        # run.sh reads this to locate the BDN-mounted checkpoint.
+        "INIT_MODEL_MOUNT": INIT_MODEL_MOUNT,
+    },
+    cache_config=definitions.CacheConfig(
+        enabled=True,
+    ),
+    checkpointing_config=definitions.CheckpointingConfig(
+        enabled=True,
+    ),
+)
+
+training_compute = definitions.Compute(
+    node_count=1,
+    accelerator=truss_config.AcceleratorSpec(
+        accelerator=truss_config.Accelerator.H100,
+        count=1,
+    ),
+)
+
+training_job = definitions.TrainingJob(
+    image=definitions.Image(base_image=BASE_IMAGE),
+    compute=training_compute,
+    runtime=training_runtime,
+    # Mount the base checkpoint via BDN instead of downloading it in run.sh.
+    # `auth_secret_name` authenticates against the gated HF repo.
+    weights=[
+        WeightsSource(
+            source=f"hf://{INIT_MODEL}",
+            mount_location=INIT_MODEL_MOUNT,
+            auth_secret_name="hf_access_token",
+        ),
+    ],
+)
+
+training_project = definitions.TrainingProject(
+    name="Nemotron-3.5-ASR-Streaming Finetuned", job=training_job
+)

examples/nemotron-asr-nemo/prepare_data.py

@@ -0,0 +1,121 @@
+"""Prepare a NeMo ASR manifest for fine-tuning Nemotron 3.5 ASR.
+
+This uses the small, public AN4 corpus as a stand-in for your own data so the
+example runs end-to-end quickly. Swap in your own audio + transcripts by
+producing JSON-lines manifests with the same schema:
+
+    {"audio_filepath": "/abs/path/clip.wav", "duration": 4.27,
+     "text": "reference transcript", "lang": "en-US", "target_lang": "en-US"}
+
+Two details matter most for Nemotron 3.5 ASR:
+  * Every clip carries a ``target_lang`` tag - this drives the model's
+    prompt-based language conditioning. Use a locale the model recognizes
+    (e.g. en-US, es-ES, el-GR, bg-BG).
+  * Transcripts should match the base model's text style: punctuated and
+    properly cased. AN4 is lowercase/unpunctuated, so expect it only as a
+    smoke test, not a quality benchmark.
+"""
+
+import argparse
+import glob
+import json
+import os
+import subprocess
+import tarfile
+import urllib.request
+
+AN4_URL = "https://dldata-public.s3.us-east-2.amazonaws.com/an4_sphere.tar.gz"
+
+
+def download_and_extract_an4(data_dir: str) -> str:
+    source_data_dir = os.path.join(data_dir, "an4")
+    if os.path.exists(source_data_dir):
+        print(f"AN4 already present at {source_data_dir}, skipping download.")
+        return source_data_dir
+
+    tar_path = os.path.join(data_dir, "an4_sphere.tar.gz")
+    print(f"Downloading AN4 from {AN4_URL} ...")
+    urllib.request.urlretrieve(AN4_URL, tar_path)
+    print("Extracting AN4 ...")
+    with tarfile.open(tar_path) as tar:
+        tar.extractall(path=data_dir)
+    return source_data_dir
+
+
+def build_manifest(transcripts_path, manifest_path, target_wavs_dir, target_lang):
+    import librosa
+
+    with open(transcripts_path, "r") as fin, open(manifest_path, "w") as fout:
+        for line in fin:
+            # Lines look like: <s> transcript </s> (fileID)
+            transcript = line[: line.find("(") - 1].lower()
+            transcript = transcript.replace("<s>", "").replace("</s>", "").strip()
+
+            file_id = line[line.find("(") + 1 : -2]
+            audio_path = os.path.join(target_wavs_dir, file_id + ".wav")
+            duration = librosa.core.get_duration(path=audio_path)
+
+            metadata = {
+                "audio_filepath": audio_path,
+                "duration": duration,
+                "text": transcript,
+                "lang": target_lang,
+                "target_lang": target_lang,
+            }
+            json.dump(metadata, fout)
+            fout.write("\n")
+    print(f"Wrote manifest: {manifest_path}")
+
+
+def main(args):
+    data_dir = args.data_dir
+    os.makedirs(data_dir, exist_ok=True)
+
+    source_data_dir = download_and_extract_an4(data_dir)
+    target_data_dir = os.path.join(data_dir, "an4_converted")
+    target_wavs_dir = os.path.join(target_data_dir, "wavs")
+    os.makedirs(target_wavs_dir, exist_ok=True)
+
+    # Convert the .sph source files to mono 16kHz .wav files.
+    sph_list = glob.glob(os.path.join(source_data_dir, "**/*.sph"), recursive=True)
+    print(f"Converting {len(sph_list)} .sph files to .wav ...")
+    for sph_path in sph_list:
+        wav_path = os.path.join(
+            target_wavs_dir,
+            os.path.splitext(os.path.basename(sph_path))[0] + ".wav",
+        )
+        subprocess.run(["sox", sph_path, "-r", "16000", "-c", "1", wav_path], check=True)
+
+    build_manifest(
+        os.path.join(source_data_dir, "etc/an4_train.transcription"),
+        os.path.join(target_data_dir, "train_manifest.json"),
+        target_wavs_dir,
+        args.target_lang,
+    )
+    build_manifest(
+        os.path.join(source_data_dir, "etc/an4_test.transcription"),
+        os.path.join(target_data_dir, "test_manifest.json"),
+        target_wavs_dir,
+        args.target_lang,
+    )
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Prepare AN4 manifests for NeMo ASR.")
+    parser.add_argument(
+        "--data_dir",
+        type=str,
+        default=os.environ.get("DATA_DIR", "./data"),
+        help="Directory to download and process the dataset into.",
+    )
+    parser.add_argument(
+        "--target_lang",
+        type=str,
+        default="en-US",
+        help="Language tag applied to every clip (drives prompt conditioning).",
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    main(parse_args())

examples/nemotron-asr-nemo/requirements.txt

@@ -0,0 +1,4 @@
+# The NeMo container already provides nemo_toolkit[asr] and its full ASR stack
+# (lightning, numba, kaldialign, librosa, ...). Only add the small extra used to
+# speed up the Hugging Face checkpoint download in the local-dev fallback.
+hf_transfer