Audio: add Qwen3-ASR adapter

nemo_curator/models/asr/base.py

@@ -42,8 +42,10 @@ class ASRResult:
             when ``skipped`` is true. Defaults to ``"empty_audio"`` in the stage.
         unsupported_language: Optional normalized language code used by the
             stage to annotate items excluded by its language allowlist.
-        extras: Adapter-specific diagnostics outside the canonical shape; the
-            stage never reads inside this dict.
+        extras: Adapter-specific, manifest-serializable diagnostics. When the
+            stage's ``extras_key`` is enabled, it writes a shallow copy of this
+            dictionary under that one nested output field without interpreting
+            individual keys.
     """
 
     text: str

nemo_curator/models/asr/qwen_asr.py

@@ -0,0 +1,222 @@
+# 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.
+
+"""Qwen3-ASR vLLM implementation of the shared ASR adapter.
+
+This uses the same ``Qwen3ASRModel.LLM`` construction and vLLM engine settings
+as the nkoluguri reference. ``ASRStage`` owns mono conversion and resampling;
+the adapter hands one prepared batch to one ``transcribe`` call and maps
+results back to ``ASRResult`` positions.
+"""
+
+from __future__ import annotations
+
+import gc
+from copy import deepcopy
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+import torch
+from huggingface_hub import snapshot_download
+from loguru import logger
+
+from nemo_curator.models.asr.base import ASRResult
+from nemo_curator.utils.vllm_utils import merge_vllm_kwargs
+
+_DEFAULT_QWEN3_ASR_MODEL = "Qwen/Qwen3-ASR-0.6B"
+
+# Qwen's audio processor needs >=200 samples for STFT padding. 1600 samples
+# (100 ms at 16 kHz) is a conservative floor that also matches the Qwen-Omni
+# preprocessing path.
+_MIN_SAMPLES = 1600
+
+
+def _qwen_asr_model_cls() -> Any:  # noqa: ANN401
+    try:
+        from qwen_asr import Qwen3ASRModel
+    except ImportError as exc:
+        msg = "QwenASRAdapter requires the audio_cuda12 and vllm extras: uv sync --extra audio_cuda12 --extra vllm"
+        raise ImportError(msg) from exc
+    return Qwen3ASRModel
+
+
+@dataclass
+class QwenASRAdapter:
+    """Run vLLM-backed Qwen3-ASR over Curator waveform items.
+
+    Every valid item in one adapter call goes to a single ``transcribe`` call,
+    so the caller's batch boundary is the model's batch boundary.
+    ``max_inference_batch_size`` is the library's own internal cap and is passed
+    through at construction.
+
+    ``revision`` is accepted to satisfy the shared adapter constructor and is
+    forwarded to both weight prefetch and the vLLM model loader.
+
+    ``vllm_kwargs`` exposes additional engine settings, following the existing
+    Qwen-Omni adapter convention. Adapter-owned settings cannot be overridden
+    through this mapping. Its default is empty, so normal construction exactly
+    matches the nkoluguri reference engine arguments.
+    """
+
+    model_id: str = _DEFAULT_QWEN3_ASR_MODEL
+    revision: str | None = None
+    gpu_memory_utilization: float = 0.7
+    max_new_tokens: int = 4096
+    max_inference_batch_size: int = 128
+    vllm_kwargs: dict[str, Any] = field(default_factory=dict)
+    _model: Any = field(default=None, init=False, repr=False)
+
+    def __post_init__(self) -> None:
+        if not self.model_id:
+            msg = "QwenASRAdapter.model_id must be non-empty"
+            raise ValueError(msg)
+        if not 0.0 < float(self.gpu_memory_utilization) <= 1.0:
+            msg = f"QwenASRAdapter.gpu_memory_utilization must be in (0, 1], got {self.gpu_memory_utilization}"
+            raise ValueError(msg)
+        if (
+            not isinstance(self.max_new_tokens, int)
+            or isinstance(self.max_new_tokens, bool)
+            or self.max_new_tokens <= 0
+        ):
+            msg = f"QwenASRAdapter.max_new_tokens must be a positive integer, got {self.max_new_tokens!r}"
+            raise ValueError(msg)
+        if (
+            not isinstance(self.max_inference_batch_size, int)
+            or isinstance(self.max_inference_batch_size, bool)
+            or self.max_inference_batch_size <= 0
+        ):
+            msg = (
+                "QwenASRAdapter.max_inference_batch_size must be a positive integer, "
+                f"got {self.max_inference_batch_size!r}"
+            )
+            raise ValueError(msg)
+        self.vllm_kwargs = deepcopy(dict(self.vllm_kwargs))
+
+    def _model_owned_vllm_kwargs(self) -> dict[str, Any]:
+        """Return the qwen-asr constructor arguments owned by this adapter."""
+        return {
+            "model": self.model_id,
+            "revision": self.revision,
+            "gpu_memory_utilization": self.gpu_memory_utilization,
+            "max_inference_batch_size": self.max_inference_batch_size,
+            "max_new_tokens": self.max_new_tokens,
+            "trust_remote_code": True,
+            "enforce_eager": True,
+            "enable_prefix_caching": True,
+            "prefix_caching_hash_algo": "xxhash",
+        }
+
+    @classmethod
+    def download_weights_on_node(cls, model_id: str, revision: str | None = None) -> None:
+        """Populate the local Hugging Face cache without allocating a GPU."""
+        snapshot_download(model_id, revision=revision)
+
+    def load_model(self, *, num_gpus: int) -> None:
+        """Load one worker-local Qwen3-ASR model through its vLLM backend."""
+        if self._model is not None:
+            return
+        if not isinstance(num_gpus, int) or isinstance(num_gpus, bool) or num_gpus != 1:
+            msg = f"QwenASRAdapter requires exactly one integer GPU, got {num_gpus!r}"
+            raise ValueError(msg)
+
+        logger.info(
+            "Loading QwenASRAdapter model={} gpu_mem={} max_new_tokens={} max_batch={}",
+            self.model_id,
+            self.gpu_memory_utilization,
+            self.max_new_tokens,
+            self.max_inference_batch_size,
+        )
+        model_kwargs = merge_vllm_kwargs(
+            self.vllm_kwargs,
+            self._model_owned_vllm_kwargs(),
+            owner_description="adapter-owned arguments",
+        )
+        if model_kwargs["revision"] is None:
+            del model_kwargs["revision"]
+        try:
+            self._model = _qwen_asr_model_cls().LLM(**model_kwargs)
+        except Exception:
+            self.unload_model()
+            raise
+        logger.info("QwenASRAdapter ready ({})", self.model_id)
+
+    def unload_model(self) -> None:
+        """Release the worker-local model and CUDA cache state."""
+        self._model = None
+        gc.collect()
+        try:
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+        except Exception as exc:  # noqa: BLE001
+            logger.debug("CUDA cache clear skipped: {}", exc)
+
+    @staticmethod
+    def _waveform(item: dict[str, Any]) -> np.ndarray:
+        waveform = np.asarray(item.get("waveform"), dtype=np.float32)
+        if waveform.ndim != 1:
+            msg = f"ASRStage must provide a mono 1-D waveform, got shape {waveform.shape}"
+            raise ValueError(msg)
+        return waveform
+
+    def transcribe_batch(self, items: list[dict[str, Any]]) -> list[ASRResult]:
+        """Transcribe one adapter call while preserving input order."""
+        if not items:
+            return []
+
+        valid_indices: list[int] = []
+        audio_inputs: list[tuple[np.ndarray, int]] = []
+        languages: list[str | None] = []
+        for index, item in enumerate(items):
+            waveform = self._waveform(item)
+            source_rate = int(item.get("sample_rate") or 0)
+            if waveform.size < _MIN_SAMPLES or source_rate <= 0:
+                continue
+            valid_indices.append(index)
+            audio_inputs.append((waveform, source_rate))
+            languages.append(item.get("language"))
+
+        results = [ASRResult(text="", skipped=True) for _ in items]
+        if not audio_inputs:
+            logger.warning(
+                "QwenASRAdapter: all {} items were shorter than {} samples or lacked a sample rate",
+                len(items),
+                _MIN_SAMPLES,
+            )
+            return results
+        if len(audio_inputs) < len(items):
+            logger.warning(
+                "QwenASRAdapter: skipping {}/{} items shorter than {} samples",
+                len(items) - len(audio_inputs),
+                len(items),
+                _MIN_SAMPLES,
+            )
+
+        outputs = self._model.transcribe(audio=audio_inputs, language=languages)
+
+        outputs = list(outputs or [])
+        if len(outputs) != len(valid_indices):
+            msg = f"Qwen3-ASR returned {len(outputs)} transcriptions for {len(valid_indices)} valid inputs"
+            raise RuntimeError(msg)
+
+        for index, output in zip(valid_indices, outputs, strict=True):
+            text = getattr(output, "text", output)
+            text = "" if text is None else str(text)
+            detected_language = getattr(output, "language", "") or ""
+            results[index] = ASRResult(
+                text=text,
+                skipped=not text.strip(),
+                extras={"detected_language": str(detected_language)} if detected_language else {},
+            )
+        return results

nemo_curator/models/asr/qwen_omni.py

@@ -27,7 +27,7 @@
 from loguru import logger
 
 from nemo_curator.models.asr.base import ASRResult
-from nemo_curator.utils.vllm_utils import create_vllm_llm
+from nemo_curator.utils.vllm_utils import create_vllm_llm, merge_vllm_kwargs
 
 if TYPE_CHECKING:
     import numpy as np
@@ -70,7 +70,6 @@ def _require_qwen_omni_stack(*, context: str) -> None:
 _QWEN_OMNI_SAMPLE_RATE = 16000
 _MIN_QWEN_AUDIO_SAMPLES = 1600
 _PROMPT_CONTENT_ORDERS = frozenset({"text_audio", "audio_text"})
-_RESERVED_VLLM_KWARGS = frozenset({"model", "revision", "tensor_parallel_size"})
 
 
 def _default_vllm_kwargs() -> dict[str, Any]:
@@ -154,10 +153,6 @@ def __post_init__(self) -> None:
             raise ValueError(msg)
         self.vllm_kwargs = deepcopy(dict(self.vllm_kwargs))
         self.sampling_kwargs = deepcopy(dict(self.sampling_kwargs))
-        reserved_vllm_kwargs = sorted(_RESERVED_VLLM_KWARGS.intersection(self.vllm_kwargs))
-        if reserved_vllm_kwargs:
-            msg = f"vllm_kwargs cannot override stage-owned arguments: {', '.join(reserved_vllm_kwargs)}"
-            raise ValueError(msg)
         if "max_tokens" in self.sampling_kwargs:
             msg = "sampling_kwargs cannot override adapter-owned max_tokens; use max_output_tokens"
             raise ValueError(msg)
@@ -166,6 +161,14 @@ def __post_init__(self) -> None:
         self._llm: Any = None
         self._sampling_params: Any = None
 
+    def _stage_owned_vllm_kwargs(self, *, num_gpus: int | None) -> dict[str, Any]:
+        """Return vLLM constructor arguments supplied by the stage contract."""
+        return {
+            "model": self.model_id,
+            "revision": self.revision,
+            "tensor_parallel_size": num_gpus,
+        }
+
     @staticmethod
     def _load_text(text: str | None, file_path: str | None) -> str | None:
         if file_path:
@@ -206,10 +209,14 @@ def load_model(self, *, num_gpus: int) -> None:
             + (f"  revision={self.revision}" if self.revision is not None else "")
         )
 
-        engine_kwargs = dict(self.vllm_kwargs)
-        engine_kwargs["tensor_parallel_size"] = num_gpus
-        if self.revision is not None:
-            engine_kwargs["revision"] = self.revision
+        engine_kwargs = merge_vllm_kwargs(
+            self.vllm_kwargs,
+            self._stage_owned_vllm_kwargs(num_gpus=num_gpus),
+            owner_description="stage-owned arguments",
+        )
+        del engine_kwargs["model"]
+        if engine_kwargs["revision"] is None:
+            del engine_kwargs["revision"]
 
         try:
             proc_kwargs: dict[str, Any] = {}

nemo_curator/stages/audio/inference/asr/stage.py

@@ -27,6 +27,7 @@
 
 import hydra.utils
 import numpy as np
+import soundfile
 import torch
 import torchaudio
 from loguru import logger
@@ -112,8 +113,9 @@ def _set_note(task_data: dict[str, Any], stage_name: str, value: str) -> None:
 class ASRStage(ProcessingStage[AudioTask, AudioTask]):
     """Audio speech-recognition stage with a pluggable adapter.
 
-    The stage writes only ``pred_text_key`` plus the optional control columns
-    ``_skipme`` and ``additional_notes``.
+    The stage writes ``pred_text_key`` and optional control columns ``_skipme``
+    and ``additional_notes``. When ``extras_key`` is configured, it also writes
+    non-empty adapter metadata as one nested dictionary under that key.
     """
 
     # Adapter selection.
@@ -132,6 +134,7 @@ class ASRStage(ProcessingStage[AudioTask, AudioTask]):
     default_language: str | None = None
     supported_language_codes: list[str] | None = None
     pred_text_key: str = "pred_text"
+    extras_key: str | None = None
 
     skip_if_output_exists: bool = False
 
@@ -149,6 +152,14 @@ def __post_init__(self) -> None:
         if self.pred_text_key in {_SKIP_ME_KEY, _NOTES_KEY}:
             msg = f"ASRStage.pred_text_key cannot use reserved control column {self.pred_text_key!r}"
             raise ValueError(msg)
+        if self.extras_key is not None:
+            self.extras_key = self.extras_key.strip()
+            if not self.extras_key:
+                msg = "ASRStage.extras_key must be non-empty or None"
+                raise ValueError(msg)
+            if self.extras_key in {self.pred_text_key, _SKIP_ME_KEY, _NOTES_KEY}:
+                msg = f"ASRStage.extras_key cannot collide with another output column: {self.extras_key!r}"
+                raise ValueError(msg)
         if int(self.batch_size) <= 0:
             msg = f"ASRStage.batch_size must be > 0, got {self.batch_size}"
             raise ValueError(msg)
@@ -236,7 +247,10 @@ def inputs(self) -> tuple[list[str], list[str]]:
         return [], [self.audio_filepath_key]
 
     def outputs(self) -> tuple[list[str], list[str]]:
-        return [], [self.pred_text_key, _SKIP_ME_KEY, _NOTES_KEY]
+        optional_outputs = [self.pred_text_key, _SKIP_ME_KEY, _NOTES_KEY]
+        if self.extras_key is not None:
+            optional_outputs.append(self.extras_key)
+        return [], optional_outputs
 
     def _resolve_language(self, task: AudioTask) -> str | None:
         code = self._resolve_language_code(task)
@@ -278,11 +292,15 @@ def _build_items(self, tasks: list[AudioTask]) -> list[dict[str, Any]]:
     def _load_audio(audio_filepath: str) -> tuple[np.ndarray, int]:
         """Open one resampled file inside the ASR worker.
 
-        ``ResampleAudioStage`` guarantees mono audio, so squeezing its channel
-        dimension matches the file-backed tagging pipeline contract.
+        ``soundfile`` avoids making PCM WAV decoding depend on TorchCodec's
+        optional CUDA/FFmpeg shared libraries. SoundFile returns multichannel
+        audio as sample-major, so transpose it to the channel-first shape used
+        by ``_prepare_waveform``.
         """
-        waveform, sample_rate = torchaudio.load(audio_filepath)
-        return waveform.squeeze(0).numpy(), sample_rate
+        waveform, sample_rate = soundfile.read(audio_filepath, dtype="float32")
+        if waveform.ndim == _CHANNEL_FIRST_DIMENSIONS:
+            waveform = waveform.T
+        return np.ascontiguousarray(waveform, dtype=np.float32), sample_rate
 
     def _prepare_waveform(self, waveform: object, sample_rate: object) -> np.ndarray:
         """Return contiguous mono float32 samples at ``target_sample_rate``."""
@@ -429,6 +447,11 @@ def assemble(
         skipped_count = 0
         for task, item, result in zip(tasks, items, results, strict=True):
             task.data[self.pred_text_key] = result.text
+            if self.extras_key is not None:
+                if result.extras:
+                    task.data[self.extras_key] = dict(result.extras)
+                else:
+                    task.data.pop(self.extras_key, None)
             unsupported_language = result.unsupported_language
             missing_language = self._supported_language_codes is not None and not item["language_code"]
             if missing_language: