audio_language.py 16.4 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""
4
This example shows how to use vLLM for running offline inference
5
with the correct prompt format on audio language models.
6
7
8
9

For most models, the prompt format should follow corresponding examples
on HuggingFace model repository.
"""
10

11
import os
12
from dataclasses import asdict
13
from typing import Any, NamedTuple
14
15

from huggingface_hub import snapshot_download
16
17
from transformers import AutoTokenizer

18
from vllm import LLM, EngineArgs, SamplingParams
19
from vllm.assets.audio import AudioAsset
20
from vllm.lora.request import LoRARequest
21
from vllm.utils.argparse_utils import FlexibleArgumentParser
22

23
audio_assets = [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
24
25
26
question_per_audio_count = {
    0: "What is 1+1?",
    1: "What is recited in the audio?",
27
    2: "What sport and what nursery rhyme are referenced?",
28
}
29

30
31
32

class ModelRequestData(NamedTuple):
    engine_args: EngineArgs
33
34
35
36
37
    prompt: str | None = None
    prompt_token_ids: dict[str, list[int]] | None = None
    multi_modal_data: dict[str, Any] | None = None
    stop_token_ids: list[int] | None = None
    lora_requests: list[LoRARequest] | None = None
38
39


40
41
42
43
# NOTE: The default `max_num_seqs` and `max_model_len` may result in OOM on
# lower-end GPUs.
# Unless specified, these settings have been tested to work on a single L4.

44

45
46
47
# AudioFlamingo3
def run_audioflamingo3(question: str, audio_count: int) -> ModelRequestData:
    model_name = "nvidia/audio-flamingo-3-hf"
Patrick von Platen's avatar
Patrick von Platen committed
48
49
    engine_args = EngineArgs(
        model=model_name,
50
        max_model_len=4096,
Patrick von Platen's avatar
Patrick von Platen committed
51
52
53
54
55
        max_num_seqs=2,
        limit_mm_per_prompt={"audio": audio_count},
        enforce_eager=True,
    )

56
57
    # AudioFlamingo3 uses <sound> token for audio
    audio_placeholder = "<sound>" * audio_count
Patrick von Platen's avatar
Patrick von Platen committed
58

59
60
61
62
63
64
65
    prompt = (
        "<|im_start|>system\n"
        "You are a helpful assistant.<|im_end|>\n"
        "<|im_start|>user\n"
        f"{audio_placeholder}{question}<|im_end|>\n"
        "<|im_start|>assistant\n"
    )
Patrick von Platen's avatar
Patrick von Platen committed
66
67
68

    return ModelRequestData(
        engine_args=engine_args,
69
        prompt=prompt,
Patrick von Platen's avatar
Patrick von Platen committed
70
71
72
    )


Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
# Gemma3N
def run_gemma3n(question: str, audio_count: int) -> ModelRequestData:
    model_name = "google/gemma-3n-E2B-it"
    engine_args = EngineArgs(
        model=model_name,
        max_model_len=2048,
        max_num_batched_tokens=2048,
        max_num_seqs=2,
        limit_mm_per_prompt={"audio": audio_count},
        enforce_eager=True,
    )
    prompt = f"<start_of_turn>user\n<audio_soft_token>{question}"
    "<end_of_turn>\n<start_of_turn>model\n"
    return ModelRequestData(
        engine_args=engine_args,
        prompt=prompt,
    )


92
93
# Granite Speech
def run_granite_speech(question: str, audio_count: int) -> ModelRequestData:
94
    # NOTE - the setting in this example are somewhat different from what is
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
    # optimal for granite speech, and it is generally recommended to use beam
    # search. Check the model README for suggested settings.
    # https://huggingface.co/ibm-granite/granite-speech-3.3-8b
    model_name = "ibm-granite/granite-speech-3.3-8b"

    engine_args = EngineArgs(
        model=model_name,
        trust_remote_code=True,
        max_model_len=2048,
        max_num_seqs=2,
        enable_lora=True,
        max_lora_rank=64,
        limit_mm_per_prompt={"audio": audio_count},
    )

    # The model has an audio-specific lora directly in its model dir;
    # it should be enabled whenever you pass audio inputs to the model.
    speech_lora_path = model_name
    audio_placeholder = "<|audio|>" * audio_count
    prompts = f"<|start_of_role|>system<|end_of_role|>Knowledge Cutoff Date: April 2024.\nToday's Date: December 19, 2024.\nYou are Granite, developed by IBM. You are a helpful AI assistant<|end_of_text|>\n<|start_of_role|>user<|end_of_role|>{audio_placeholder}{question}<|end_of_text|>\n<|start_of_role|>assistant<|end_of_role|>"  # noqa: E501

    return ModelRequestData(
        engine_args=engine_args,
        prompt=prompts,
        lora_requests=[LoRARequest("speech", 1, speech_lora_path)],
    )


123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
# MiDashengLM
def run_midashenglm(question: str, audio_count: int):
    model_name = "mispeech/midashenglm-7b"

    engine_args = EngineArgs(
        model=model_name,
        trust_remote_code=True,
        max_model_len=4096,
        max_num_seqs=5,
        limit_mm_per_prompt={"audio": audio_count},
    )

    audio_in_prompt = "".join(
        ["<|audio_bos|><|AUDIO|><|audio_eos|>" for idx in range(audio_count)]
    )

    default_system = "You are a helpful language and speech assistant."

    prompt = (
        f"<|im_start|>system\n{default_system}<|im_end|>\n"
        "<|im_start|>user\n"
        f"{audio_in_prompt}{question}<|im_end|>\n"
        "<|im_start|>assistant\n"
    )
    return ModelRequestData(
        engine_args=engine_args,
        prompt=prompt,
    )


153
# MiniCPM-O
154
def run_minicpmo(question: str, audio_count: int) -> ModelRequestData:
155
    model_name = "openbmb/MiniCPM-o-2_6"
156
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
157
158
159
160
    engine_args = EngineArgs(
        model=model_name,
        trust_remote_code=True,
        max_model_len=4096,
161
        max_num_seqs=2,
162
163
        limit_mm_per_prompt={"audio": audio_count},
    )
164

165
    stop_tokens = ["<|im_end|>", "<|endoftext|>"]
166
167
168
169
    stop_token_ids = [tokenizer.convert_tokens_to_ids(i) for i in stop_tokens]

    audio_placeholder = "(<audio>./</audio>)" * audio_count
    audio_chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n<|spk_bos|><|spk|><|spk_eos|><|tts_bos|>' }}{% endif %}"  # noqa: E501
170
171
172
173
174
175
176
    messages = [{"role": "user", "content": f"{audio_placeholder}\n{question}"}]
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        chat_template=audio_chat_template,
    )
177
178
179
180
181
182

    return ModelRequestData(
        engine_args=engine_args,
        prompt=prompt,
        stop_token_ids=stop_token_ids,
    )
183
184


185
# Phi-4-multimodal-instruct
186
def run_phi4mm(question: str, audio_count: int) -> ModelRequestData:
187
188
189
190
191
192
193
194
    """
    Phi-4-multimodal-instruct supports both image and audio inputs. Here, we
    show how to process audio inputs.
    """
    model_path = snapshot_download("microsoft/Phi-4-multimodal-instruct")
    # Since the vision-lora and speech-lora co-exist with the base model,
    # we have to manually specify the path of the lora weights.
    speech_lora_path = os.path.join(model_path, "speech-lora")
195
    placeholders = "".join([f"<|audio_{i + 1}|>" for i in range(audio_count)])
196

197
    prompts = f"<|user|>{placeholders}{question}<|end|><|assistant|>"
198

199
    engine_args = EngineArgs(
200
201
        model=model_path,
        trust_remote_code=True,
202
        max_model_len=12800,
203
204
205
        max_num_seqs=2,
        enable_lora=True,
        max_lora_rank=320,
206
        limit_mm_per_prompt={"audio": audio_count},
207
208
    )

209
210
211
212
213
    return ModelRequestData(
        engine_args=engine_args,
        prompt=prompts,
        lora_requests=[LoRARequest("speech", 1, speech_lora_path)],
    )
214
215


216
# Qwen2-Audio
217
def run_qwen2_audio(question: str, audio_count: int) -> ModelRequestData:
218
219
    model_name = "Qwen/Qwen2-Audio-7B-Instruct"

220
221
222
223
224
225
    engine_args = EngineArgs(
        model=model_name,
        max_model_len=4096,
        max_num_seqs=5,
        limit_mm_per_prompt={"audio": audio_count},
    )
226

227
228
229
230
231
232
    audio_in_prompt = "".join(
        [
            f"Audio {idx + 1}: <|audio_bos|><|AUDIO|><|audio_eos|>\n"
            for idx in range(audio_count)
        ]
    )
233

234
235
236
237
238
239
    prompt = (
        "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
        "<|im_start|>user\n"
        f"{audio_in_prompt}{question}<|im_end|>\n"
        "<|im_start|>assistant\n"
    )
240
241
242
243
244

    return ModelRequestData(
        engine_args=engine_args,
        prompt=prompt,
    )
245
246


247
248
249
250
251
252
253
254
255
256
257
# Qwen2.5-Omni
def run_qwen2_5_omni(question: str, audio_count: int):
    model_name = "Qwen/Qwen2.5-Omni-7B"

    engine_args = EngineArgs(
        model=model_name,
        max_model_len=4096,
        max_num_seqs=5,
        limit_mm_per_prompt={"audio": audio_count},
    )

258
259
260
    audio_in_prompt = "".join(
        ["<|audio_bos|><|AUDIO|><|audio_eos|>\n" for idx in range(audio_count)]
    )
261
262
263
264

    default_system = (
        "You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
        "Group, capable of perceiving auditory and visual inputs, as well as "
265
266
        "generating text and speech."
    )
267

268
269
270
271
272
273
    prompt = (
        f"<|im_start|>system\n{default_system}<|im_end|>\n"
        "<|im_start|>user\n"
        f"{audio_in_prompt}{question}<|im_end|>\n"
        "<|im_start|>assistant\n"
    )
274
275
276
277
278
279
    return ModelRequestData(
        engine_args=engine_args,
        prompt=prompt,
    )


280
# Ultravox 0.5-1B
281
def run_ultravox(question: str, audio_count: int) -> ModelRequestData:
282
    model_name = "fixie-ai/ultravox-v0_5-llama-3_2-1b"
283

284
    tokenizer = AutoTokenizer.from_pretrained(model_name)
285
286
287
288
    messages = [{"role": "user", "content": "<|audio|>\n" * audio_count + question}]
    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
289

290
291
292
293
294
295
296
297
298
299
300
301
    engine_args = EngineArgs(
        model=model_name,
        max_model_len=4096,
        max_num_seqs=5,
        trust_remote_code=True,
        limit_mm_per_prompt={"audio": audio_count},
    )

    return ModelRequestData(
        engine_args=engine_args,
        prompt=prompt,
    )
302
303


304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
# Voxtral
# Make sure to install mistral-common[audio].
def run_voxtral(question: str, audio_count: int) -> ModelRequestData:
    from mistral_common.audio import Audio
    from mistral_common.protocol.instruct.chunk import (
        AudioChunk,
        RawAudio,
        TextChunk,
    )
    from mistral_common.protocol.instruct.messages import (
        UserMessage,
    )
    from mistral_common.protocol.instruct.request import ChatCompletionRequest
    from mistral_common.tokens.tokenizers.mistral import MistralTokenizer

    model_name = "mistralai/Voxtral-Mini-3B-2507"
    tokenizer = MistralTokenizer.from_hf_hub(model_name)

    engine_args = EngineArgs(
        model=model_name,
        max_model_len=8192,
        max_num_seqs=2,
        limit_mm_per_prompt={"audio": audio_count},
        config_format="mistral",
        load_format="mistral",
        tokenizer_mode="mistral",
        enforce_eager=True,
        enable_chunked_prefill=False,
    )

    text_chunk = TextChunk(text=question)
    audios = [
        Audio.from_file(str(audio_assets[i].get_local_path()), strict=False)
        for i in range(audio_count)
    ]
    audio_chunks = [
        AudioChunk(input_audio=RawAudio.from_audio(audio)) for audio in audios
    ]

    messages = [UserMessage(content=[*audio_chunks, text_chunk])]

    req = ChatCompletionRequest(messages=messages, model=model_name)

    tokens = tokenizer.encode_chat_completion(req)
    prompt_ids, audios = tokens.tokens, tokens.audios

    audios_and_sr = [(au.audio_array, au.sampling_rate) for au in audios]

    multi_modal_data = {"audio": audios_and_sr}

    return ModelRequestData(
        engine_args=engine_args,
        prompt_token_ids=prompt_ids,
        multi_modal_data=multi_modal_data,
    )


361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
# GLM-ASR
def run_glmasr(question: str, audio_count: int) -> ModelRequestData:
    model_name = "zai-org/GLM-ASR-Nano-2512"

    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

    # GLM-ASR uses <|pad|> token for audio
    audio_placeholder = "<|pad|>" * audio_count

    messages = [{"role": "user", "content": f"{audio_placeholder}{question}"}]
    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )

    engine_args = EngineArgs(
        model=model_name,
        trust_remote_code=True,
        max_model_len=4096,
        max_num_seqs=2,
        limit_mm_per_prompt={"audio": audio_count},
    )

    return ModelRequestData(
        engine_args=engine_args,
        prompt=prompt,
    )


389
# Whisper
390
def run_whisper(question: str, audio_count: int) -> ModelRequestData:
391
    assert audio_count == 1, "Whisper only support single audio input per prompt"
392
393
394
395
    model_name = "openai/whisper-large-v3-turbo"

    prompt = "<|startoftranscript|>"

396
397
398
399
400
401
402
403
404
405
406
    engine_args = EngineArgs(
        model=model_name,
        max_model_len=448,
        max_num_seqs=5,
        limit_mm_per_prompt={"audio": audio_count},
    )

    return ModelRequestData(
        engine_args=engine_args,
        prompt=prompt,
    )
407
408
409


model_example_map = {
410
    "audioflamingo3": run_audioflamingo3,
Nicolò Lucchesi's avatar
Nicolò Lucchesi committed
411
    "gemma3n": run_gemma3n,
412
    "glmasr": run_glmasr,
413
    "granite_speech": run_granite_speech,
414
    "midashenglm": run_midashenglm,
415
    "minicpmo": run_minicpmo,
416
    "phi4_mm": run_phi4mm,
417
    "qwen2_audio": run_qwen2_audio,
418
    "qwen2_5_omni": run_qwen2_5_omni,
419
    "ultravox": run_ultravox,
420
    "voxtral": run_voxtral,
421
    "whisper": run_whisper,
422
}
423
424


425
426
def parse_args():
    parser = FlexibleArgumentParser(
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
        description="Demo on using vLLM for offline inference with "
        "audio language models"
    )
    parser.add_argument(
        "--model-type",
        "-m",
        type=str,
        default="ultravox",
        choices=model_example_map.keys(),
        help='Huggingface "model_type".',
    )
    parser.add_argument(
        "--num-prompts", type=int, default=1, help="Number of prompts to run."
    )
    parser.add_argument(
        "--num-audios",
        type=int,
        default=1,
        choices=[0, 1, 2],
        help="Number of audio items per prompt.",
    )
    parser.add_argument(
        "--seed",
        type=int,
451
        default=0,
452
453
        help="Set the seed when initializing `vllm.LLM`.",
    )
454
455
456
457
458
459
460
    parser.add_argument(
        "--tensor-parallel-size",
        "-tp",
        type=int,
        default=None,
        help="Tensor parallel size to override the model's default setting. ",
    )
461
462
463
464

    return parser.parse_args()


465
466
467
468
469
def main(args):
    model = args.model_type
    if model not in model_example_map:
        raise ValueError(f"Model type {model} is not supported.")

470
471
472
473
474
475
    if args.tensor_parallel_size is not None and args.tensor_parallel_size < 1:
        raise ValueError(
            f"tensor_parallel_size must be a positive integer, "
            f"got {args.tensor_parallel_size}"
        )

476
    audio_count = args.num_audios
477
478
479
    req_data = model_example_map[model](
        question_per_audio_count[audio_count], audio_count
    )
480

481
482
483
    # Disable other modalities to save memory
    default_limits = {"image": 0, "video": 0, "audio": 0}
    req_data.engine_args.limit_mm_per_prompt = default_limits | dict(
484
485
        req_data.engine_args.limit_mm_per_prompt or {}
    )
486

487
    engine_args = asdict(req_data.engine_args) | {"seed": args.seed}
488
489
    if args.tensor_parallel_size is not None:
        engine_args["tensor_parallel_size"] = args.tensor_parallel_size
490
491
    llm = LLM(**engine_args)

492
493
    # We set temperature to 0.2 so that outputs can be different
    # even when all prompts are identical when running batch inference.
494
495
496
    sampling_params = SamplingParams(
        temperature=0.2, max_tokens=64, stop_token_ids=req_data.stop_token_ids
    )
497

498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
    def get_input(start, end):
        mm_data = req_data.multi_modal_data
        if not mm_data:
            mm_data = {}
            if end - start > 0:
                mm_data = {
                    "audio": [
                        asset.audio_and_sample_rate for asset in audio_assets[start:end]
                    ]
                }

        inputs = {"multi_modal_data": mm_data}

        if req_data.prompt:
            inputs["prompt"] = req_data.prompt
        else:
            inputs["prompt_token_ids"] = req_data.prompt_token_ids

        return inputs

    # Batch inference
519
    assert args.num_prompts > 0
520
521
522
    if audio_count != 1:
        inputs = get_input(0, audio_count)
        inputs = [inputs] * args.num_prompts
Patrick von Platen's avatar
Patrick von Platen committed
523
    else:
524
525
526
527
528
529
530
        # For single audio input, we need to vary the audio input
        # to avoid deduplication in vLLM engine.
        inputs = []
        for i in range(args.num_prompts):
            start = i % len(audio_assets)
            inp = get_input(start, start + 1)
            inputs.append(inp)
Patrick von Platen's avatar
Patrick von Platen committed
531

532
    # Add LoRA request if applicable
533
534
535
    lora_request = (
        req_data.lora_requests * args.num_prompts if req_data.lora_requests else None
    )
536
537
538
539
540
541

    outputs = llm.generate(
        inputs,
        sampling_params=sampling_params,
        lora_request=lora_request,
    )
542
543
544
545
546
547
548

    for o in outputs:
        generated_text = o.outputs[0].text
        print(generated_text)


if __name__ == "__main__":
549
    args = parse_args()
550
    main(args)