@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 217 lines 6.47 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """This script demonstrates how to use the vLLM API server to perform audio
4 transcription with the `openai/whisper-large-v3` model.
5
6 Before running this script, you must start the vLLM server with the following command:
7
8 vllm serve openai/whisper-large-v3
9
10 Requirements:
11 - vLLM with audio support
12 - openai Python SDK
13 - httpx for streaming support
14
15 The script performs:
16 1. Synchronous transcription using OpenAI-compatible API.
17 2. Streaming transcription using raw HTTP request to the vLLM server.
18 """
19
20 import argparse
21 import asyncio
22
23 from openai import AsyncOpenAI, OpenAI
24
25 from vllm.assets.audio import AudioAsset
26
27
28 def sync_openai(
29 audio_path: str,
30 client: OpenAI,
31 model: str,
32 *,
33 repetition_penalty: float = 1.3,
34 hotwords: str = None,
35 prompt: str | None = None,
36 ):
37 """Perform synchronous transcription using OpenAI-compatible API.
38
39 The optional ``prompt`` is the OpenAI-API ``prompt`` field (style /
40 vocabulary hint). It is wired through model-by-model: Whisper uses it
41 as a ``<|prev|>`` continuation hint, Qwen3-ASR maps it into the
42 chat-template ``system`` turn. Models that do not consume it accept
43 it without effect.
44 """
45 with open(audio_path, "rb") as f:
46 transcription = client.audio.transcriptions.create(
47 file=f,
48 model=model,
49 language="en",
50 prompt=prompt or "",
51 response_format="json",
52 temperature=0.0,
53 # Additional sampling params not provided by OpenAI API.
54 extra_body=dict(
55 seed=4419,
56 repetition_penalty=repetition_penalty,
57 hotwords=hotwords,
58 ),
59 )
60 print("transcription result [sync]:", transcription.text)
61
62
63 async def stream_openai_response(
64 audio_path: str,
65 client: AsyncOpenAI,
66 model: str,
67 hotwords: str = None,
68 prompt: str | None = None,
69 ):
70 """Perform asynchronous transcription using OpenAI-compatible API."""
71 print("\ntranscription result [stream]:", end=" ")
72 with open(audio_path, "rb") as f:
73 transcription = await client.audio.transcriptions.create(
74 file=f,
75 model=model,
76 language="en",
77 prompt=prompt or "",
78 response_format="json",
79 temperature=0.0,
80 # Additional sampling params not provided by OpenAI API.
81 extra_body=dict(
82 seed=420,
83 top_p=0.6,
84 hotwords=hotwords,
85 ),
86 stream=True,
87 )
88 async for chunk in transcription:
89 if chunk.choices:
90 content = chunk.choices[0].get("delta", {}).get("content")
91 print(content, end="", flush=True)
92
93 print() # Final newline after stream ends
94
95
96 def stream_api_response(audio_path: str, model: str, openai_api_base: str):
97 """Stream a transcription using raw HTTP requests to the vLLM API server."""
98 import json
99 import os
100
101 import requests
102
103 api_url = f"{openai_api_base}/audio/transcriptions"
104 headers = {"User-Agent": "Transcription-Client"}
105 with open(audio_path, "rb") as f:
106 files = {"file": (os.path.basename(audio_path), f)}
107 data = {
108 "stream": "true",
109 "model": model,
110 "language": "en",
111 "response_format": "json",
112 }
113
114 print("\ntranscription result [stream]:", end=" ")
115 response = requests.post(
116 api_url, headers=headers, files=files, data=data, stream=True
117 )
118 for chunk in response.iter_lines(
119 chunk_size=8192, decode_unicode=False, delimiter=b"\n"
120 ):
121 if chunk:
122 data = chunk[len("data: ") :]
123 data = json.loads(data.decode("utf-8"))
124 data = data["choices"][0]
125 delta = data["delta"]["content"]
126 print(delta, end="", flush=True)
127
128 finish_reason = data.get("finish_reason")
129 if finish_reason is not None:
130 print(f"\n[Stream finished reason: {finish_reason}]")
131 break
132
133
134 def main(args):
135 mary_had_lamb = str(AudioAsset("mary_had_lamb").get_local_path())
136 winning_call = str(AudioAsset("winning_call").get_local_path())
137
138 # Modify OpenAI's API key and API base to use vLLM's API server.
139 openai_api_key = "EMPTY"
140 openai_api_base = "http://localhost:8000/v1"
141 client = OpenAI(
142 api_key=openai_api_key,
143 base_url=openai_api_base,
144 )
145
146 model = client.models.list().data[0].id
147 print(f"Using model: {model}")
148
149 # Run the synchronous function
150 sync_openai(
151 audio_path=args.audio_path if args.audio_path else mary_had_lamb,
152 client=client,
153 model=model,
154 repetition_penalty=args.repetition_penalty,
155 hotwords=args.hotwords,
156 prompt=args.prompt,
157 )
158
159 # Run the asynchronous function
160 if "openai" in model:
161 client = AsyncOpenAI(
162 api_key=openai_api_key,
163 base_url=openai_api_base,
164 )
165 asyncio.run(
166 stream_openai_response(
167 args.audio_path if args.audio_path else winning_call,
168 client,
169 model,
170 hotwords=args.hotwords,
171 prompt=args.prompt,
172 )
173 )
174 else:
175 stream_api_response(
176 args.audio_path if args.audio_path else winning_call,
177 model,
178 openai_api_base,
179 )
180
181
182 if __name__ == "__main__":
183 # setup argparser
184 parser = argparse.ArgumentParser(
185 description="OpenAI Transcription Client using vLLM API Server"
186 )
187 parser.add_argument(
188 "--audio_path",
189 type=str,
190 default=None,
191 help="The path to the audio file to transcribe.",
192 )
193 parser.add_argument(
194 "--repetition_penalty",
195 type=float,
196 default=1.3,
197 help="repetition penalty",
198 )
199 parser.add_argument(
200 "--hotwords",
201 type=str,
202 default=None,
203 help="hotwords",
204 )
205 parser.add_argument(
206 "--prompt",
207 type=str,
208 default=None,
209 help=(
210 "Optional `prompt` (OpenAI transcription API: style/vocabulary "
211 "hint). Wired model-by-model: Whisper uses it as a `<|prev|>` "
212 "continuation hint, Qwen3-ASR maps it into the chat-template "
213 "system turn."
214 ),
215 )
216 args = parser.parse_args()
217 main(args)