| 1 | #!/usr/bin/env python3 |
| 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | """Isolated GSM8K evaluation script for vLLM serve endpoint.""" |
| 5 | |
| 6 | import argparse |
| 7 | import ast |
| 8 | import asyncio |
| 9 | import json |
| 10 | import os |
| 11 | import tempfile |
| 12 | import time |
| 13 | from collections.abc import Generator |
| 14 | |
| 15 | import aiohttp |
| 16 | import numpy as np |
| 17 | import regex as re |
| 18 | import requests |
| 19 | from tqdm.asyncio import tqdm |
| 20 | |
| 21 | from vllm.assets.base import VLLM_S3_BUCKET_URL |
| 22 | |
| 23 | INVALID = -9999999 |
| 24 | |
| 25 | |
| 26 | def download_and_cache_file(url: str, filename: str | None = None) -> str: |
| 27 | """Download and cache a file from a URL.""" |
| 28 | if filename is None: |
| 29 | filename = os.path.join(tempfile.gettempdir(), url.split("/")[-1]) |
| 30 | |
| 31 | if os.path.exists(filename): |
| 32 | return filename |
| 33 | |
| 34 | print(f"Downloading from {url} to {filename}") |
| 35 | response = requests.get(url, stream=True) |
| 36 | response.raise_for_status() |
| 37 | |
| 38 | with open(filename, "wb") as f: |
| 39 | for chunk in response.iter_content(chunk_size=1024): |
| 40 | f.write(chunk) |
| 41 | |
| 42 | return filename |
| 43 | |
| 44 | |
| 45 | def load_gsm8k_data() -> tuple[list[dict], list[dict]]: |
| 46 | """Load GSM8K train and test data.""" |
| 47 | train_url = f"{VLLM_S3_BUCKET_URL}/ci-datasets/gsm8k/train.jsonl" |
| 48 | test_url = f"{VLLM_S3_BUCKET_URL}/ci-datasets/gsm8k/test.jsonl" |
| 49 | |
| 50 | train_file = download_and_cache_file(train_url) |
| 51 | test_file = download_and_cache_file(test_url) |
| 52 | |
| 53 | train_data = list(read_jsonl(train_file)) |
| 54 | test_data = list(read_jsonl(test_file)) |
| 55 | |
| 56 | return train_data, test_data |
| 57 | |
| 58 | |
| 59 | def read_jsonl(filename: str) -> Generator[dict, None, None]: |
| 60 | """Read a JSONL file.""" |
| 61 | with open(filename) as fin: |
| 62 | for line in fin: |
| 63 | if not line.startswith("#"): |
| 64 | yield json.loads(line) |
| 65 | |
| 66 | |
| 67 | def get_answer_value(answer_str: str) -> int: |
| 68 | """Extract the numerical answer from the response.""" |
| 69 | answer_str = answer_str.replace(",", "") |
| 70 | numbers = re.findall(r"\d+", answer_str) |
| 71 | if len(numbers) < 1: |
| 72 | return INVALID |
| 73 | try: |
| 74 | return ast.literal_eval(numbers[-1]) |
| 75 | except SyntaxError: |
| 76 | return INVALID |
| 77 | |
| 78 | |
| 79 | async def call_vllm_api( |
| 80 | session: aiohttp.ClientSession, |
| 81 | prompt: str, |
| 82 | temperature: float, |
| 83 | max_tokens: int, |
| 84 | stop: list[str] | None = None, |
| 85 | url: str | None = None, |
| 86 | seed: int | None = None, |
| 87 | ) -> tuple[str, int]: |
| 88 | """Call vLLM's OpenAI-compatible completions endpoint. |
| 89 | |
| 90 | Returns: |
| 91 | Tuple of (response_text, completion_tokens) |
| 92 | |
| 93 | """ |
| 94 | data = { |
| 95 | "prompt": prompt, |
| 96 | "temperature": temperature, |
| 97 | "max_tokens": max_tokens, |
| 98 | "stop": stop, |
| 99 | } |
| 100 | if seed is not None: |
| 101 | data["seed"] = seed |
| 102 | |
| 103 | try: |
| 104 | async with session.post(f"{url}/v1/completions", json=data) as response: |
| 105 | response.raise_for_status() |
| 106 | result = await response.json() |
| 107 | text = result["choices"][0]["text"] |
| 108 | completion_tokens = result.get("usage", {}).get("completion_tokens", 0) |
| 109 | return text, completion_tokens |
| 110 | except Exception as e: |
| 111 | print(f"Error calling vLLM API ({type(e).__name__}): {e}") |
| 112 | return "", 0 |
| 113 | |
| 114 | |
| 115 | async def call_vllm_chat_api( |
| 116 | session: aiohttp.ClientSession, |
| 117 | model: str, |
| 118 | prompt: str, |
| 119 | temperature: float, |
| 120 | max_tokens: int, |
| 121 | stop: list[str] | None = None, |
| 122 | url: str | None = None, |
| 123 | seed: int | None = None, |
| 124 | ) -> tuple[str, int]: |
| 125 | """Call vLLM's OpenAI-compatible chat completions endpoint.""" |
| 126 | data = { |
| 127 | "model": model, |
| 128 | "messages": [{"role": "user", "content": prompt}], |
| 129 | "temperature": temperature, |
| 130 | "max_tokens": max_tokens, |
| 131 | "stop": stop, |
| 132 | } |
| 133 | if seed is not None: |
| 134 | data["seed"] = seed |
| 135 | |
| 136 | try: |
| 137 | async with session.post(f"{url}/v1/chat/completions", json=data) as response: |
| 138 | response.raise_for_status() |
| 139 | result = await response.json() |
| 140 | text = result["choices"][0]["message"]["content"] or "" |
| 141 | completion_tokens = result.get("usage", {}).get("completion_tokens", 0) |
| 142 | return text, completion_tokens |
| 143 | except Exception as e: |
| 144 | print(f"Error calling vLLM chat API ({type(e).__name__}): {e}") |
| 145 | return "", 0 |
| 146 | |
| 147 | |
| 148 | def _build_gsm8k_prompts( |
| 149 | num_questions: int = 1319, |
| 150 | num_shots: int = 5, |
| 151 | gen_prefix: str = "", |
| 152 | ) -> tuple[list[str], list[int]]: |
| 153 | """Build few-shot GSM8K completion prompts and ground-truth labels.""" |
| 154 | if num_questions == 0: |
| 155 | return [], [] |
| 156 | train_data, test_data = load_gsm8k_data() |
| 157 | num_questions = min(num_questions, len(test_data)) |
| 158 | |
| 159 | few_shot_examples = "" |
| 160 | for i in range(num_shots): |
| 161 | few_shot_examples += ( |
| 162 | f"Question: {train_data[i]['question']}\n" |
| 163 | f"Answer:{gen_prefix} {train_data[i]['answer']}\n\n" |
| 164 | ) |
| 165 | |
| 166 | prompts = [] |
| 167 | labels = [] |
| 168 | for i in range(num_questions): |
| 169 | prompts.append( |
| 170 | few_shot_examples |
| 171 | + f"Question: {test_data[i]['question']}\nAnswer:{gen_prefix}" |
| 172 | ) |
| 173 | labels.append(get_answer_value(test_data[i]["answer"])) |
| 174 | |
| 175 | assert all(label != INVALID for label in labels), "Some labels are invalid" |
| 176 | return prompts, labels |
| 177 | |
| 178 | |
| 179 | def _score_gsm8k( |
| 180 | states: list[str], |
| 181 | output_tokens: list[int], |
| 182 | labels: list[int], |
| 183 | num_shots: int, |
| 184 | max_tokens: int, |
| 185 | latency: float, |
| 186 | ) -> dict[str, float | int]: |
| 187 | """Score GSM8K responses and return a results dict.""" |
| 188 | num_questions = len(labels) |
| 189 | preds = [get_answer_value(state) for state in states] |
| 190 | accuracy = np.mean(np.array(preds) == np.array(labels)) |
| 191 | invalid_rate = np.mean(np.array(preds) == INVALID) |
| 192 | total_output_tokens = sum(output_tokens) |
| 193 | tokens_per_second = total_output_tokens / latency if latency > 0 else 0.0 |
| 194 | |
| 195 | return { |
| 196 | "accuracy": accuracy, |
| 197 | "invalid_rate": invalid_rate, |
| 198 | "latency": latency, |
| 199 | "questions_per_second": num_questions / latency if latency > 0 else 0.0, |
| 200 | "total_output_tokens": total_output_tokens, |
| 201 | "tokens_per_second": tokens_per_second, |
| 202 | "num_questions": num_questions, |
| 203 | "num_shots": num_shots, |
| 204 | "max_tokens": max_tokens, |
| 205 | "timestamp": time.time(), |
| 206 | } |
| 207 | |
| 208 | |
| 209 | def evaluate_gsm8k( |
| 210 | num_questions: int = 1319, |
| 211 | num_shots: int = 5, |
| 212 | max_tokens: int = 256, |
| 213 | model: str | None = None, |
| 214 | use_chat_completions: bool = False, |
| 215 | host: str = "http://127.0.0.1", |
| 216 | port: int = 8000, |
| 217 | temperature: float = 0.0, |
| 218 | seed: int | None = 42, |
| 219 | request_timeout_seconds: float = 600, |
| 220 | gen_prefix: str = "", |
| 221 | max_concurrency: int | None = None, |
| 222 | ) -> dict[str, float | int]: |
| 223 | """Evaluate GSM8K accuracy using vLLM serve endpoint. |
| 224 | |
| 225 | Returns dict with accuracy, invalid_rate, latency, etc. |
| 226 | """ |
| 227 | base_url = f"{host}:{port}" |
| 228 | prompts, labels = _build_gsm8k_prompts(num_questions, num_shots, gen_prefix) |
| 229 | num_questions = len(prompts) |
| 230 | |
| 231 | async def run_async_evaluation(): |
| 232 | states: list[str] = [""] * num_questions |
| 233 | output_tokens: list[int] = [0] * num_questions |
| 234 | |
| 235 | async def get_answer(session: aiohttp.ClientSession, i: int) -> tuple[str, int]: |
| 236 | stop = ["Question", "Assistant:", "<|separator|>"] |
| 237 | if use_chat_completions: |
| 238 | if model is None: |
| 239 | raise ValueError("model is required for chat completions") |
| 240 | answer, tokens = await call_vllm_chat_api( |
| 241 | session=session, |
| 242 | model=model, |
| 243 | prompt=prompts[i], |
| 244 | temperature=temperature, |
| 245 | max_tokens=max_tokens, |
| 246 | stop=stop, |
| 247 | url=base_url, |
| 248 | seed=seed, |
| 249 | ) |
| 250 | else: |
| 251 | answer, tokens = await call_vllm_api( |
| 252 | session=session, |
| 253 | prompt=prompts[i], |
| 254 | temperature=temperature, |
| 255 | max_tokens=max_tokens, |
| 256 | stop=stop, |
| 257 | url=base_url, |
| 258 | seed=seed, |
| 259 | ) |
| 260 | states[i] = answer |
| 261 | output_tokens[i] = tokens |
| 262 | return answer, tokens |
| 263 | |
| 264 | timeout = aiohttp.ClientTimeout(total=request_timeout_seconds) |
| 265 | connector = ( |
| 266 | aiohttp.TCPConnector(limit=max_concurrency) |
| 267 | if max_concurrency is not None |
| 268 | else None |
| 269 | ) |
| 270 | async with aiohttp.ClientSession( |
| 271 | timeout=timeout, connector=connector |
| 272 | ) as session: |
| 273 | tasks = [get_answer(session, i) for i in range(num_questions)] |
| 274 | await tqdm.gather(*tasks, desc="Evaluating") |
| 275 | |
| 276 | return states, output_tokens |
| 277 | |
| 278 | print(f"Running GSM8K evaluation: {num_questions} questions, {num_shots}-shot") |
| 279 | |
| 280 | tic = time.perf_counter() |
| 281 | states, output_tokens = asyncio.run(run_async_evaluation()) |
| 282 | latency = time.perf_counter() - tic |
| 283 | |
| 284 | return _score_gsm8k(states, output_tokens, labels, num_shots, max_tokens, latency) |
| 285 | |
| 286 | |
| 287 | def evaluate_gsm8k_offline( |
| 288 | llm, |
| 289 | num_questions: int = 1319, |
| 290 | num_shots: int = 5, |
| 291 | max_tokens: int = 256, |
| 292 | temperature: float = 0.0, |
| 293 | gen_prefix: str = "", |
| 294 | use_chat_completions: bool = False, |
| 295 | chat_template_kwargs: dict[str, object] | None = None, |
| 296 | ) -> dict[str, float | int]: |
| 297 | """Evaluate GSM8K accuracy using an offline vllm.LLM object. |
| 298 | |
| 299 | Same prompts and scoring as evaluate_gsm8k(), but runs generation |
| 300 | directly via llm.generate() instead of calling a server over HTTP. |
| 301 | |
| 302 | When ``use_chat_completions=True``, prompts go through the chat template via |
| 303 | ``llm.chat()`` instead of raw completion (for instruction-tuned models). |
| 304 | ``chat_template_kwargs`` are forwarded to ``llm.chat()`` when provided. |
| 305 | """ |
| 306 | from vllm import SamplingParams |
| 307 | |
| 308 | prompts, labels = _build_gsm8k_prompts(num_questions, num_shots, gen_prefix) |
| 309 | sampling_params = SamplingParams( |
| 310 | temperature=temperature, |
| 311 | max_tokens=max_tokens, |
| 312 | stop=["Question", "Assistant:", "<|separator|>"], |
| 313 | ) |
| 314 | mode = "chat" if use_chat_completions else "completion" |
| 315 | print( |
| 316 | f"Running offline GSM8K evaluation: {len(prompts)} questions, " |
| 317 | f"{num_shots}-shot, {mode}" |
| 318 | ) |
| 319 | |
| 320 | tic = time.perf_counter() |
| 321 | if use_chat_completions: |
| 322 | conversations = [[{"role": "user", "content": p}] for p in prompts] |
| 323 | outputs = llm.chat( |
| 324 | conversations, |
| 325 | sampling_params, |
| 326 | chat_template_kwargs=chat_template_kwargs, |
| 327 | ) |
| 328 | else: |
| 329 | outputs = llm.generate(prompts, sampling_params) |
| 330 | latency = time.perf_counter() - tic |
| 331 | |
| 332 | states = [o.outputs[0].text for o in outputs] |
| 333 | output_tokens = [len(o.outputs[0].token_ids) for o in outputs] |
| 334 | |
| 335 | return _score_gsm8k(states, output_tokens, labels, num_shots, max_tokens, latency) |
| 336 | |
| 337 | |
| 338 | def main() -> None: |
| 339 | parser = argparse.ArgumentParser(description="GSM8K evaluation for vLLM serve") |
| 340 | parser.add_argument( |
| 341 | "--num-shots", type=int, default=5, help="Number of few-shot examples" |
| 342 | ) |
| 343 | parser.add_argument( |
| 344 | "--num-questions", |
| 345 | type=int, |
| 346 | default=1319, |
| 347 | help="Number of questions to evaluate", |
| 348 | ) |
| 349 | parser.add_argument( |
| 350 | "--max-tokens", type=int, default=256, help="Max tokens for generation" |
| 351 | ) |
| 352 | parser.add_argument("--host", type=str, default="http://127.0.0.1", help="Host URL") |
| 353 | parser.add_argument("--port", type=int, default=8000, help="Port number") |
| 354 | parser.add_argument( |
| 355 | "--temperature", type=float, default=0.0, help="Temperature for generation" |
| 356 | ) |
| 357 | parser.add_argument( |
| 358 | "--seed", type=int, default=42, help="Random seed for reproducibility" |
| 359 | ) |
| 360 | parser.add_argument( |
| 361 | "--max-concurrency", |
| 362 | type=int, |
| 363 | help="Maximum number of concurrent requests", |
| 364 | ) |
| 365 | parser.add_argument( |
| 366 | "--request-timeout-seconds", |
| 367 | type=float, |
| 368 | default=600, |
| 369 | help="Timeout for each request, including time waiting for a connection", |
| 370 | ) |
| 371 | parser.add_argument("--save-results", type=str, help="Save results to JSON file") |
| 372 | |
| 373 | args = parser.parse_args() |
| 374 | |
| 375 | result = evaluate_gsm8k( |
| 376 | num_questions=args.num_questions, |
| 377 | num_shots=args.num_shots, |
| 378 | max_tokens=args.max_tokens, |
| 379 | host=args.host, |
| 380 | port=args.port, |
| 381 | temperature=args.temperature, |
| 382 | seed=args.seed, |
| 383 | max_concurrency=args.max_concurrency, |
| 384 | request_timeout_seconds=args.request_timeout_seconds, |
| 385 | ) |
| 386 | |
| 387 | # Print results to terminal |
| 388 | print("\nResults:") |
| 389 | print(f"Accuracy: {result['accuracy']:.3f}") |
| 390 | print(f"Invalid responses: {result['invalid_rate']:.3f}") |
| 391 | print(f"Total latency: {result['latency']:.3f} s") |
| 392 | print(f"Questions per second: {result['questions_per_second']:.3f}") |
| 393 | print(f"Total output tokens: {result['total_output_tokens']}") |
| 394 | print(f"Output tokens per second: {result['tokens_per_second']:.3f}") |
| 395 | |
| 396 | # Optional file saving |
| 397 | if args.save_results: |
| 398 | with open(args.save_results, "w") as f: |
| 399 | json.dump(result, f, indent=2) |
| 400 | print(f"Results saved to {args.save_results}") |
| 401 | |
| 402 | |
| 403 | if __name__ == "__main__": |
| 404 | main() |