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>
94 lines
3.1 KiB
Python
94 lines
3.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Offline Audex (Nemotron-Labs-Audex-2B) audio understanding example.
|
|
|
|
Speech (or general audio) + text instruction → text, through the
|
|
single-stage ``audex_thinker_only`` pipeline: NV-Whisper
|
|
encoder + relu2 projector + 2B dense LM on ``checkpoint_folder_full``.
|
|
|
|
The prompt format mirrors the official audioqa script: a ChatML user turn
|
|
holding the ``<so_embedding>`` placeholder (expanded by the processor to
|
|
750 embedding positions per 30 s clip) plus the instruction, with a closed
|
|
``<think></think>`` priming.
|
|
|
|
Examples (without --audio-files, vLLM's public ``mary_had_lamb`` asset is
|
|
transcribed):
|
|
|
|
# Transcribe WAVs (ASR):
|
|
python examples/offline_inference/audex/audio_qa.py \\
|
|
--audio-files a.wav b.wav
|
|
|
|
# Free-form audio QA:
|
|
python examples/offline_inference/audex/audio_qa.py \\
|
|
--audio-files a.wav --question "What language is being spoken?"
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
from common import default_deploy_config, load_audio_file, request_index
|
|
from vllm.assets.audio import AudioAsset
|
|
|
|
from vllm_omni import Omni
|
|
|
|
ASR_QUESTION = "Transcribe the input speech."
|
|
|
|
|
|
def build_prompt(question: str) -> str:
|
|
return f"<|im_start|>user\n<so_embedding>\n{question}<|im_end|>\n<|im_start|>assistant\n<think></think>"
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="Offline Audex audio understanding")
|
|
parser.add_argument("--model", type=str, default="nvidia/Nemotron-Labs-Audex-2B")
|
|
parser.add_argument(
|
|
"--audio-files",
|
|
type=str,
|
|
nargs="+",
|
|
default=None,
|
|
help="Input clips. Defaults to vLLM's mary_had_lamb audio asset.",
|
|
)
|
|
parser.add_argument("--question", type=str, default=ASR_QUESTION)
|
|
parser.add_argument(
|
|
"--deploy-config",
|
|
type=str,
|
|
# The model root's default deploy yaml is the TTS pipeline; audio
|
|
# understanding needs the single-stage thinker-only pipeline.
|
|
default=default_deploy_config("audex_thinker_only.yaml"),
|
|
help="Deploy yaml (defaults to the audex_thinker_only pipeline).",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
|
|
|
|
prompt_text = build_prompt(args.question)
|
|
if args.audio_files:
|
|
clips = [load_audio_file(path) for path in args.audio_files]
|
|
else:
|
|
clips = [AudioAsset("mary_had_lamb").audio_and_sample_rate]
|
|
prompts = []
|
|
for audio, sr in clips:
|
|
prompts.append(
|
|
{
|
|
"prompt": prompt_text,
|
|
"multi_modal_data": {"audio": (audio, sr)},
|
|
}
|
|
)
|
|
|
|
outputs = engine.generate(prompts)
|
|
|
|
ordered = sorted(outputs, key=request_index)
|
|
names = [Path(p).name for p in args.audio_files] if args.audio_files else ["mary_had_lamb"]
|
|
for name, req_output in zip(names, ordered):
|
|
text = req_output.outputs[0].text if req_output.outputs else ""
|
|
print(f"{name}: {text.strip()}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|