diff --git a/nemo_curator/models/faster_whisper_asr.py b/nemo_curator/models/faster_whisper_asr.py new file mode 100644 index 0000000000..0bb8222d73 --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 60c46fbf0c..6cd584b3bf 100644 --- a/pyproject.toml +++ b/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 = [ diff --git a/tests/models/test_faster_whisper_asr.py b/tests/models/test_faster_whisper_asr.py new file mode 100644 index 0000000000..3ac8142de8 --- /dev/null +++ b/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)]) diff --git a/uv.lock b/uv.lock index d18db75994..3a996ea49e 100644 --- a/uv.lock +++ b/uv.lock @@ -5115,6 +5115,7 @@ all = [ { name = "data-designer" }, { name = "easydict" }, { name = "einops" }, + { name = "faster-whisper", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin'" }, { name = "fasttext" }, { name = "ftfy" }, { name = "google-cloud-translate" }, @@ -5179,6 +5180,7 @@ all = [ ] audio-common = [ { name = "accelerate" }, + { name = "faster-whisper", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin'" }, { name = "librosa" }, { name = "nemo-text-processing", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin'" }, { name = "nemo-toolkit", extra = ["asr"], marker = "platform_machine == 'x86_64' and sys_platform != 'darwin'" }, @@ -5196,6 +5198,7 @@ audio-common = [ ] audio-cpu = [ { name = "accelerate" }, + { name = "faster-whisper", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin'" }, { name = "librosa" }, { name = "nemo-text-processing", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin'" }, { name = "nemo-toolkit", extra = ["asr"], marker = "platform_machine == 'x86_64' and sys_platform != 'darwin'" }, @@ -5216,6 +5219,7 @@ audio-cuda12 = [ { name = "accelerate" }, { name = "cuda-python", version = "12.9.4", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "cuda-python", version = "12.9.5", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64') or sys_platform != 'linux'" }, + { name = "faster-whisper", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin'" }, { name = "gpustat" }, { name = "librosa" }, { name = "nemo-text-processing", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin'" }, @@ -5537,6 +5541,7 @@ requires-dist = [ { name = "datasets", specifier = ">=4.0.0" }, { name = "easydict", marker = "extra == 'video-cpu'" }, { name = "einops", marker = "extra == 'video-cpu'" }, + { name = "faster-whisper", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin' and extra == 'audio-common'", specifier = ">=1.1.0" }, { name = "fasttext", marker = "extra == 'text-cpu'", specifier = "==0.9.3" }, { name = "fsspec" }, { name = "ftfy", marker = "extra == 'text-cpu'" },