Audio: add Faster-Whisper ASR adapter

nemo_curator/models/faster_whisper_asr.py

@@ -0,0 +1,136 @@
+# 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.
+
+"""Faster-Whisper implementation of the shared ASR adapter."""
+
+from __future__ import annotations
+
+import gc
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+
+from nemo_curator.models.asr.base import ASRResult
+
+_TARGET_SAMPLE_RATE = 16_000
+_LANGUAGE_ALIASES = {
+    "fil": "tl",
+    "in": "id",
+    "iw": "he",
+    "ji": "yi",
+    "jv": "jw",
+    "nb": "no",
+}
+
+
+def _faster_whisper_stack() -> tuple[type, Any]:
+    try:
+        from faster_whisper import WhisperModel
+        from faster_whisper.utils import download_model
+    except ImportError as error:
+        msg = "Faster-Whisper ASR requires the 'faster-whisper' package"
+        raise ImportError(msg) from error
+    return WhisperModel, download_model
+
+
+def _whisper_model_class() -> type:
+    return _faster_whisper_stack()[0]
+
+
+def _download_whisper_model(model_id: str, revision: str | None) -> None:
+    _faster_whisper_stack()[1](model_id, revision=revision)
+
+
+@dataclass
+class FasterWhisperASR:
+    """Run Faster-Whisper on mono 16 kHz waveforms prepared by ``ASRStage``."""
+
+    model_id: str = "large-v3"
+    revision: str | None = None
+    compute_type: str = "float16"
+    beam_size: int = 5
+    vad_filter: bool = True
+    without_timestamps: bool = True
+    cpu_compute_type: str = "int8"
+    _model: Any | None = field(default=None, init=False, repr=False)
+
+    def __post_init__(self) -> None:
+        if not self.model_id:
+            msg = "FasterWhisperASR.model_id must be non-empty"
+            raise ValueError(msg)
+
+    @classmethod
+    def download_weights_on_node(cls, model_id: str, revision: str | None = None) -> None:
+        """Populate Faster-Whisper's node-local cache without allocating a GPU."""
+        _download_whisper_model(model_id, revision)
+
+    def load_model(self, *, num_gpus: int) -> None:
+        if self._model is not None:
+            return
+        if num_gpus < 0:
+            msg = "num_gpus must be non-negative"
+            raise ValueError(msg)
+
+        device = "cuda" if num_gpus else "cpu"
+        compute_type = self.compute_type if num_gpus else self.cpu_compute_type
+        self._model = _whisper_model_class()(
+            self.model_id,
+            device=device,
+            compute_type=compute_type,
+            revision=self.revision,
+        )
+
+    def unload_model(self) -> None:
+        self._model = None
+        gc.collect()
+        try:
+            import torch
+
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+        except ImportError:
+            pass
+
+    def transcribe_batch(self, items: list[dict[str, Any]]) -> list[ASRResult]:
+        if self._model is None:
+            msg = "FasterWhisperASR is not initialized; call load_model() first"
+            raise RuntimeError(msg)
+
+        results: list[ASRResult] = []
+        for item in items:
+            waveform = np.asarray(item.get("waveform"), dtype=np.float32)
+            if waveform.size == 0:
+                results.append(ASRResult(text="", skipped=True, skip_reason="empty_audio"))
+                continue
+            if waveform.ndim != 1:
+                msg = f"ASRStage must provide a mono 1-D waveform, got shape {waveform.shape}"
+                raise ValueError(msg)
+            sample_rate = int(item.get("sample_rate") or 0)
+            if sample_rate != _TARGET_SAMPLE_RATE:
+                msg = f"ASRStage must provide {_TARGET_SAMPLE_RATE} Hz audio; received {sample_rate} Hz"
+                raise ValueError(msg)
+
+            raw_language = str(item.get("language_code") or "").strip().lower()
+            language = _LANGUAGE_ALIASES.get(raw_language, raw_language) or None
+            segments, _ = self._model.transcribe(
+                np.ascontiguousarray(waveform),
+                language=language,
+                beam_size=self.beam_size,
+                vad_filter=self.vad_filter,
+                without_timestamps=self.without_timestamps,
+            )
+            text = " ".join(segment.text for segment in segments).strip()
+            results.append(ASRResult(text=text, extras={"language_code": language}))
+        return results

pyproject.toml

@@ -121,6 +121,7 @@ audio_common = [
     "nemo_text_processing; platform_machine == 'x86_64' and platform_system != 'Darwin'",  # pulls pynini, which ships x86_64-only wheels and fails to build on aarch64
     "opencc-python-reimplemented",
     "pyannote-audio>=4.0.0; platform_machine == 'x86_64' and platform_system != 'Darwin'",
+    "faster-whisper>=1.1.0; platform_machine == 'x86_64' and platform_system != 'Darwin'",
     "whisperx>=3.8.4; platform_machine == 'x86_64' and platform_system != 'Darwin'",
 ]
 audio_cpu = [

tests/models/test_faster_whisper_asr.py

@@ -0,0 +1,87 @@
+# 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.
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pytest
+
+from nemo_curator.models.asr.base import ASRAdapter
+from nemo_curator.models.faster_whisper_asr import FasterWhisperASR
+
+
+def _item(*, sample_rate: int = 16_000, language_code: str = "en") -> dict[str, object]:
+    return {
+        "waveform": np.zeros(160, dtype=np.float32),
+        "sample_rate": sample_rate,
+        "language_code": language_code,
+    }
+
+
+def test_adapter_conforms_to_shared_protocol() -> None:
+    assert isinstance(FasterWhisperASR(), ASRAdapter)
+
+
+def test_download_weights_uses_download_only() -> None:
+    with patch("nemo_curator.models.faster_whisper_asr._download_whisper_model") as download_model:
+        FasterWhisperASR.download_weights_on_node("large-v3", revision="abc123")
+
+    download_model.assert_called_once_with("large-v3", "abc123")
+
+
+def test_load_model_uses_stage_owned_gpu_count() -> None:
+    model_class = MagicMock()
+    adapter = FasterWhisperASR(compute_type="float16", revision="abc123")
+
+    with patch("nemo_curator.models.faster_whisper_asr._whisper_model_class", return_value=model_class):
+        adapter.load_model(num_gpus=1)
+
+    model_class.assert_called_once_with("large-v3", device="cuda", compute_type="float16", revision="abc123")
+
+
+def test_adapter_transcribes_waveform_and_maps_language_alias() -> None:
+    adapter = FasterWhisperASR()
+    adapter._model = MagicMock()
+    adapter._model.transcribe.return_value = (
+        iter([SimpleNamespace(text="hello"), SimpleNamespace(text="world")]),
+        object(),
+    )
+
+    results = adapter.transcribe_batch([_item(language_code="fil")])
+
+    assert [result.text for result in results] == ["hello world"]
+    assert results[0].extras["language_code"] == "tl"
+    assert adapter._model.transcribe.call_args.kwargs["language"] == "tl"
+
+
+def test_adapter_preserves_empty_positions() -> None:
+    adapter = FasterWhisperASR()
+    adapter._model = MagicMock()
+    item = _item()
+    item["waveform"] = np.zeros(0, dtype=np.float32)
+
+    results = adapter.transcribe_batch([item])
+
+    assert results[0].skipped is True
+    assert results[0].skip_reason == "empty_audio"
+    adapter._model.transcribe.assert_not_called()
+
+
+def test_adapter_requires_upstream_resampling() -> None:
+    adapter = FasterWhisperASR()
+    adapter._model = MagicMock()
+
+    with pytest.raises(ValueError, match="ASRStage must provide 16000 Hz"):
+        adapter.transcribe_batch([_item(sample_rate=8_000)])