From 8aca432df56efd1cd1d44de5fb33701258e717ec Mon Sep 17 00:00:00 2001 From: Aaron <29749331+aarnphm@users.noreply.github.com> Date: Mon, 6 Nov 2023 20:52:49 -0500 Subject: [PATCH] chore(examples): add logprobs check on PyTorch Signed-off-by: Aaron <29749331+aarnphm@users.noreply.github.com> --- examples/openai_client.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/examples/openai_client.py b/examples/openai_client.py index dfd6e44e..ff87cb0c 100644 --- a/examples/openai_client.py +++ b/examples/openai_client.py @@ -2,6 +2,9 @@ from __future__ import annotations import os import openai +import importlib.util + +SUPPORT_LOGPROBS = importlib.util.find_spec('vllm') is not None openai.api_base = os.getenv('OPENLLM_ENDPOINT', 'http://localhost:3000') + '/v1' openai.api_key = 'na' @@ -15,7 +18,7 @@ print('\n' +'-'*50 + ' /v1/completions' + ' [stream=False] ' + '-'*50 + "\n") print(openai.Completion.create(model=MODEL, prompt='Write a tagline for an ice cream shop.', max_tokens=128)) print('\n' +'-'*50 + ' /v1/completions' + ' [stream=True] ' + '-'*50 + "\n") -for chunk in openai.Completion.create(model=MODEL, prompt='Say this is a test', max_tokens=12, temperature=0.8, stream=True, logprobs=2): +for chunk in openai.Completion.create(model=MODEL, prompt='Say this is a test', max_tokens=12, temperature=0.8, stream=True, logprobs=2 if SUPPORT_LOGPROBS else None): print(chunk) print('\n' +'-'*50 + ' /v1/chat/completions' + ' [stream=False] ' + '-'*50 + "\n")