main
py 67 lines 1.83 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import os
4
5 import lm_eval
6 import openai
7
8 BASE_URL = f"http://localhost:{os.environ.get('PROXY_PORT', 8192)}/v1"
9 NUM_CONCURRENT = 100
10 TASK = "gsm8k"
11 FILTER = "exact_match,strict-match"
12 RTOL = 0.05
13
14 # Model-specific expected values
15 EXPECTED_VALUES = {
16 "Qwen/Qwen3-0.6B": 0.41,
17 }
18
19 SIMPLE_PROMPT = (
20 "The best part about working on vLLM is that I got to meet so many people across "
21 "various different organizations like UCB, Google, and Meta which means",
22 )
23
24 # Get model name from environment variable
25 MODEL_NAME = os.environ.get("TEST_MODEL", "Qwen/Qwen3-0.6B")
26
27
28 def run_simple_prompt():
29 client = openai.OpenAI(api_key="EMPTY", base_url=BASE_URL)
30 completion = client.completions.create(model=MODEL_NAME, prompt=SIMPLE_PROMPT)
31
32 print("-" * 50)
33 print(f"Completion results for {MODEL_NAME}:")
34 print(completion)
35 print("-" * 50)
36
37
38 def test_accuracy():
39 """Run the end to end accuracy test."""
40 run_simple_prompt()
41
42 model_args = (
43 f"model={MODEL_NAME},"
44 f"base_url={BASE_URL}/completions,"
45 f"num_concurrent={NUM_CONCURRENT},tokenized_requests=False"
46 )
47 results = lm_eval.simple_evaluate(
48 model="local-completions",
49 model_args=model_args,
50 tasks=TASK,
51 )
52
53 measured_value = results["results"][TASK][FILTER]
54 expected_value = EXPECTED_VALUES.get(MODEL_NAME)
55
56 print(f"Measured accuracy value: {measured_value}\n")
57 if expected_value is None:
58 print(
59 f"Warning: No expected value found for {MODEL_NAME}. "
60 "Skipping accuracy check."
61 )
62 return
63
64 assert (
65 measured_value - RTOL < expected_value
66 and measured_value + RTOL > expected_value
67 ), f"Expected: {expected_value} | Measured: {measured_value}"