main
py 70 lines 1.98 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Example online usage of sequence reward models.
4
5 Run `vllm serve <model> --runner pooling`
6 to start up the server in vLLM. e.g.
7
8 vllm serve Skywork/Skywork-Reward-V2-Qwen3-0.6B
9
10 The key distinction between sequence classification and token classification
11 lies in their output granularity: sequence classification produces a single
12 result for an entire input sequence, whereas token classification yields a
13 result for each individual token within the sequence.
14 """
15
16 import argparse
17 import pprint
18
19 import requests
20
21
22 def post_http_request(prompt: dict, api_url: str) -> requests.Response:
23 headers = {"User-Agent": "Test Client"}
24 response = requests.post(api_url, headers=headers, json=prompt)
25 return response
26
27
28 def parse_args():
29 parser = argparse.ArgumentParser()
30 parser.add_argument("--host", type=str, default="localhost")
31 parser.add_argument("--port", type=int, default=8000)
32
33 return parser.parse_args()
34
35
36 def main(args):
37 base_url = f"http://{args.host}:{args.port}"
38 models_url = base_url + "/v1/models"
39 pooing_url = base_url + "/pooling"
40
41 response = requests.get(models_url)
42 model = response.json()["data"][0]["id"]
43
44 # Input like Completions API
45 prompt = {"model": model, "input": "vLLM is great!"}
46 pooling_response = post_http_request(prompt=prompt, api_url=pooing_url)
47 print("-" * 50)
48 print("Pooling Response:")
49 pprint.pprint(pooling_response.json())
50 print("-" * 50)
51
52 # Input like Chat API
53 prompt = {
54 "model": model,
55 "messages": [
56 {
57 "role": "user",
58 "content": [{"type": "text", "text": "vLLM is great!"}],
59 }
60 ],
61 }
62 pooling_response = post_http_request(prompt=prompt, api_url=pooing_url)
63 print("Pooling Response:")
64 pprint.pprint(pooling_response.json())
65 print("-" * 50)
66
67
68 if __name__ == "__main__":
69 args = parse_args()
70 main(args)