A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 05:45 UTC
Add output token counting to gsm8k eval (#28594)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Nov 14, 2025 at 04:32 UTC
c9a3a02149d83cc2840769228c4e591d39351bb6
1 file changed
+23
-9
tests/evals/gsm8k/gsm8k_eval.py
+23
-9
@@ -83,8 +83,12 @@ async def call_vllm_api(
83
stop: list[str] | None = None,
84
url: str | None = None,
85
seed: int | None = None,
86
-) -> str:
87
- """Call vLLM's OpenAI-compatible completions endpoint."""
86
+) -> tuple[str, int]:
87
+ """Call vLLM's OpenAI-compatible completions endpoint.
88
+
89
+ Returns:
90
+ Tuple of (response_text, completion_tokens)
91
+ """
92
data = {
93
"prompt": prompt,
94
"temperature": temperature,
@@ -98,10 +102,12 @@ async def call_vllm_api(
102
async with session.post(f"{url}/v1/completions", json=data) as response:
103
response.raise_for_status()
104
result = await response.json()
101
- return result["choices"][0]["text"]
105
+ text = result["choices"][0]["text"]
106
+ completion_tokens = result.get("usage", {}).get("completion_tokens", 0)
107
+ return text, completion_tokens
108
except Exception as e:
109
print(f"Error calling vLLM API: {e}")
104
- return ""
110
+ return "", 0
111
112
113
def evaluate_gsm8k(
@@ -146,10 +152,11 @@ def evaluate_gsm8k(
152
# Run evaluation
153
async def run_async_evaluation():
154
states: list[str] = [""] * num_questions
155
+ output_tokens: list[int] = [0] * num_questions
156
150
- async def get_answer(session: aiohttp.ClientSession, i: int) -> str:
157
+ async def get_answer(session: aiohttp.ClientSession, i: int) -> tuple[str, int]:
158
prompt = few_shot_examples + questions[i]
152
- answer = await call_vllm_api(
159
+ answer, tokens = await call_vllm_api(
160
session=session,
161
prompt=prompt,
162
temperature=temperature,
@@ -159,7 +166,8 @@ def evaluate_gsm8k(
166
seed=seed,
167
)
168
states[i] = answer
162
- return answer
169
+ output_tokens[i] = tokens
170
+ return answer, tokens
171
172
async with aiohttp.ClientSession(
173
timeout=aiohttp.ClientTimeout(total=600)
@@ -167,24 +175,28 @@ def evaluate_gsm8k(
175
tasks = [get_answer(session, i) for i in range(num_questions)]
176
await tqdm.gather(*tasks, desc="Evaluating")
177
170
- return states
178
+ return states, output_tokens
179
180
print(f"Running GSM8K evaluation: {num_questions} questions, {num_shots}-shot")
181
182
tic = time.perf_counter()
175
- states = asyncio.run(run_async_evaluation())
183
+ states, output_tokens = asyncio.run(run_async_evaluation())
184
latency = time.perf_counter() - tic
185
186
# Compute metrics
187
preds = [get_answer_value(state) for state in states]
188
accuracy = np.mean(np.array(preds) == np.array(labels))
189
invalid_rate = np.mean(np.array(preds) == INVALID)
190
+ total_output_tokens = sum(output_tokens)
191
+ tokens_per_second = total_output_tokens / latency if latency > 0 else 0.0
192
193
result = {
194
"accuracy": accuracy,
195
"invalid_rate": invalid_rate,
196
"latency": latency,
197
"questions_per_second": num_questions / latency,
198
+ "total_output_tokens": total_output_tokens,
199
+ "tokens_per_second": tokens_per_second,
200
"num_questions": num_questions,
201
"num_shots": num_shots,
202
"max_tokens": max_tokens,
@@ -236,6 +248,8 @@ def main() -> None:
248
print(f"Invalid responses: {result['invalid_rate']:.3f}")
249
print(f"Total latency: {result['latency']:.3f} s")
250
print(f"Questions per second: {result['questions_per_second']:.3f}")
251
+ print(f"Total output tokens: {result['total_output_tokens']}")
252
+ print(f"Output tokens per second: {result['tokens_per_second']:.3f}")
253
254
# Optional file saving
255
if args.save_results: