| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """This example shows how to use vLLM for running offline inference with |
| 4 | the explicit/implicit prompt format on enc-dec LMMs for text generation. |
| 5 | """ |
| 6 | |
| 7 | import os |
| 8 | import time |
| 9 | from collections.abc import Sequence |
| 10 | from typing import NamedTuple |
| 11 | |
| 12 | from vllm import LLM, EngineArgs, PromptType, SamplingParams |
| 13 | from vllm.assets.audio import AudioAsset |
| 14 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 15 | |
| 16 | |
| 17 | class ModelRequestData(NamedTuple): |
| 18 | engine_args: EngineArgs |
| 19 | prompts: Sequence[PromptType] |
| 20 | |
| 21 | |
| 22 | def run_whisper(): |
| 23 | os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" |
| 24 | |
| 25 | engine_args = EngineArgs( |
| 26 | model="openai/whisper-large-v3-turbo", |
| 27 | max_model_len=448, |
| 28 | max_num_seqs=16, |
| 29 | limit_mm_per_prompt={"audio": 1}, |
| 30 | dtype="half", |
| 31 | ) |
| 32 | |
| 33 | prompts = [ |
| 34 | { # Test implicit prompt |
| 35 | "prompt": "<|startoftranscript|>", |
| 36 | "multi_modal_data": { |
| 37 | "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate, |
| 38 | }, |
| 39 | }, |
| 40 | { # Test explicit encoder/decoder prompt |
| 41 | "encoder_prompt": { |
| 42 | "prompt": "", |
| 43 | "multi_modal_data": { |
| 44 | "audio": AudioAsset("winning_call").audio_and_sample_rate, |
| 45 | }, |
| 46 | }, |
| 47 | "decoder_prompt": "<|startoftranscript|>", |
| 48 | }, |
| 49 | ] |
| 50 | |
| 51 | return ModelRequestData( |
| 52 | engine_args=engine_args, |
| 53 | prompts=prompts, |
| 54 | ) |
| 55 | |
| 56 | |
| 57 | def run_fireredasr2(): |
| 58 | """FireRedASR2 – Automatic Speech Recognition model. |
| 59 | |
| 60 | This model uses a Conformer encoder + Qwen2 LLM decoder architecture |
| 61 | for speech-to-text transcription. Audio is passed via the implicit |
| 62 | prompt format with the ``<|AUDIO|>`` placeholder token. |
| 63 | """ |
| 64 | engine_args = EngineArgs( |
| 65 | model="allendou/FireRedASR2-LLM-vllm", |
| 66 | max_model_len=448, |
| 67 | max_num_seqs=16, |
| 68 | limit_mm_per_prompt={"audio": 1}, |
| 69 | ) |
| 70 | |
| 71 | prompt_str = ( |
| 72 | "<|im_start|>user\n<|AUDIO|>请转写音频为文字<|im_end|>\n<|im_start|>assistant\n" |
| 73 | ) |
| 74 | |
| 75 | prompts = [ |
| 76 | { # Implicit prompt with audio |
| 77 | "prompt": prompt_str, |
| 78 | "multi_modal_data": { |
| 79 | "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate, |
| 80 | }, |
| 81 | }, |
| 82 | { # Another audio sample |
| 83 | "prompt": prompt_str, |
| 84 | "multi_modal_data": { |
| 85 | "audio": AudioAsset("winning_call").audio_and_sample_rate, |
| 86 | }, |
| 87 | }, |
| 88 | ] |
| 89 | |
| 90 | return ModelRequestData( |
| 91 | engine_args=engine_args, |
| 92 | prompts=prompts, |
| 93 | ) |
| 94 | |
| 95 | |
| 96 | model_example_map = { |
| 97 | "fireredasr2": run_fireredasr2, |
| 98 | "whisper": run_whisper, |
| 99 | } |
| 100 | |
| 101 | |
| 102 | def parse_args(): |
| 103 | parser = FlexibleArgumentParser( |
| 104 | description="Demo on using vLLM for offline inference with " |
| 105 | "vision language models for text generation" |
| 106 | ) |
| 107 | parser.add_argument( |
| 108 | "--model-type", |
| 109 | "-m", |
| 110 | type=str, |
| 111 | default="whisper", |
| 112 | choices=model_example_map.keys(), |
| 113 | help='Huggingface "model_type".', |
| 114 | ) |
| 115 | parser.add_argument( |
| 116 | "--seed", |
| 117 | type=int, |
| 118 | default=0, |
| 119 | help="Set the seed when initializing `vllm.LLM`.", |
| 120 | ) |
| 121 | return parser.parse_args() |
| 122 | |
| 123 | |
| 124 | def main(args): |
| 125 | model = args.model_type |
| 126 | if model not in model_example_map: |
| 127 | raise ValueError(f"Model type {model} is not supported.") |
| 128 | |
| 129 | req_data = model_example_map[model]() |
| 130 | |
| 131 | # Disable other modalities to save memory |
| 132 | engine_args = req_data.engine_args |
| 133 | default_limits = {"image": 0, "video": 0, "audio": 0} |
| 134 | limit_mm_per_prompt = default_limits | (engine_args.limit_mm_per_prompt or {}) |
| 135 | engine_args.limit_mm_per_prompt = limit_mm_per_prompt |
| 136 | engine_args.seed = args.seed |
| 137 | llm = LLM.from_engine_args(engine_args) |
| 138 | |
| 139 | prompts = req_data.prompts |
| 140 | |
| 141 | # Create a sampling params object. |
| 142 | sampling_params = SamplingParams( |
| 143 | temperature=0, |
| 144 | top_p=1.0, |
| 145 | max_tokens=64, |
| 146 | skip_special_tokens=False, |
| 147 | ) |
| 148 | |
| 149 | start = time.time() |
| 150 | |
| 151 | # Generate output tokens from the prompts. The output is a list of |
| 152 | # RequestOutput objects that contain the prompt, generated |
| 153 | # text, and other information. |
| 154 | outputs = llm.generate(prompts, sampling_params) |
| 155 | |
| 156 | # Print the outputs. |
| 157 | for output in outputs: |
| 158 | prompt = output.prompt |
| 159 | generated_text = output.outputs[0].text |
| 160 | print(f"Decoder prompt: {prompt!r}, Generated text: {generated_text!r}") |
| 161 | |
| 162 | duration = time.time() - start |
| 163 | |
| 164 | print("Duration:", duration) |
| 165 | print("RPS:", len(prompts) / duration) |
| 166 | |
| 167 | |
| 168 | if __name__ == "__main__": |
| 169 | args = parse_args() |
| 170 | main(args) |