main
py 64 lines 1.82 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """An example shows how to generate chat completions from reasoning models
4 like DeepSeekR1.
5
6 To run this example, you need to start the vLLM server
7 with the reasoning parser:
8
9 ```bash
10 vllm serve deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B \
11 --reasoning-parser deepseek_r1
12 ```
13
14 This example demonstrates how to generate chat completions from reasoning models
15 using the OpenAI Python client library.
16 """
17
18 from openai import OpenAI
19
20 # Modify OpenAI's API key and API base to use vLLM's API server.
21 openai_api_key = "EMPTY"
22 openai_api_base = "http://localhost:8000/v1"
23
24
25 def main():
26 client = OpenAI(
27 api_key=openai_api_key,
28 base_url=openai_api_base,
29 )
30
31 models = client.models.list()
32 model = models.data[0].id
33
34 # Round 1
35 messages = [{"role": "user", "content": "9.11 and 9.8, which is greater?"}]
36 # ruff: noqa: E501
37 # For granite, add: `extra_body={"chat_template_kwargs": {"thinking": True}}`
38 response = client.chat.completions.create(model=model, messages=messages)
39
40 reasoning = response.choices[0].message.reasoning
41 content = response.choices[0].message.content
42
43 print("reasoning for Round 1:", reasoning)
44 print("content for Round 1:", content)
45
46 # Round 2
47 messages.append({"role": "assistant", "content": content})
48 messages.append(
49 {
50 "role": "user",
51 "content": "How many Rs are there in the word 'strawberry'?",
52 }
53 )
54 response = client.chat.completions.create(model=model, messages=messages)
55
56 reasoning = response.choices[0].message.reasoning
57 content = response.choices[0].message.content
58
59 print("reasoning for Round 2:", reasoning)
60 print("content for Round 2:", content)
61
62
63 if __name__ == "__main__":
64 main()