5014fea4a0
Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com> Signed-off-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
107 lines
4.1 KiB
Python
107 lines
4.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Offline E2E smoke test for Audex audio understanding (thinker-only).
|
|
|
|
Single-stage ``audex_thinker_only`` pipeline on
|
|
``checkpoint_folder_full``: WAV (+ instruction) in, text out. The input clip
|
|
is vLLM's public ``mary_had_lamb`` audio asset (downloaded on first use), so
|
|
its transcript is known.
|
|
|
|
WER is a trend metric (recorded in the assertion message, no numeric hard
|
|
gate); the hard gates are structural: non-empty, non-degenerate text.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
|
|
import pytest
|
|
from vllm.assets.audio import AudioAsset
|
|
|
|
from tests.helpers.mark import hardware_test
|
|
from tests.helpers.runtime import OmniRunner
|
|
from tests.helpers.stage_config import get_deploy_config_path
|
|
|
|
MODEL = "nvidia/Nemotron-Labs-Audex-2B"
|
|
MODEL_DIR_ENV = "VLLM_OMNI_AUDEX_MODEL_DIR"
|
|
|
|
# Full transcript of the mary_had_lamb asset (Edison's phonograph recording);
|
|
# a shortened reference would count the extra spoken words as insertions and
|
|
# inflate WER past the trend threshold.
|
|
_REFERENCE = (
|
|
"the first words i spoke in the original phonograph a little piece of practical poetry "
|
|
"mary had a little lamb its fleece was white as snow and everywhere that mary went the "
|
|
"lamb was sure to go"
|
|
)
|
|
|
|
ASR_PROMPT = (
|
|
"<|im_start|>user\n<so_embedding>\nTranscribe the input speech.<|im_end|>\n<|im_start|>assistant\n<think></think>"
|
|
)
|
|
|
|
_audex_deployment = get_deploy_config_path("audex_thinker_only.yaml")
|
|
_audex_model = os.environ.get(MODEL_DIR_ENV) or MODEL
|
|
_OMNI_RUNNER_PARAMS = [
|
|
pytest.param(
|
|
(_audex_model, _audex_deployment, {"async_chunk": False}),
|
|
id="thinker_only",
|
|
),
|
|
]
|
|
# 30B-A3B variant: opt-in (pulls ~60 GB); enable with AUDEX_E2E_30B=1.
|
|
_AUDEX_30B_MODEL = os.environ.get("VLLM_OMNI_AUDEX_30B_MODEL_DIR") or "nvidia/Nemotron-Labs-Audex-30B-A3B"
|
|
if os.environ.get("AUDEX_E2E_30B") == "1":
|
|
_OMNI_RUNNER_PARAMS.append(
|
|
pytest.param(
|
|
(_AUDEX_30B_MODEL, get_deploy_config_path("audex_thinker_only_30b.yaml"), {"async_chunk": False}),
|
|
id="thinker_only_30b",
|
|
)
|
|
)
|
|
pytestmark = [
|
|
pytest.mark.slow,
|
|
pytest.mark.tts,
|
|
pytest.mark.parametrize("omni_runner", _OMNI_RUNNER_PARAMS, indirect=True),
|
|
]
|
|
|
|
|
|
def _normalize(text: str) -> list[str]:
|
|
return "".join(c.lower() if c.isalnum() or c.isspace() else " " for c in text).split()
|
|
|
|
|
|
def _wer(hyp: str, ref: str) -> float:
|
|
hyp_words, ref_words = _normalize(hyp), _normalize(ref)
|
|
if not ref_words:
|
|
return 0.0
|
|
rows = list(range(len(hyp_words) + 1))
|
|
for i, ref_word in enumerate(ref_words, 1):
|
|
prev, rows[0] = rows[0], i
|
|
for j, hyp_word in enumerate(hyp_words, 1):
|
|
cur = min(
|
|
rows[j] + 1,
|
|
rows[j - 1] + 1,
|
|
prev + (ref_word != hyp_word),
|
|
)
|
|
prev, rows[j] = rows[j], cur
|
|
return rows[-1] / len(ref_words)
|
|
|
|
|
|
@pytest.mark.advanced_model
|
|
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
|
def test_audex_offline_asr_smoke(omni_runner: OmniRunner, run_level: str) -> None:
|
|
"""Transcribing the public audio asset yields coherent, non-degenerate text."""
|
|
audio, sr = AudioAsset("mary_had_lamb").audio_and_sample_rate
|
|
outputs = omni_runner.omni.generate([{"prompt": ASR_PROMPT, "multi_modal_data": {"audio": (audio, sr)}}])
|
|
|
|
assert len(outputs) == 1
|
|
text = outputs[0].outputs[0].text if outputs[0].outputs else ""
|
|
assert isinstance(text, str)
|
|
|
|
if run_level in {"advanced_model", "full_model"}:
|
|
words = _normalize(text)
|
|
assert words, f"Empty transcription: {text!r}"
|
|
# Degenerate repetition guard: no word may dominate the output.
|
|
most_common = max(words.count(w) for w in set(words))
|
|
assert most_common <= max(3, len(words) // 2), f"Degenerate transcript: {text!r}"
|
|
wer = _wer(text, _REFERENCE)
|
|
# Trend metric (recorded, not gated): substantially-correct check only.
|
|
assert wer <= 0.9, f"Transcript unrelated to reference (WER={wer:.2f}): {text!r}"
|
|
print(f"[trend] Audex S2T WER vs reference: {wer:.3f} ({text!r})")
|