Audio: add Qwen-Omni in-process ASR adapter#1967
Conversation
Greptile SummaryThis PR introduces a Qwen3-Omni in-process ASR adapter (
Confidence Score: 4/5Safe to merge with one targeted fix: tasks with no language code are silently emitted with an empty transcription and no skip marker when a language allowlist is configured. The adapter, stage, and executor plumbing are well-structured and the critical paths (partial GPU state cleanup, strict zip on vLLM outputs, lazy optional imports, revision forwarding) are all correct and tested. One confirmed defect: nemo_curator/stages/audio/inference/asr/stage.py — specifically the Important Files Changed
Sequence DiagramsequenceDiagram
participant YAML as pipeline.yaml
participant Stage as ASRStage
participant Adapter as QwenOmniASRAdapter
participant vLLM as vLLM Engine
YAML->>Stage: process_batch(tasks)
Stage->>Stage: _partition_inference_tasks()
Stage->>Stage: validate_input(task)
Stage->>Stage: _build_items() → items[]
Stage->>Stage: run_inference(items)
Note over Stage: Language allowlist filtering
Stage->>Stage: _load_audio(filepath) per item
Stage->>Adapter: transcribe_batch(adapter_items)
Adapter->>Adapter: _prepare_batch() → Turn 1 inputs
Adapter->>vLLM: _generate(t1_inputs)
vLLM-->>Adapter: "RequestOutput[] (strict=True zip)"
Adapter->>Adapter: mark empty outputs as skipped
opt followup_prompt set
Adapter->>Adapter: _prepare_turn2_batch() → Turn 2 inputs
Adapter->>vLLM: _generate(t2_inputs)
vLLM-->>Adapter: RequestOutput[]
end
Adapter-->>Stage: ASRResult[] (text, secondary_text, skipped)
Stage->>Stage: assemble() → write pred_text_key, _skipme, notes
Stage-->>YAML: tasks[] (mutated in-place)
Reviews (52): Last reviewed commit: "Restore main MCP lock version" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| with open(out_path, "a", encoding="utf-8") as f: | ||
| f.write(json.dumps(task.data, ensure_ascii=False) + "\n") |
There was a problem hiding this comment.
json.dumps will crash when keep_waveform=True
task.data can contain a numpy ndarray (keyed by waveform_key) when the upstream InferenceQwenOmniStage is configured with keep_waveform: true. json.dumps has no numpy serializer and raises TypeError: Object of type ndarray is not JSON serializable, crashing the entire shard write for all tasks in that batch. Either strip the waveform here before serialising, or document that keep_waveform must be false when this writer is used (and add a validation guard in setup or __post_init__).
There was a problem hiding this comment.
Thanks for flagging. This is already addressed on the current head (936fa17f):
_manifest_datafirst drops keys named indrop_manifest_keys(defaults to("waveform",)) so the configuredwaveform_keynever reaches serialisation, regardless ofkeep_waveform.- Anything else with
.shapeand.dtype(numpy ndarrays, torch tensors, etc.) is dropped via a duck-typing guard beforejson.dumpsis called. - The remaining
json.dumpscall is wrapped intry/except TypeError, so a previously-unseen non-serialisable value raises a focusedTypeErrorwith the offending key instead of crashing the shard.
Citation: nemo_curator/stages/audio/io/sharded_manifest_writer.py:96-111. Resolving as already-fixed.
| cfg = value if OmegaConf.is_config(value) else OmegaConf.create(value) | ||
| if "_target_" in cfg: | ||
| return hydra.utils.instantiate(cfg) | ||
| raw = OmegaConf.to_container(cfg, resolve=True) | ||
| return Resources(**raw) | ||
| msg = f"Invalid resources override: {value!r}" |
There was a problem hiding this comment.
logger.info(f"Hydra config:\n{OmegaConf.to_yaml(cfg)}") prints the full resolved config, including any hf_token passed as a Hydra override. The credential ends up in every log sink (stdout, files, observability stacks) in plaintext. Consider redacting the hf_token field before logging — for example, by building a sanitised copy of the config dict — or logging only a subset of non-sensitive keys.
There was a problem hiding this comment.
Thanks for flagging. This was addressed in commit 936fa17f:
- The startup log call no longer uses raw
OmegaConf.to_yaml(cfg). It now uses_safe_config_yaml(cfg)(tutorialmain.py:346), which builds a redacted copy of the config before rendering to YAML. _redact_secret_valueswalks the config recursively (main.py:170-179) and replaces values of any key matching_SECRET_KEY_NAMES(which explicitly includeshf_token,password,secret_key,token,credentials, …) or_SECRET_KEY_PARTSsubstrings with<redacted>.- Trailing-suffix matching also catches any custom secret named
*_token,*_secret, or*_password.
Note: the regression test that covered this behaviour (tests/stages/audio/inference/test_qwen_omni_tutorial.py::test_safe_config_yaml_redacts_hf_token_but_keeps_token_counts) was removed in this revision per @sarahyurick's "we shouldn't need pytests for tutorials" comment. The helper code itself is unchanged.
Citations: tutorials/audio/qwen_omni_inprocess/main.py:60-72, :160-184, :346. Resolving as already-fixed.
| completed.add(shard_key) | ||
| return completed | ||
|
|
||
| @staticmethod |
There was a problem hiding this comment.
None return from extractfile
tarfile.TarFile.extractfile returns None for members that are not regular files (hard links, directory entries embedded in some tar formats). The preceding tar_info.isfile() guard does not cover all cases where extractfile may return None — calling .read() on a None result would raise AttributeError. Add a None check before .read().
|
|
||
| def _get_prompt_text(self, language: str | None) -> str: | ||
| """Return the EN-specific prompt for English, otherwise the default prompt.""" | ||
| if language and language == "English" and self.en_prompt_text: |
There was a problem hiding this comment.
language and is redundant: if language == "English" evaluates to True the string is already truthy, so the extra truthiness check is dead code and can mislead readers into thinking an empty-string case needs guarding here.
| if language and language == "English" and self.en_prompt_text: | |
| if language == "English" and self.en_prompt_text: |
sarahyurick
left a comment
sarahyurick
left a comment
There was a problem hiding this comment.
I did an initial pass to start familiarizing myself with the PR for now. Left some minor comments.
| model_id: str = _QWEN3_OMNI_MODEL_ID, | ||
| prompt_text: str = "Transcribe the audio.", | ||
| en_prompt_text: str | None = None, | ||
| followup_prompt: str = "Now listen to the audio again and add any false starts, filler words and preserve colloquial words (like lemme, gonna, wanna, etc) as is spoken in the audio.", |
There was a problem hiding this comment.
Instead of having the long string here, let's create a script variable called _FOLLOWUP_PROMPT or similar.
| prefix_caching_hash_algo="xxhash", | ||
| ) | ||
|
|
||
| from transformers import Qwen3OmniMoeProcessor |
| self._sampling_params = None | ||
| gc.collect() | ||
| try: | ||
| import torch |
| def _prepare_turn2_single( | ||
| self, waveform_16k: np.ndarray, pred_text: str, language: str | None = None, | ||
| ) -> dict[str, Any] | None: | ||
| from qwen_omni_utils import process_mm_info |
| """ | ||
|
|
||
| name: str = "sharded_manifest_writer" | ||
| output_dir: str = "" |
There was a problem hiding this comment.
| output_dir: str = "" | |
| output_dir: str |
instead of checking for it in the post init.
| from nemo_curator.stages.audio.inference.qwen_omni import InferenceQwenOmniStage | ||
| from nemo_curator.stages.audio.io.nemo_tarred_reader import NemoTarredAudioReader, NemoTarShardReaderStage | ||
| from nemo_curator.stages.audio.io.sharded_manifest_writer import ShardedManifestWriterStage | ||
| from tutorials.audio.qwen_omni_inprocess.main import ( |
There was a problem hiding this comment.
We shouldn't need pytests for tutorials.
| and prefetches common HuggingFace model attributes without hardcoding a | ||
| full Granary v2 post-processing graph in this entry point. | ||
| """ | ||
| from huggingface_hub import hf_hub_download, snapshot_download |
| If you do not have `uv`, use pip: | ||
|
|
||
| ```bash | ||
| pip install -e ".[audio_cuda12]" | ||
| ``` |
There was a problem hiding this comment.
We should only encourage uv and not pip.
| from qwen_omni_utils import process_mm_info | ||
| from transformers import Qwen3OmniMoeProcessor |
There was a problem hiding this comment.
audio_cuda12 installs
qwen_omni_utils is only shipped with the audio_cuda12 optional extra, but from qwen_omni_utils import process_mm_info and from transformers import Qwen3OmniMoeProcessor are both imported at module level, outside any guard. On a standard Curator installation (including the Mac/ARM case the PR description explicitly claims will work), import nemo_curator.models.qwen_omni fails immediately with ImportError: No module named 'qwen_omni_utils'. The vllm import immediately below correctly uses try/except ImportError + VLLM_AVAILABLE; these two imports need the same treatment — either fold them into that same try block, or defer them into setup() where VLLM_AVAILABLE is already checked.
| """ | ||
|
|
||
| name: str = "nemo_tar_shard_discovery" | ||
| yaml_path: str = "" |
There was a problem hiding this comment.
| yaml_path: str = "" | |
| yaml_path: str |
This can be empty instead of checking it in the post init.
| """ | ||
|
|
||
| name: str = "nemo_tarred_audio_reader" | ||
| yaml_path: str = "" |
There was a problem hiding this comment.
| yaml_path: str = "" | |
| yaml_path: str |
Same comment as above.
NemoTarShardDiscoveryStage and NemoTarredAudioReader previously declared yaml_path as `yaml_path: str = ""` and validated it via a manual empty- string check in __post_init__. Per reviewer feedback, this is cleaner as a required dataclass field: __init__ enforces the requirement directly, removing the sentinel and the manual check. Reordered yaml_path before defaulted fields in both dataclasses (Python requires non-default fields first). NemoTarShardDiscoveryStage's __post_init__ existed only for the empty-string check and is removed; NemoTarredAudioReader's __post_init__ keeps super().__init__() and self._stages wiring, only the check is removed. All three construction sites (the test, hydra _target_ resolution, and NemoTarredAudioReader's own __post_init__) already use keyword arguments, so the required-field reorder is transparent. The observable difference for a caller that forgets yaml_path is now TypeError from __init__ (before any object exists) instead of ValueError from __post_init__. Addresses: - NVIDIA-NeMo#1967 (comment) - NVIDIA-NeMo#1967 (comment) Signed-off-by: aaftaabv@gmail.com <aaftaabv@gmail.com>
…1.0 pins The audio_cuda12 PR previously added three [tool.uv] override-dependencies entries (huggingface-hub==0.36.0, transformers==4.57.6, accelerate==1.12.0) described as overriding qwen-asr's incompatible declared pins. Audit of the resolved graph showed that rationale was inaccurate: - qwen-asr 0.0.6 declares transformers==4.57.6 / accelerate==1.12.0 natively in its requires-dist, so no override is needed for those. - The remaining transformers/accelerate constraints across the graph (nemo-toolkit ~=4.57.0, vllm >=4.56.0,<5, pyannote-audio >=4.48.3, whisperx >=4.48.0) all naturally permit 4.57.6 / 1.12.0. - The hf-hub conflict (data-designer-engine >=1.0.1 vs whisperx <1.0.0) was already resolved on main by hf-hub>=0.34,<1.0; this PR's tightening to ==0.36.0 was redundant. All three entries removed; the pre-existing hf-hub override restored to its original range. Verified via `uv lock` that the resolved versions of transformers, accelerate, huggingface-hub, qwen-asr, nagisa, soynlp, vllm, torch, torchaudio, torchvision, torchcodec, nixl-cu12, xgrammar, fasttext, data-designer-engine, pyannote-audio, whisperx, nemo-toolkit, qwen-omni-utils, fsspec, numpy, protobuf, setuptools are byte-identical before and after (591 packages resolved). This PR now adds zero new override-dependencies entries, eliminating the cross-modality blast radius that was flagged. Also softened the remaining hard `==` pins inside audio_common / audio_cuda12: - nagisa ==0.2.11 -> >=0.2.11,<0.3 (8-year-old stable lib; qwen-asr still transitively pins 0.2.11) - qwen-asr ==0.0.6 -> >=0.0.6,<0.1 (only 0.0.6 exists today, but a future security-patch 0.0.7 auto-flows in if/when published; <0.1 caps speculative jumps) - fasttext ==0.9.3 -> >=0.9.3,<0.10 (10-year-old maintenance-mode lib; releases every 2-4 years) The 7-line preamble comment above qwen-asr was trimmed to one line in the same edit (reviewer asked for less verbose). uv.lock is regenerated; all resolved package versions remain identical to the pre-edit state. Addresses: - NVIDIA-NeMo#1967 (comment) - NVIDIA-NeMo#1967 (comment) - NVIDIA-NeMo#1967 (comment) - NVIDIA-NeMo#1967 (comment) - NVIDIA-NeMo#1967 (comment) Signed-off-by: aaftaabv@gmail.com <aaftaabv@gmail.com>
NemoTarShardDiscoveryStage and NemoTarredAudioReader previously declared yaml_path as `yaml_path: str = ""` and validated it via a manual empty- string check in __post_init__. Per reviewer feedback, this is cleaner as a required dataclass field: __init__ enforces the requirement directly, removing the sentinel and the manual check. Reordered yaml_path before defaulted fields in both dataclasses (Python requires non-default fields first). NemoTarShardDiscoveryStage's __post_init__ existed only for the empty-string check and is removed; NemoTarredAudioReader's __post_init__ keeps super().__init__() and self._stages wiring, only the check is removed. All three construction sites (the test, hydra _target_ resolution, and NemoTarredAudioReader's own __post_init__) already use keyword arguments, so the required-field reorder is transparent. The observable difference for a caller that forgets yaml_path is now TypeError from __init__ (before any object exists) instead of ValueError from __post_init__. Addresses: - NVIDIA-NeMo#1967 (comment) - NVIDIA-NeMo#1967 (comment) Signed-off-by: Aaftab V <aaftabv@nvidia.com>
…1.0 pins The audio_cuda12 PR previously added three [tool.uv] override-dependencies entries (huggingface-hub==0.36.0, transformers==4.57.6, accelerate==1.12.0) described as overriding qwen-asr's incompatible declared pins. Audit of the resolved graph showed that rationale was inaccurate: - qwen-asr 0.0.6 declares transformers==4.57.6 / accelerate==1.12.0 natively in its requires-dist, so no override is needed for those. - The remaining transformers/accelerate constraints across the graph (nemo-toolkit ~=4.57.0, vllm >=4.56.0,<5, pyannote-audio >=4.48.3, whisperx >=4.48.0) all naturally permit 4.57.6 / 1.12.0. - The hf-hub conflict (data-designer-engine >=1.0.1 vs whisperx <1.0.0) was already resolved on main by hf-hub>=0.34,<1.0; this PR's tightening to ==0.36.0 was redundant. All three entries removed; the pre-existing hf-hub override restored to its original range. Verified via `uv lock` that the resolved versions of transformers, accelerate, huggingface-hub, qwen-asr, nagisa, soynlp, vllm, torch, torchaudio, torchvision, torchcodec, nixl-cu12, xgrammar, fasttext, data-designer-engine, pyannote-audio, whisperx, nemo-toolkit, qwen-omni-utils, fsspec, numpy, protobuf, setuptools are byte-identical before and after (591 packages resolved). This PR now adds zero new override-dependencies entries, eliminating the cross-modality blast radius that was flagged. Also softened the remaining hard `==` pins inside audio_common / audio_cuda12: - nagisa ==0.2.11 -> >=0.2.11,<0.3 (8-year-old stable lib; qwen-asr still transitively pins 0.2.11) - qwen-asr ==0.0.6 -> >=0.0.6,<0.1 (only 0.0.6 exists today, but a future security-patch 0.0.7 auto-flows in if/when published; <0.1 caps speculative jumps) - fasttext ==0.9.3 -> >=0.9.3,<0.10 (10-year-old maintenance-mode lib; releases every 2-4 years) The 7-line preamble comment above qwen-asr was trimmed to one line in the same edit (reviewer asked for less verbose). uv.lock is regenerated; all resolved package versions remain identical to the pre-edit state. Addresses: - NVIDIA-NeMo#1967 (comment) - NVIDIA-NeMo#1967 (comment) - NVIDIA-NeMo#1967 (comment) - NVIDIA-NeMo#1967 (comment) - NVIDIA-NeMo#1967 (comment) Signed-off-by: Aaftab V <aaftabv@nvidia.com>
…scaling Per Praateek's review feedback on tutorials/audio/qwen_omni_inprocess/ qwen_omni_inprocess.yaml (NVIDIA-NeMo#1967), controlled-throughput benchmarking is not a tutorial concern. Remove the explicit worker-count overrides (reader_num_workers, reader_num_workers_per_node, omni_num_workers) from the tutorial config and the NemoTarredAudioReader / NemoTarShardReaderStage / InferenceQwenOmniStage classes. The ProcessingStage base class defaults (num_workers() -> None, xenna_stage_spec() -> {}) let Xenna's autoscaler size each stage from the cluster resources, which is the documented happy-path behaviour of the audio pipelines. Changes: - tutorials/audio/qwen_omni_inprocess/qwen_omni_inprocess.yaml: remove the MULTI-NODE SCALING block (reader_num_workers, reader_num_workers_per_node, omni_num_workers) and the three Hydra interpolations that fed them into the reader and qwen-omni stages. - nemo_curator/stages/audio/io/nemo_tarred_reader.py: drop reader_num_workers / reader_num_workers_per_node fields on the composite NemoTarredAudioReader; drop num_workers_override / num_workers_per_node fields, their __post_init__ validation, the num_workers() override, and the xenna_stage_spec() override on NemoTarShardReaderStage. - nemo_curator/stages/audio/inference/qwen_omni.py: drop the num_workers_override field plus the num_workers() and xenna_stage_spec() overrides on InferenceQwenOmniStage. Remove the now-unused `Any` import. - tutorials/audio/qwen_omni_inprocess/README.md: drop the reader_num_workers / reader_num_workers_per_node / omni_num_workers rows from the configuration-knob table. - tests/stages/audio/inference/test_qwen_omni.py: remove test_worker_override_specs (no longer applicable). Addresses: - NVIDIA-NeMo#1967 (comment) Signed-off-by: Aaftab V <aaftabv@nvidia.com>
05264b0 to
2c42888
Compare
NemoTarShardDiscoveryStage and NemoTarredAudioReader previously declared yaml_path as `yaml_path: str = ""` and validated it via a manual empty- string check in __post_init__. Per reviewer feedback, this is cleaner as a required dataclass field: __init__ enforces the requirement directly, removing the sentinel and the manual check. Reordered yaml_path before defaulted fields in both dataclasses (Python requires non-default fields first). NemoTarShardDiscoveryStage's __post_init__ existed only for the empty-string check and is removed; NemoTarredAudioReader's __post_init__ keeps super().__init__() and self._stages wiring, only the check is removed. All three construction sites (the test, hydra _target_ resolution, and NemoTarredAudioReader's own __post_init__) already use keyword arguments, so the required-field reorder is transparent. The observable difference for a caller that forgets yaml_path is now TypeError from __init__ (before any object exists) instead of ValueError from __post_init__. Addresses: - NVIDIA-NeMo#1967 (comment) - NVIDIA-NeMo#1967 (comment) Signed-off-by: Aaftab V <aaftabv@nvidia.com>
…1.0 pins The audio_cuda12 PR previously added three [tool.uv] override-dependencies entries (huggingface-hub==0.36.0, transformers==4.57.6, accelerate==1.12.0) described as overriding qwen-asr's incompatible declared pins. Audit of the resolved graph showed that rationale was inaccurate: - qwen-asr 0.0.6 declares transformers==4.57.6 / accelerate==1.12.0 natively in its requires-dist, so no override is needed for those. - The remaining transformers/accelerate constraints across the graph (nemo-toolkit ~=4.57.0, vllm >=4.56.0,<5, pyannote-audio >=4.48.3, whisperx >=4.48.0) all naturally permit 4.57.6 / 1.12.0. - The hf-hub conflict (data-designer-engine >=1.0.1 vs whisperx <1.0.0) was already resolved on main by hf-hub>=0.34,<1.0; this PR's tightening to ==0.36.0 was redundant. All three entries removed; the pre-existing hf-hub override restored to its original range. Verified via `uv lock` that the resolved versions of transformers, accelerate, huggingface-hub, qwen-asr, nagisa, soynlp, vllm, torch, torchaudio, torchvision, torchcodec, nixl-cu12, xgrammar, fasttext, data-designer-engine, pyannote-audio, whisperx, nemo-toolkit, qwen-omni-utils, fsspec, numpy, protobuf, setuptools are byte-identical before and after (591 packages resolved). This PR now adds zero new override-dependencies entries, eliminating the cross-modality blast radius that was flagged. Also softened the remaining hard `==` pins inside audio_common / audio_cuda12: - nagisa ==0.2.11 -> >=0.2.11,<0.3 (8-year-old stable lib; qwen-asr still transitively pins 0.2.11) - qwen-asr ==0.0.6 -> >=0.0.6,<0.1 (only 0.0.6 exists today, but a future security-patch 0.0.7 auto-flows in if/when published; <0.1 caps speculative jumps) - fasttext ==0.9.3 -> >=0.9.3,<0.10 (10-year-old maintenance-mode lib; releases every 2-4 years) The 7-line preamble comment above qwen-asr was trimmed to one line in the same edit (reviewer asked for less verbose). uv.lock is regenerated; all resolved package versions remain identical to the pre-edit state. Addresses: - NVIDIA-NeMo#1967 (comment) - NVIDIA-NeMo#1967 (comment) - NVIDIA-NeMo#1967 (comment) - NVIDIA-NeMo#1967 (comment) - NVIDIA-NeMo#1967 (comment) Signed-off-by: Aaftab V <aaftabv@nvidia.com>
…scaling Per Praateek's review feedback on tutorials/audio/qwen_omni_inprocess/ qwen_omni_inprocess.yaml (NVIDIA-NeMo#1967), controlled-throughput benchmarking is not a tutorial concern. Remove the explicit worker-count overrides (reader_num_workers, reader_num_workers_per_node, omni_num_workers) from the tutorial config and the NemoTarredAudioReader / NemoTarShardReaderStage / InferenceQwenOmniStage classes. The ProcessingStage base class defaults (num_workers() -> None, xenna_stage_spec() -> {}) let Xenna's autoscaler size each stage from the cluster resources, which is the documented happy-path behaviour of the audio pipelines. Changes: - tutorials/audio/qwen_omni_inprocess/qwen_omni_inprocess.yaml: remove the MULTI-NODE SCALING block (reader_num_workers, reader_num_workers_per_node, omni_num_workers) and the three Hydra interpolations that fed them into the reader and qwen-omni stages. - nemo_curator/stages/audio/io/nemo_tarred_reader.py: drop reader_num_workers / reader_num_workers_per_node fields on the composite NemoTarredAudioReader; drop num_workers_override / num_workers_per_node fields, their __post_init__ validation, the num_workers() override, and the xenna_stage_spec() override on NemoTarShardReaderStage. - nemo_curator/stages/audio/inference/qwen_omni.py: drop the num_workers_override field plus the num_workers() and xenna_stage_spec() overrides on InferenceQwenOmniStage. Remove the now-unused `Any` import. - tutorials/audio/qwen_omni_inprocess/README.md: drop the reader_num_workers / reader_num_workers_per_node / omni_num_workers rows from the configuration-knob table. - tests/stages/audio/inference/test_qwen_omni.py: remove test_worker_override_specs (no longer applicable). Addresses: - NVIDIA-NeMo#1967 (comment) Signed-off-by: Aaftab V <aaftabv@nvidia.com>
2c42888 to
1bc03a7
Compare
| entry_count += 1 | ||
| entry = json.loads(stripped) | ||
| audio_path = str(entry[self.filepath_key]) | ||
| for lookup_key in self._manifest_lookup_keys(audio_path): |
There was a problem hiding this comment.
KeyError crashes entire shard manifest read
entry[self.filepath_key] raises an unhandled KeyError if any manifest line is missing the configured filepath_key (default "audio_filepath"). Because _read_manifest has no per-entry try-except, a single malformed line aborts parsing of the entire manifest — the function exits via exception before returning entries, so all audio in that shard's tar becomes inaccessible and the shard is reported as failed to the executor. Wrapping the body of the loop in a try-except (or using entry.get(self.filepath_key) with a skip on None) would confine the failure to the bad entry and let the rest of the manifest be processed.
Bring the SDP-V2 design doc's stage->adapter split (see "Replaceability: Stage -> Adapter" in the design doc) into PR NVIDIA-NeMo#1967 for the Qwen-Omni first-pass ASR slot. Pure refactor: no behaviour change, no metric-key change, no runtime-performance change expected. Layout changes -------------- New: nemo_curator/adapters/__init__.py nemo_curator/adapters/asr/__init__.py nemo_curator/adapters/asr/base.py - ASRAdapter Protocol + ASRResult dataclass nemo_curator/adapters/asr/qwen_omni.py - QwenOmniASRAdapter (vLLM setup, two-turn generate; core inference logic moved verbatim from the deleted nemo_curator/models/qwen_omni.py) nemo_curator/stages/audio/inference/asr/stage.py - generic ASRStage matching the design doc shape (Tier-1 adapter_target + model_id + stage knobs; Tier-2 opaque adapter_kwargs; adapter class resolved at setup() via hydra.utils.get_class, matching Curator framework convention in nemo_curator/config/run.py). Deleted: nemo_curator/stages/audio/inference/qwen_omni.py - replaced by ASRStage. nemo_curator/models/qwen_omni.py - body moved into QwenOmniASRAdapter. Other audio inference stages (asr_nemo.py, sortformer, pyannote, whisperx_vad, vad_segmentation, nemo_asr_align) are intentionally left unchanged - the doc's broader stage-adapter rollout is a follow-up PR series, not in PR NVIDIA-NeMo#1967's scope. YAML / tutorial --------------- tutorials/audio/qwen_omni_inprocess/qwen_omni_inprocess.yaml rewritten to the design-doc shape: _target_ now ASRStage, with new top-level adapter_target field (Tier-1 swap line) and an adapter_kwargs block (Tier-2 opaque) carrying the Qwen-Omni-specific knobs. The stage's name is explicitly overridden to QwenOmni_inference so perf_summary_merged.json's stages[QwenOmni_inference] key stays stable across the refactor, allowing direct apples-to-apples comparison against the existing 4e8021a / ac6598e / d8b1d751 / 8cd168ea baselines. Tests ----- tests/stages/audio/inference/test_qwen_omni.py rewritten end-to-end: 21 tests covering both ASRStage (process_batch contract, language resolution, skip handling, A5-alias-dedup metric filter preservation, adapter_target validation, outputs() with/without secondary_text_key, setup() via mocked hydra.utils.get_class, setup_on_node prefetch behaviour) and QwenOmniASRAdapter (transcribe_batch packaging contract, single-turn drops secondary_text, empty-vLLM-output helpers). Behaviour preservation ---------------------- Verified equivalent to pre-refactor for every observable surface: - task.data output keys (qwen3_prediction_s1, qwen3_prediction_s2, _skip_me, waveform drop on keep_waveform=False) - vLLM LLM(...) ctor args + generate() call sequence (moved verbatim) - Turn-1 + Turn-2 disfluency flow (gated by followup_prompt) - setup_on_node snapshot_download weight prefetch - Metric key shape including A5-fix alias-dedup filter - All P0 / A2 / A4 / A5 anomaly fixes from earlier commits Throughput-neutrality will be Kratos-validated next. Signed-off-by: Aaftab V <aaftabv@nvidia.com>
| """ | ||
|
|
||
| from nemo_curator.adapters.asr.base import ASRAdapter, ASRResult | ||
| from nemo_curator.adapters.asr.qwen_omni import QwenOmniASRAdapter |
There was a problem hiding this comment.
audio_cuda12 package imports
nemo_curator/adapters/asr/__init__.py eagerly imports QwenOmniASRAdapter from qwen_omni.py, which in turn has bare top-level imports of qwen_omni_utils and Qwen3OmniMoeProcessor outside any try/except guard. Any code that does from nemo_curator.adapters.asr import ASRAdapter — including test-only or CPU-only Curator consumers — will immediately get ImportError: No module named 'qwen_omni_utils' on a standard install.
The __init__.py should only re-export the protocol and result dataclass from base.py. The concrete adapter is resolved lazily at runtime via hydra.utils.get_class inside ASRStage.setup; there is no need to make it a package-level public symbol.
There was a problem hiding this comment.
+1 should we be doing lazy imports? For example: https://github.com/NVIDIA-NeMo/Curator/blob/main/nemo_curator/stages/text/classifiers/__init__.py
There was a problem hiding this comment.
+1 — implemented.
adapters/asr/__init__.py: ASRAdapter / ASRResult stay eager from base.py; QwenOmniASRAdapter is lazy-loaded via PEP 562 __getattr__ (same pattern as nemo_curator/stages/text/classifiers/__init__.py). Importing nemo_curator.adapters.asr no longer pulls in qwen_omni at package init.
qwen_omni.py: qwen_omni_utils and Qwen3OmniMoeProcessor are now behind try/except ImportError (like vllm); setup() raises a single clear error if audio_cuda12 deps are missing.
Tests: tests/adapters/asr/test_package_lazy_import.py asserts the package import does not load qwen_omni until QwenOmniASRAdapter is accessed.
Hydra YAML still uses nemo_curator.adapters.asr.qwen_omni.QwenOmniASRAdapter; ASRStage still imports only from base.
mohammadaaftabv
changed the title
| self._llm = LLM( | ||
| model=self.model_id, | ||
| trust_remote_code=True, | ||
| gpu_memory_utilization=self.gpu_memory_utilization, | ||
| tensor_parallel_size=tp_size, | ||
| limit_mm_per_prompt={"image": 1, "video": 1, "audio": int(self.limit_mm_per_prompt_audio)}, | ||
| max_num_seqs=self.max_num_seqs, | ||
| max_model_len=self.max_model_len, | ||
| seed=int(self.seed), | ||
| enable_prefix_caching=bool(self.enable_prefix_caching), | ||
| prefix_caching_hash_algo=str(self.prefix_caching_hash_algo), | ||
| ) | ||
|
|
||
| self._processor = Qwen3OmniMoeProcessor.from_pretrained(self.model_id) |
There was a problem hiding this comment.
revision field is documented as a Tier-1 model-version pin and is correctly forwarded in prefetch_weights via snapshot_download. However, it is silently dropped when building the actual inference objects in setup(): LLM is constructed without revision=, and Qwen3OmniMoeProcessor.from_pretrained is called without revision=. This means a deployment that sets revision to pin a specific checkpoint will download the intended revision but load whatever HuggingFace Hub resolves by default at inference time — the two could diverge if the remote tag moves or the local cache contains multiple revisions.
| self._llm = LLM( | |
| model=self.model_id, | |
| trust_remote_code=True, | |
| gpu_memory_utilization=self.gpu_memory_utilization, | |
| tensor_parallel_size=tp_size, | |
| limit_mm_per_prompt={"image": 1, "video": 1, "audio": int(self.limit_mm_per_prompt_audio)}, | |
| max_num_seqs=self.max_num_seqs, | |
| max_model_len=self.max_model_len, | |
| seed=int(self.seed), | |
| enable_prefix_caching=bool(self.enable_prefix_caching), | |
| prefix_caching_hash_algo=str(self.prefix_caching_hash_algo), | |
| ) | |
| self._processor = Qwen3OmniMoeProcessor.from_pretrained(self.model_id) | |
| llm_kwargs: dict[str, Any] = {} | |
| if self.revision is not None: | |
| llm_kwargs["revision"] = self.revision | |
| self._llm = LLM( | |
| model=self.model_id, | |
| trust_remote_code=True, | |
| gpu_memory_utilization=self.gpu_memory_utilization, | |
| tensor_parallel_size=tp_size, | |
| limit_mm_per_prompt={"image": 1, "video": 1, "audio": int(self.limit_mm_per_prompt_audio)}, | |
| max_num_seqs=self.max_num_seqs, | |
| max_model_len=self.max_model_len, | |
| seed=int(self.seed), | |
| enable_prefix_caching=bool(self.enable_prefix_caching), | |
| prefix_caching_hash_algo=str(self.prefix_caching_hash_algo), | |
| **llm_kwargs, | |
| ) | |
| proc_kwargs: dict[str, Any] = {} | |
| if self.revision is not None: | |
| proc_kwargs["revision"] = self.revision | |
| self._processor = Qwen3OmniMoeProcessor.from_pretrained(self.model_id, **proc_kwargs) |
Keep adapter diagnostics opaque by modeling skip reasons and unsupported-language annotations as typed stage-facing fields. Signed-off-by: aaftaabv@gmail.com <aaftaabv@gmail.com>
Document the minimal YAML-driven manifest transcription path so users can exercise the new adapter without duplicating a tutorial runner. Signed-off-by: aaftaabv@gmail.com <aaftaabv@gmail.com>
Warn users that the libsndfile-backed loader skips unsupported AAC/M4A rows so they can transcode inputs and verify output cardinality. Signed-off-by: aaftaabv@gmail.com <aaftaabv@gmail.com>
Keep stage orchestration separate from model inference while addressing review feedback and preserving lazy optional dependencies. Signed-off-by: aaftaabv@gmail.com <aaftaabv@gmail.com>
Point lifecycle tests at the relocated adapter and remove review-requested section-banner noise. Signed-off-by: aaftaabv@gmail.com <aaftaabv@gmail.com>
Expose skip and note fields through the stage contract so pipeline introspection matches the data assembled for failed or filtered rows. Signed-off-by: aaftaabv@gmail.com <aaftaabv@gmail.com>
Signed-off-by: aaftaabv@gmail.com <aaftaabv@gmail.com>
Signed-off-by: aaftaabv@gmail.com <aaftaabv@gmail.com>
Signed-off-by: aaftaabv@gmail.com <aaftaabv@gmail.com>
Signed-off-by: aaftaabv@gmail.com <aaftaabv@gmail.com>
7a1fe4a to
a1dcef4
Compare
| adapter_items = [ | ||
| { | ||
| "waveform": self._load_audio(str(items[index]["audio_filepath"])), | ||
| "language": items[index]["language"], | ||
| "language_code": items[index]["language_code"], | ||
| "reference_text": items[index]["reference_text"], | ||
| "task_id": items[index]["task_id"], | ||
| } | ||
| for index in supported_indices | ||
| ] |
There was a problem hiding this comment.
_load_audio is called inside a bare list comprehension with no per-item exception handling. torchaudio.load raises RuntimeError (file not found) or various audio-library errors on a corrupt file; either propagates out of run_inference and then out of process_batch, losing all successfully-preprocessed items in the batch.
The adapter-side _prepare_single deliberately isolates per-item failures (catch, log warning, return None), but the stage-side loading has no equivalent guard. A single missing resampled audio file — caused by a disk-full write failure or any race on the prior ResampleAudioStage — silently discards the entire 32-item batch instead of skipping that one utterance.
| [[package]] | ||
| name = "mcp" | ||
| version = "1.28.1" | ||
| version = "1.26.0" |
There was a problem hiding this comment.
This shouldn't be downgraded. Probably a bad merge with main?
| "transformers>=4.57.6,<5.0", | ||
| "vllm>=0.18.1,<0.19; (platform_machine == 'x86_64' and platform_system != 'Darwin')", |
There was a problem hiding this comment.
no objection here. cc: @praateekmahajan If you have strong preferences on modality specific vllm/transformers requirements.
There was a problem hiding this comment.
Currently the resolution is vllm<0.19 (from inference server), so why do we need this?
I'm still not clear why we need a new audio_vllm category
|
|
||
|
|
||
| @dataclass | ||
| class ASRStage(ProcessingStage[AudioTask, AudioTask]): |
There was a problem hiding this comment.
Any thoughts on In-process vs inference server for this task? not blocking right now, but in other models/cases we've seen in-process be slower than inference server for generative workloads
There was a problem hiding this comment.
Per preliminary investigations, repackaging audios into json format and sending for inference can cause huge overheads. So, we will for the coming future when audio is input to a gpu stage use inprocess exclusively. Server way of inference works for text data since it is much lighter compared to audio data.
There was a problem hiding this comment.
Payload being heavy only matters if you have C clients sending requests one at a time where C=G as in GPU workers, so the GPU workers are sitting idle while while request is being sent over the wire. However C (being on CPU) can be >> G and we can send request in batched format per C concurrently, therefore GPU is always working even though the next payload is being "prepared".
We saw this with PDFs too where payload is heavier than text but still got ~40% improvement,.
There was a problem hiding this comment.
Fair point Prateek. This warrants more investigation into whether inprocess Vs inference server for audio modality. There are theoretical points supporting both speeding up and slowing down of inference of audio data if we switch to inference server way of things. This investigation is not currently in our team's priority since we are planning to release a new pipeline, granary-v2 soon as part of nemo curator main. Currently entire audio modality in main is inprocess gpu inferences. So, we will surely look into this in the future. Adding this discussion to the doc for long term tracking of this idea.
| ) | ||
| self._sampling_params = SamplingParams(**sampling_kwargs) | ||
| self._processor = Qwen3OmniMoeProcessor.from_pretrained(self.model_id, **proc_kwargs) | ||
| self._prep_pool = ThreadPoolExecutor(max_workers=self.prep_workers) |
Onur is busy with other works right now and has asked me to ask Sarah for reviews.
| def _is_language_supported(self, item: dict[str, Any]) -> bool: | ||
| if self._supported_language_codes is None: | ||
| return True | ||
| code = str(item.get("language_code", "") or "").strip().lower() | ||
| return bool(code) and code in self._supported_language_codes |
There was a problem hiding this comment.
When supported_language_codes is configured (e.g. ["en"]) and a task arrives with no source_lang_key field and no default_language, _resolve_language_code returns None. _build_items therefore stores language_code: None. Inside run_inference, _is_language_supported evaluates str(None or "") == "" → bool("") → False, so the item is filtered before the adapter. The default ASRResult is then constructed with unsupported_language = str(None or "").strip().lower() which is "". Back in assemble (line 396), if unsupported_language: is falsy for the empty string, so no diagnostic note and no _skipme key are written. The task silently receives pred_text = "" with skipped=False, indistinguishable from a valid empty transcript.
Concrete failure: set supported_language_codes: [en], omit default_language, and ingest any task whose manifest line has no source_lang key.
| **self.adapter_kwargs, | ||
| ) | ||
| try: | ||
| adapter.setup(num_gpus=self._adapter_gpu_count()) |
There was a problem hiding this comment.
Look at vllm_utils which has some retry logic to load models. There is a race condition when trying to load multiple vllm models together on a large cluster.
|
|
||
|
|
||
| @contextmanager | ||
| def _mock_qwen_setup( |
There was a problem hiding this comment.
If everything is mocked how will we catch if some version upgrade breaks something?
There was a problem hiding this comment.
Please improve tests to be more realistic, so that we catch actual errors rather than everything being mocked which just makes sure "correct code" is called
| execution_mode=streaming | ||
| ``` | ||
|
|
||
| Use Xenna batch mode only as a fallback when streaming runs out of memory for |
There was a problem hiding this comment.
Do we recommend batch mode?
There was a problem hiding this comment.
No thats why only when seeing OOM issues on xenna_streaming, related to processing large manifest on small hardware we recommend using xenna batch mode since it can help with not loading too many audio files at once in the audio resampling stage. Audio resampling stage is a cpu stage that autoscales on xenna basis free number of cpu threads and not basis available free ram. And resample stage is pretty much an audio staple that needs entire audio file in ram at once per instance.
|
|
||
| ## Select the executor | ||
|
|
||
| The default `backend: ray_data` matches the reference runner. When using |
There was a problem hiding this comment.
Let's default to Ray Data. if you see slow down please ask on core channel
7 participants
Summary
Scope boundary
This PR contains no sharded manifest writer, audio reader, payload lifecycle, global bucketing, or performance instrumentation. Those concerns are isolated in separate PRs.
Executor note
Xenna streaming is the intended default. Use Xenna batch mode only as a fallback when streaming runs out of memory for the workload.
Test plan
64 passed)main