main
py 55 lines 1.47 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Call a token-in/token-out server started with:
4
5 vllm serve Qwen/Qwen3-0.6B --enable-scale-out
6 """
7
8 import httpx
9 from transformers import AutoTokenizer
10
11 GEN_ENDPOINT = "http://localhost:8000/inference/v1/generate"
12 DUMMY_API_KEY = "empty"
13 MODEL_NAME = "Qwen/Qwen3-0.6B"
14
15 transport = httpx.HTTPTransport()
16 headers = {"Authorization": f"Bearer {DUMMY_API_KEY}"}
17 client = httpx.Client(
18 transport=transport,
19 base_url=GEN_ENDPOINT,
20 timeout=600,
21 headers=headers,
22 )
23 messages = [
24 {"role": "system", "content": "You are a helpful assistant."},
25 {"role": "user", "content": "How many countries are in the EU?"},
26 ]
27
28
29 def main(client):
30 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
31 token_ids = tokenizer.apply_chat_template(
32 messages,
33 add_generation_prompt=True,
34 enable_thinking=False,
35 return_dict=True,
36 ).input_ids
37 payload = {
38 "model": MODEL_NAME,
39 "token_ids": token_ids,
40 "sampling_params": {"max_tokens": 24, "temperature": 0.2, "detokenize": False},
41 "stream": False,
42 }
43 resp = client.post(GEN_ENDPOINT, json=payload)
44 resp.raise_for_status()
45 data = resp.json()
46 print(data)
47 print("-" * 50)
48 print("Token generation results:")
49 res = tokenizer.decode(data["choices"][0]["token_ids"])
50 print(res)
51 print("-" * 50)
52
53
54 if __name__ == "__main__":
55 main(client)