Unverified Commit 71075029 authored by CYJiang's avatar CYJiang Committed by GitHub
Browse files

[Doc] Support --stream arg in openai_completion_client.py script (#18388)


Signed-off-by: default avatargoogs1025 <googs1025@gmail.com>
parent ca86a7cf
...@@ -12,6 +12,9 @@ from enum import Enum ...@@ -12,6 +12,9 @@ from enum import Enum
from openai import BadRequestError, OpenAI from openai import BadRequestError, OpenAI
from pydantic import BaseModel from pydantic import BaseModel
openai_api_key = "EMPTY"
openai_api_base = "http://localhost:8000/v1"
# Guided decoding by Choice (list of possible options) # Guided decoding by Choice (list of possible options)
def guided_choice_completion(client: OpenAI, model: str): def guided_choice_completion(client: OpenAI, model: str):
...@@ -134,8 +137,8 @@ def extra_backend_options_completion(client: OpenAI, model: str): ...@@ -134,8 +137,8 @@ def extra_backend_options_completion(client: OpenAI, model: str):
def main(): def main():
client: OpenAI = OpenAI( client: OpenAI = OpenAI(
base_url="http://localhost:8000/v1", base_url=openai_api_base,
api_key="-", api_key=openai_api_key,
) )
model = client.models.list().data[0].id model = client.models.list().data[0].id
......
...@@ -7,11 +7,14 @@ from openai import OpenAI ...@@ -7,11 +7,14 @@ from openai import OpenAI
# to enforce the format of a tool call response, but it could be used for # to enforce the format of a tool call response, but it could be used for
# any structured output within a subset of the response. # any structured output within a subset of the response.
openai_api_key = "EMPTY"
openai_api_base = "http://localhost:8000/v1"
def main(): def main():
client = OpenAI( client = OpenAI(
base_url="http://localhost:8000/v1", base_url=openai_api_base,
api_key="-", api_key=openai_api_key,
) )
messages = [{ messages = [{
......
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
import argparse
from openai import OpenAI from openai import OpenAI
# Modify OpenAI's API key and API base to use vLLM's API server. # Modify OpenAI's API key and API base to use vLLM's API server.
...@@ -7,7 +9,15 @@ openai_api_key = "EMPTY" ...@@ -7,7 +9,15 @@ openai_api_key = "EMPTY"
openai_api_base = "http://localhost:8000/v1" openai_api_base = "http://localhost:8000/v1"
def main(): def parse_args():
parser = argparse.ArgumentParser(description="Client for vLLM API server")
parser.add_argument("--stream",
action="store_true",
help="Enable streaming response")
return parser.parse_args()
def main(args):
client = OpenAI( client = OpenAI(
# defaults to os.environ.get("OPENAI_API_KEY") # defaults to os.environ.get("OPENAI_API_KEY")
api_key=openai_api_key, api_key=openai_api_key,
...@@ -18,18 +28,17 @@ def main(): ...@@ -18,18 +28,17 @@ def main():
model = models.data[0].id model = models.data[0].id
# Completion API # Completion API
stream = False
completion = client.completions.create( completion = client.completions.create(
model=model, model=model,
prompt="A robot may not injure a human being", prompt="A robot may not injure a human being",
echo=False, echo=False,
n=2, n=2,
stream=stream, stream=args.stream,
logprobs=3) logprobs=3)
print("-" * 50) print("-" * 50)
print("Completion results:") print("Completion results:")
if stream: if args.stream:
for c in completion: for c in completion:
print(c) print(c)
else: else:
...@@ -38,4 +47,5 @@ def main(): ...@@ -38,4 +47,5 @@ def main():
if __name__ == "__main__": if __name__ == "__main__":
main() args = parse_args()
main(args)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment