| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import io |
| 5 | import json |
| 6 | import os |
| 7 | import sys |
| 8 | import time |
| 9 | import traceback |
| 10 | from dataclasses import dataclass, field |
| 11 | |
| 12 | import aiohttp |
| 13 | import huggingface_hub.constants |
| 14 | from tqdm.asyncio import tqdm |
| 15 | from transformers import AutoTokenizer, PythonBackend, TokenizersBackend |
| 16 | |
| 17 | # NOTE(simon): do not import vLLM here so the benchmark script |
| 18 | # can run without vLLM installed. |
| 19 | |
| 20 | AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60) |
| 21 | |
| 22 | |
| 23 | @dataclass |
| 24 | class RequestFuncInput: |
| 25 | prompt: str |
| 26 | api_url: str |
| 27 | prompt_len: int |
| 28 | output_len: int |
| 29 | model: str |
| 30 | model_name: str | None = None |
| 31 | logprobs: int | None = None |
| 32 | extra_body: dict | None = None |
| 33 | multi_modal_content: dict | list[dict] | None = None |
| 34 | ignore_eos: bool = False |
| 35 | language: str | None = None |
| 36 | request_id: str | None = None |
| 37 | |
| 38 | |
| 39 | @dataclass |
| 40 | class RequestFuncOutput: |
| 41 | generated_text: str = "" |
| 42 | success: bool = False |
| 43 | latency: float = 0.0 |
| 44 | output_tokens: int = 0 |
| 45 | ttft: float = 0.0 # Time to first token |
| 46 | itl: list[float] = field(default_factory=list) # list of inter-token latencies |
| 47 | tpot: float = 0.0 # avg next-token latencies |
| 48 | prompt_len: int = 0 |
| 49 | error: str = "" |
| 50 | |
| 51 | |
| 52 | async def async_request_tgi( |
| 53 | request_func_input: RequestFuncInput, |
| 54 | pbar: tqdm | None = None, |
| 55 | ) -> RequestFuncOutput: |
| 56 | api_url = request_func_input.api_url |
| 57 | assert api_url.endswith("generate_stream") |
| 58 | |
| 59 | async with aiohttp.ClientSession( |
| 60 | trust_env=True, timeout=AIOHTTP_TIMEOUT |
| 61 | ) as session: |
| 62 | params = { |
| 63 | "max_new_tokens": request_func_input.output_len, |
| 64 | "do_sample": True, |
| 65 | "temperature": 0.01, # TGI does not accept 0.0 temperature. |
| 66 | "top_p": 0.99, # TGI does not accept 1.0 top_p. |
| 67 | "truncate": request_func_input.prompt_len, |
| 68 | "ignore_eos_token": request_func_input.ignore_eos, |
| 69 | } |
| 70 | payload = { |
| 71 | "inputs": request_func_input.prompt, |
| 72 | "parameters": params, |
| 73 | } |
| 74 | headers = None |
| 75 | if request_func_input.request_id: |
| 76 | headers = {"x-request-id": request_func_input.request_id} |
| 77 | output = RequestFuncOutput() |
| 78 | output.prompt_len = request_func_input.prompt_len |
| 79 | if request_func_input.ignore_eos: |
| 80 | output.output_tokens = request_func_input.output_len |
| 81 | else: |
| 82 | output.output_tokens = None |
| 83 | |
| 84 | ttft = 0.0 |
| 85 | st = time.perf_counter() |
| 86 | most_recent_timestamp = st |
| 87 | try: |
| 88 | async with session.post( |
| 89 | url=api_url, json=payload, headers=headers |
| 90 | ) as response: |
| 91 | if response.status == 200: |
| 92 | async for chunk_bytes in response.content: |
| 93 | chunk_bytes = chunk_bytes.strip() |
| 94 | if not chunk_bytes: |
| 95 | continue |
| 96 | chunk_bytes = chunk_bytes.decode("utf-8") |
| 97 | |
| 98 | # NOTE: Sometimes TGI returns a ping response without |
| 99 | # any data, we should skip it. |
| 100 | if chunk_bytes.startswith(":"): |
| 101 | continue |
| 102 | chunk = chunk_bytes.removeprefix("data:") |
| 103 | |
| 104 | data = json.loads(chunk) |
| 105 | timestamp = time.perf_counter() |
| 106 | # First token |
| 107 | if ttft == 0.0: |
| 108 | ttft = timestamp - st |
| 109 | output.ttft = ttft |
| 110 | |
| 111 | # Decoding phase |
| 112 | else: |
| 113 | output.itl.append(timestamp - most_recent_timestamp) |
| 114 | |
| 115 | most_recent_timestamp = timestamp |
| 116 | |
| 117 | output.latency = most_recent_timestamp - st |
| 118 | output.success = True |
| 119 | output.generated_text = data["generated_text"] |
| 120 | else: |
| 121 | output.error = response.reason or "" |
| 122 | output.success = False |
| 123 | except Exception: |
| 124 | output.success = False |
| 125 | exc_info = sys.exc_info() |
| 126 | output.error = "".join(traceback.format_exception(*exc_info)) |
| 127 | |
| 128 | if pbar: |
| 129 | pbar.update(1) |
| 130 | return output |
| 131 | |
| 132 | |
| 133 | async def async_request_trt_llm( |
| 134 | request_func_input: RequestFuncInput, |
| 135 | pbar: tqdm | None = None, |
| 136 | ) -> RequestFuncOutput: |
| 137 | api_url = request_func_input.api_url |
| 138 | assert api_url.endswith("generate_stream") |
| 139 | |
| 140 | async with aiohttp.ClientSession( |
| 141 | trust_env=True, timeout=AIOHTTP_TIMEOUT |
| 142 | ) as session: |
| 143 | payload = { |
| 144 | "accumulate_tokens": True, |
| 145 | "text_input": request_func_input.prompt, |
| 146 | "temperature": 0.0, |
| 147 | "top_p": 1.0, |
| 148 | "max_tokens": request_func_input.output_len, |
| 149 | "stream": True, |
| 150 | } |
| 151 | if request_func_input.ignore_eos: |
| 152 | payload["min_length"] = request_func_input.output_len |
| 153 | headers = None |
| 154 | if request_func_input.request_id: |
| 155 | headers = {"x-request-id": request_func_input.request_id} |
| 156 | output = RequestFuncOutput() |
| 157 | output.prompt_len = request_func_input.prompt_len |
| 158 | |
| 159 | ttft = 0.0 |
| 160 | st = time.perf_counter() |
| 161 | most_recent_timestamp = st |
| 162 | try: |
| 163 | async with session.post( |
| 164 | url=api_url, json=payload, headers=headers |
| 165 | ) as response: |
| 166 | if response.status == 200: |
| 167 | async for chunk_bytes in response.content: |
| 168 | chunk_bytes = chunk_bytes.strip() |
| 169 | if not chunk_bytes: |
| 170 | continue |
| 171 | |
| 172 | chunk = chunk_bytes.decode("utf-8").removeprefix("data:") |
| 173 | |
| 174 | data = json.loads(chunk) |
| 175 | output.generated_text += data["text_output"] |
| 176 | timestamp = time.perf_counter() |
| 177 | # First token |
| 178 | if ttft == 0.0: |
| 179 | ttft = timestamp - st |
| 180 | output.ttft = ttft |
| 181 | |
| 182 | # Decoding phase |
| 183 | else: |
| 184 | output.itl.append(timestamp - most_recent_timestamp) |
| 185 | |
| 186 | most_recent_timestamp = timestamp |
| 187 | |
| 188 | output.latency = most_recent_timestamp - st |
| 189 | output.success = True |
| 190 | |
| 191 | else: |
| 192 | output.error = response.reason or "" |
| 193 | output.success = False |
| 194 | except Exception: |
| 195 | output.success = False |
| 196 | exc_info = sys.exc_info() |
| 197 | output.error = "".join(traceback.format_exception(*exc_info)) |
| 198 | |
| 199 | if pbar: |
| 200 | pbar.update(1) |
| 201 | return output |
| 202 | |
| 203 | |
| 204 | async def async_request_deepspeed_mii( |
| 205 | request_func_input: RequestFuncInput, |
| 206 | pbar: tqdm | None = None, |
| 207 | ) -> RequestFuncOutput: |
| 208 | api_url = request_func_input.api_url |
| 209 | assert api_url.endswith(("completions", "profile")), ( |
| 210 | "OpenAI Completions API URL must end with 'completions' or 'profile'." |
| 211 | ) |
| 212 | |
| 213 | async with aiohttp.ClientSession( |
| 214 | trust_env=True, timeout=AIOHTTP_TIMEOUT |
| 215 | ) as session: |
| 216 | payload = { |
| 217 | "model": request_func_input.model, |
| 218 | "prompt": request_func_input.prompt, |
| 219 | "max_tokens": request_func_input.output_len, |
| 220 | "temperature": 0.01, # deepspeed-mii does not accept 0.0 temp. |
| 221 | "top_p": 1.0, |
| 222 | } |
| 223 | headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"} |
| 224 | if request_func_input.request_id: |
| 225 | headers["x-request-id"] = request_func_input.request_id |
| 226 | |
| 227 | output = RequestFuncOutput() |
| 228 | output.prompt_len = request_func_input.prompt_len |
| 229 | |
| 230 | # NOTE: DeepSpeed-MII doesn't support streaming as of Jan 28 2024, |
| 231 | # will use 0 as placeholder. |
| 232 | # See https://github.com/microsoft/DeepSpeed-MII/pull/311 |
| 233 | output.ttft = 0 |
| 234 | |
| 235 | st = time.perf_counter() |
| 236 | try: |
| 237 | async with session.post( |
| 238 | url=api_url, json=payload, headers=headers |
| 239 | ) as response: |
| 240 | if response.status == 200: |
| 241 | parsed_resp = await response.json() |
| 242 | output.latency = time.perf_counter() - st |
| 243 | if "choices" in parsed_resp: |
| 244 | output.generated_text = parsed_resp["choices"][0]["text"] |
| 245 | elif "text" in parsed_resp: |
| 246 | output.generated_text = parsed_resp["text"][0] |
| 247 | else: |
| 248 | output.error = ( |
| 249 | "Unexpected response format: " |
| 250 | "neither 'choices' nor 'text' found" |
| 251 | ) |
| 252 | output.success = False |
| 253 | output.success = True |
| 254 | else: |
| 255 | output.error = response.reason or "" |
| 256 | output.success = False |
| 257 | except Exception: |
| 258 | output.success = False |
| 259 | exc_info = sys.exc_info() |
| 260 | output.error = "".join(traceback.format_exception(*exc_info)) |
| 261 | |
| 262 | if pbar: |
| 263 | pbar.update(1) |
| 264 | return output |
| 265 | |
| 266 | |
| 267 | async def async_request_openai_completions( |
| 268 | request_func_input: RequestFuncInput, |
| 269 | pbar: tqdm | None = None, |
| 270 | ) -> RequestFuncOutput: |
| 271 | api_url = request_func_input.api_url |
| 272 | assert api_url.endswith(("completions", "profile")), ( |
| 273 | "OpenAI Completions API URL must end with 'completions' or 'profile'." |
| 274 | ) |
| 275 | |
| 276 | async with aiohttp.ClientSession( |
| 277 | trust_env=True, timeout=AIOHTTP_TIMEOUT |
| 278 | ) as session: |
| 279 | payload = { |
| 280 | "model": request_func_input.model_name |
| 281 | if request_func_input.model_name |
| 282 | else request_func_input.model, |
| 283 | "prompt": request_func_input.prompt, |
| 284 | "temperature": 0.0, |
| 285 | "repetition_penalty": 1.0, |
| 286 | "max_tokens": request_func_input.output_len, |
| 287 | "logprobs": request_func_input.logprobs, |
| 288 | "stream": True, |
| 289 | "stream_options": { |
| 290 | "include_usage": True, |
| 291 | }, |
| 292 | } |
| 293 | if request_func_input.ignore_eos: |
| 294 | payload["ignore_eos"] = request_func_input.ignore_eos |
| 295 | if request_func_input.extra_body: |
| 296 | payload.update(request_func_input.extra_body) |
| 297 | headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"} |
| 298 | if request_func_input.request_id: |
| 299 | headers["x-request-id"] = request_func_input.request_id |
| 300 | |
| 301 | output = RequestFuncOutput() |
| 302 | output.prompt_len = request_func_input.prompt_len |
| 303 | |
| 304 | generated_text = "" |
| 305 | st = time.perf_counter() |
| 306 | most_recent_timestamp = st |
| 307 | try: |
| 308 | async with session.post( |
| 309 | url=api_url, json=payload, headers=headers |
| 310 | ) as response: |
| 311 | if response.status == 200: |
| 312 | first_chunk_received = False |
| 313 | async for chunk_bytes in response.content: |
| 314 | chunk_bytes = chunk_bytes.strip() |
| 315 | if not chunk_bytes: |
| 316 | continue |
| 317 | |
| 318 | chunk = chunk_bytes.decode("utf-8").removeprefix("data: ") |
| 319 | if chunk != "[DONE]": |
| 320 | data = json.loads(chunk) |
| 321 | |
| 322 | # NOTE: Some completion API might have a last |
| 323 | # usage summary response without a token so we |
| 324 | # want to check a token was generated |
| 325 | if choices := data.get("choices"): |
| 326 | # Note that text could be empty here |
| 327 | # e.g. for special tokens |
| 328 | text = choices[0].get("text") |
| 329 | timestamp = time.perf_counter() |
| 330 | # First token |
| 331 | if not first_chunk_received: |
| 332 | first_chunk_received = True |
| 333 | output.ttft = timestamp - st |
| 334 | |
| 335 | # Decoding phase |
| 336 | else: |
| 337 | output.itl.append(timestamp - most_recent_timestamp) |
| 338 | |
| 339 | most_recent_timestamp = timestamp |
| 340 | generated_text += text or "" |
| 341 | if usage := data.get("usage"): |
| 342 | output.output_tokens = usage.get("completion_tokens") |
| 343 | if first_chunk_received: |
| 344 | output.success = True |
| 345 | else: |
| 346 | output.success = False |
| 347 | output.error = ( |
| 348 | "Never received a valid chunk to calculate TTFT." |
| 349 | "This response will be marked as failed!" |
| 350 | ) |
| 351 | output.generated_text = generated_text |
| 352 | output.latency = most_recent_timestamp - st |
| 353 | else: |
| 354 | output.error = response.reason or "" |
| 355 | output.success = False |
| 356 | except Exception: |
| 357 | output.success = False |
| 358 | exc_info = sys.exc_info() |
| 359 | output.error = "".join(traceback.format_exception(*exc_info)) |
| 360 | |
| 361 | if pbar: |
| 362 | pbar.update(1) |
| 363 | return output |
| 364 | |
| 365 | |
| 366 | async def async_request_openai_chat_completions( |
| 367 | request_func_input: RequestFuncInput, |
| 368 | pbar: tqdm | None = None, |
| 369 | ) -> RequestFuncOutput: |
| 370 | api_url = request_func_input.api_url |
| 371 | assert api_url.endswith(("chat/completions", "profile")), ( |
| 372 | "OpenAI Chat Completions API URL must end with 'chat/completions'." |
| 373 | ) |
| 374 | |
| 375 | async with aiohttp.ClientSession( |
| 376 | trust_env=True, timeout=AIOHTTP_TIMEOUT |
| 377 | ) as session: |
| 378 | content = [{"type": "text", "text": request_func_input.prompt}] |
| 379 | if request_func_input.multi_modal_content: |
| 380 | mm_content = request_func_input.multi_modal_content |
| 381 | if isinstance(mm_content, list): |
| 382 | content.extend(mm_content) |
| 383 | elif isinstance(mm_content, dict): |
| 384 | content.append(mm_content) |
| 385 | else: |
| 386 | raise TypeError( |
| 387 | "multi_modal_content must be a dict or list[dict] for openai-chat" |
| 388 | ) |
| 389 | payload = { |
| 390 | "model": request_func_input.model_name |
| 391 | if request_func_input.model_name |
| 392 | else request_func_input.model, |
| 393 | "messages": [ |
| 394 | {"role": "user", "content": content}, |
| 395 | ], |
| 396 | "temperature": 0.0, |
| 397 | "max_completion_tokens": request_func_input.output_len, |
| 398 | "stream": True, |
| 399 | "stream_options": { |
| 400 | "include_usage": True, |
| 401 | }, |
| 402 | } |
| 403 | if request_func_input.ignore_eos: |
| 404 | payload["ignore_eos"] = request_func_input.ignore_eos |
| 405 | if request_func_input.extra_body: |
| 406 | payload.update(request_func_input.extra_body) |
| 407 | headers = { |
| 408 | "Content-Type": "application/json", |
| 409 | "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", |
| 410 | } |
| 411 | if request_func_input.request_id: |
| 412 | headers["x-request-id"] = request_func_input.request_id |
| 413 | |
| 414 | output = RequestFuncOutput() |
| 415 | output.prompt_len = request_func_input.prompt_len |
| 416 | |
| 417 | generated_text = "" |
| 418 | first_chunk_received = False |
| 419 | st = time.perf_counter() |
| 420 | most_recent_timestamp = st |
| 421 | try: |
| 422 | async with session.post( |
| 423 | url=api_url, json=payload, headers=headers |
| 424 | ) as response: |
| 425 | if response.status == 200: |
| 426 | async for chunk_bytes in response.content: |
| 427 | chunk_bytes = chunk_bytes.strip() |
| 428 | if not chunk_bytes: |
| 429 | continue |
| 430 | chunk_bytes = chunk_bytes.decode("utf-8") |
| 431 | # NOTE: SSE comments (often used as pings) start with a colon. |
| 432 | # These are not JSON data payload and should be skipped. |
| 433 | if chunk_bytes.startswith(":"): |
| 434 | continue |
| 435 | |
| 436 | chunk = chunk_bytes.removeprefix("data: ") |
| 437 | |
| 438 | if chunk != "[DONE]": |
| 439 | timestamp = time.perf_counter() |
| 440 | data = json.loads(chunk) |
| 441 | |
| 442 | if choices := data.get("choices"): |
| 443 | content = choices[0]["delta"].get("content") |
| 444 | # First token |
| 445 | if not first_chunk_received: |
| 446 | first_chunk_received = True |
| 447 | output.ttft = timestamp - st |
| 448 | |
| 449 | # Decoding phase |
| 450 | else: |
| 451 | output.itl.append(timestamp - most_recent_timestamp) |
| 452 | |
| 453 | generated_text += content or "" |
| 454 | # Only token chunks advance the request end; |
| 455 | # the trailing usage chunk carries no token. |
| 456 | most_recent_timestamp = timestamp |
| 457 | elif usage := data.get("usage"): |
| 458 | output.output_tokens = usage.get("completion_tokens") |
| 459 | |
| 460 | output.generated_text = generated_text |
| 461 | if first_chunk_received: |
| 462 | output.success = True |
| 463 | else: |
| 464 | output.success = False |
| 465 | output.error = ( |
| 466 | "Never received a valid chunk to calculate TTFT." |
| 467 | "This response will be marked as failed!" |
| 468 | ) |
| 469 | output.latency = most_recent_timestamp - st |
| 470 | else: |
| 471 | output.error = response.reason or "" |
| 472 | output.success = False |
| 473 | except Exception: |
| 474 | output.success = False |
| 475 | exc_info = sys.exc_info() |
| 476 | output.error = "".join(traceback.format_exception(*exc_info)) |
| 477 | |
| 478 | if pbar: |
| 479 | pbar.update(1) |
| 480 | return output |
| 481 | |
| 482 | |
| 483 | async def async_request_openai_audio( |
| 484 | request_func_input: RequestFuncInput, |
| 485 | pbar: tqdm | None = None, |
| 486 | ) -> RequestFuncOutput: |
| 487 | # Lazy import without PlaceholderModule to avoid vllm dep. |
| 488 | import soundfile |
| 489 | |
| 490 | api_url = request_func_input.api_url |
| 491 | assert api_url.endswith(("transcriptions", "translations")), ( |
| 492 | "OpenAI Chat Completions API URL must end with 'transcriptions' " |
| 493 | ) |
| 494 | "or `translations`." |
| 495 | |
| 496 | async with aiohttp.ClientSession( |
| 497 | trust_env=True, timeout=AIOHTTP_TIMEOUT |
| 498 | ) as session: |
| 499 | content = [{"type": "text", "text": request_func_input.prompt}] |
| 500 | payload = { |
| 501 | "model": request_func_input.model_name |
| 502 | if request_func_input.model_name |
| 503 | else request_func_input.model, |
| 504 | "temperature": 0.0, |
| 505 | "max_completion_tokens": request_func_input.output_len, |
| 506 | "stream": True, |
| 507 | "language": "en", |
| 508 | # Flattened due to multipart/form-data |
| 509 | "stream_include_usage": True, |
| 510 | "stream_continuous_usage_stats": True, |
| 511 | } |
| 512 | if request_func_input.extra_body: |
| 513 | payload.update(request_func_input.extra_body) |
| 514 | headers = { |
| 515 | "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", |
| 516 | } |
| 517 | if request_func_input.request_id: |
| 518 | headers["x-request-id"] = request_func_input.request_id |
| 519 | |
| 520 | # Send audio file |
| 521 | def to_bytes(y, sr): |
| 522 | buffer = io.BytesIO() |
| 523 | soundfile.write(buffer, y, sr, format="WAV") |
| 524 | buffer.seek(0) |
| 525 | return buffer |
| 526 | |
| 527 | mm_audio = request_func_input.multi_modal_content |
| 528 | if not isinstance(mm_audio, dict) or "audio" not in mm_audio: |
| 529 | raise TypeError("multi_modal_content must be a dict containing 'audio'") |
| 530 | with to_bytes(*mm_audio["audio"]) as f: |
| 531 | form = aiohttp.FormData() |
| 532 | form.add_field("file", f, content_type="audio/wav") |
| 533 | for key, value in payload.items(): |
| 534 | form.add_field(key, str(value)) |
| 535 | |
| 536 | output = RequestFuncOutput() |
| 537 | output.prompt_len = request_func_input.prompt_len |
| 538 | |
| 539 | generated_text = "" |
| 540 | first_chunk_received = False |
| 541 | st = time.perf_counter() |
| 542 | most_recent_timestamp = st |
| 543 | try: |
| 544 | async with session.post( |
| 545 | url=api_url, data=form, headers=headers |
| 546 | ) as response: |
| 547 | if response.status == 200: |
| 548 | async for chunk_bytes in response.content: |
| 549 | chunk_bytes = chunk_bytes.strip() |
| 550 | if not chunk_bytes: |
| 551 | continue |
| 552 | |
| 553 | chunk = chunk_bytes.decode("utf-8").removeprefix("data: ") |
| 554 | if chunk != "[DONE]": |
| 555 | timestamp = time.perf_counter() |
| 556 | data = json.loads(chunk) |
| 557 | |
| 558 | if choices := data.get("choices"): |
| 559 | content = choices[0]["delta"].get("content") |
| 560 | # First token |
| 561 | if not first_chunk_received: |
| 562 | first_chunk_received = True |
| 563 | output.ttft = timestamp - st |
| 564 | |
| 565 | # Decoding phase |
| 566 | else: |
| 567 | output.itl.append( |
| 568 | timestamp - most_recent_timestamp |
| 569 | ) |
| 570 | |
| 571 | generated_text += content or "" |
| 572 | # Only token chunks advance the request |
| 573 | # end; the trailing usage chunk carries no |
| 574 | # token. |
| 575 | most_recent_timestamp = timestamp |
| 576 | elif usage := data.get("usage"): |
| 577 | output.output_tokens = usage.get( |
| 578 | "completion_tokens" |
| 579 | ) |
| 580 | |
| 581 | output.generated_text = generated_text |
| 582 | if first_chunk_received: |
| 583 | output.success = True |
| 584 | else: |
| 585 | output.success = False |
| 586 | output.error = ( |
| 587 | "Never received a valid chunk to calculate TTFT." |
| 588 | "This response will be marked as failed!" |
| 589 | ) |
| 590 | output.latency = most_recent_timestamp - st |
| 591 | else: |
| 592 | output.error = response.reason or "" |
| 593 | output.success = False |
| 594 | except Exception: |
| 595 | output.success = False |
| 596 | exc_info = sys.exc_info() |
| 597 | output.error = "".join(traceback.format_exception(*exc_info)) |
| 598 | |
| 599 | if pbar: |
| 600 | pbar.update(1) |
| 601 | return output |
| 602 | |
| 603 | |
| 604 | def get_model(pretrained_model_name_or_path: str) -> str: |
| 605 | if os.getenv("VLLM_USE_MODELSCOPE", "False").lower() == "true": |
| 606 | from modelscope import snapshot_download |
| 607 | |
| 608 | from vllm.model_executor.model_loader.weight_utils import get_lock |
| 609 | |
| 610 | # Use file lock to prevent multiple processes from |
| 611 | # downloading the same model weights at the same time. |
| 612 | with get_lock(pretrained_model_name_or_path): |
| 613 | model_path = snapshot_download( |
| 614 | model_id=pretrained_model_name_or_path, |
| 615 | local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, |
| 616 | ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"], |
| 617 | ) |
| 618 | |
| 619 | return model_path |
| 620 | return pretrained_model_name_or_path |
| 621 | |
| 622 | |
| 623 | def get_tokenizer( |
| 624 | pretrained_model_name_or_path: str, |
| 625 | tokenizer_mode: str = "auto", |
| 626 | trust_remote_code: bool = False, |
| 627 | **kwargs, |
| 628 | ) -> PythonBackend | TokenizersBackend: |
| 629 | if pretrained_model_name_or_path is not None and not os.path.exists( |
| 630 | pretrained_model_name_or_path |
| 631 | ): |
| 632 | pretrained_model_name_or_path = get_model(pretrained_model_name_or_path) |
| 633 | if tokenizer_mode == "mistral": |
| 634 | try: |
| 635 | from vllm.tokenizers.mistral import MistralTokenizer |
| 636 | except ImportError as e: |
| 637 | raise ImportError( |
| 638 | "MistralTokenizer requires vllm package.\n" |
| 639 | "Please install it with `pip install vllm` " |
| 640 | "to use mistral tokenizer mode." |
| 641 | ) from e |
| 642 | return MistralTokenizer.from_pretrained(str(pretrained_model_name_or_path)) |
| 643 | else: |
| 644 | return AutoTokenizer.from_pretrained( |
| 645 | pretrained_model_name_or_path, |
| 646 | trust_remote_code=trust_remote_code, |
| 647 | **kwargs, |
| 648 | ) |
| 649 | |
| 650 | |
| 651 | ASYNC_REQUEST_FUNCS = { |
| 652 | "tgi": async_request_tgi, |
| 653 | "vllm": async_request_openai_completions, |
| 654 | "lmdeploy": async_request_openai_completions, |
| 655 | "deepspeed-mii": async_request_deepspeed_mii, |
| 656 | "openai": async_request_openai_completions, |
| 657 | "openai-chat": async_request_openai_chat_completions, |
| 658 | "openai-audio": async_request_openai_audio, |
| 659 | "tensorrt-llm": async_request_trt_llm, |
| 660 | "scalellm": async_request_openai_completions, |
| 661 | "sglang": async_request_openai_completions, |
| 662 | "llama.cpp": async_request_openai_completions, |
| 663 | } |