vllm/examples/online_serving/openai_completion_client.py

34 lines
699 B
Python
Raw Normal View History

# SPDX-License-Identifier: Apache-2.0
from openai import OpenAI
2023-06-17 03:07:40 -07:00
# Modify OpenAI's API key and API base to use vLLM's API server.
openai_api_key = "EMPTY"
openai_api_base = "http://localhost:8000/v1"
2023-05-23 21:39:50 -07:00
client = OpenAI(
# defaults to os.environ.get("OPENAI_API_KEY")
api_key=openai_api_key,
base_url=openai_api_base,
)
2023-05-23 21:39:50 -07:00
models = client.models.list()
model = models.data[0].id
# Completion API
stream = False
completion = client.completions.create(
model=model,
prompt="A robot may not injure a human being",
echo=False,
n=2,
stream=stream,
logprobs=3)
2023-05-23 21:39:50 -07:00
print("Completion results:")
2023-05-23 21:39:50 -07:00
if stream:
for c in completion:
print(c)
else:
print(completion)