Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support logprobs for chat completion api #2292

Closed
wants to merge 1 commit into from
Closed

Support logprobs for chat completion api #2292

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 43 additions & 2 deletions vllm/entrypoints/openai/api_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -272,6 +272,8 @@ async def create_chat_completion(request: ChatCompletionRequest,
top_k=request.top_k,
ignore_eos=request.ignore_eos,
use_beam_search=request.use_beam_search,
logprobs=request.logprobs,
prompt_logprobs=request.logprobs if request.echo else None,
skip_special_tokens=request.skip_special_tokens,
spaces_between_special_tokens=spaces_between_special_tokens,
)
Expand All @@ -292,7 +294,11 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]:
role = get_role()
for i in range(request.n):
choice_data = ChatCompletionResponseStreamChoice(
index=i, delta=DeltaMessage(role=role), finish_reason=None)
index=i,
delta=DeltaMessage(role=role),
logprobs=(LogProbs()
if request.logprobs is not None else None),
finish_reason=None)
chunk = ChatCompletionStreamResponse(id=request_id,
object=chunk_object_type,
created=created_time,
Expand All @@ -314,6 +320,8 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]:
choice_data = ChatCompletionResponseStreamChoice(
index=i,
delta=DeltaMessage(content=last_msg_content),
logprobs=(LogProbs()
if request.logprobs is not None else None),
finish_reason=None)
chunk = ChatCompletionStreamResponse(
id=request_id,
Expand All @@ -339,11 +347,24 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]:
if output.finish_reason is None:
# Send token-by-token response for each request.n
delta_text = output.text[len(previous_texts[i]):]
if request.logprobs is not None:
token_ids = output.token_ids[previous_num_tokens[i]:]
top_logprobs = output.logprobs[previous_num_tokens[i]:]
offsets = len(previous_texts[i])
logprobs = create_logprobs(
token_ids=token_ids,
top_logprobs=top_logprobs,
num_output_top_logprobs=request.logprobs,
initial_text_offset=offsets,
)
else:
logprobs = None
previous_texts[i] = output.text
previous_num_tokens[i] = len(output.token_ids)
choice_data = ChatCompletionResponseStreamChoice(
index=i,
delta=DeltaMessage(content=delta_text),
logprobs=logprobs,
finish_reason=None)
chunk = ChatCompletionStreamResponse(
id=request_id,
Expand All @@ -362,7 +383,11 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]:
total_tokens=prompt_tokens + previous_num_tokens[i],
)
choice_data = ChatCompletionResponseStreamChoice(
index=i, delta=[], finish_reason=output.finish_reason)
index=i,
delta=[],
logprobs=(LogProbs()
if request.logprobs is not None else None),
finish_reason=output.finish_reason)
chunk = ChatCompletionStreamResponse(
id=request_id,
object=chunk_object_type,
Expand Down Expand Up @@ -391,11 +416,27 @@ async def completion_full_generator():
assert final_res is not None

choices = []
prompt_token_ids = final_res.prompt_token_ids
prompt_logprobs = final_res.prompt_logprobs
role = get_role()
for output in final_res.outputs:
if request.logprobs is not None:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if request.logprobs is False(or 0), it will still trigger this branch. Maybe a better way:

if request.logprobs:
    # do things
else:
   # do things

Please also update the same issue in other places.

token_ids = output.token_ids
top_logprobs = output.logprobs
if request.echo:
token_ids = prompt_token_ids + token_ids
top_logprobs = prompt_logprobs + top_logprobs
logprobs = create_logprobs(
token_ids=token_ids,
top_logprobs=top_logprobs,
num_output_top_logprobs=request.logprobs,
)
else:
logprobs = None
choice_data = ChatCompletionResponseChoice(
index=output.index,
message=ChatMessage(role=role, content=output.text),
logprobs=logprobs,
finish_reason=output.finish_reason,
)
choices.append(choice_data)
Expand Down
3 changes: 3 additions & 0 deletions vllm/entrypoints/openai/protocol.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,7 @@ class ChatCompletionRequest(BaseModel):
max_tokens: Optional[int] = None
stop: Optional[Union[str, List[str]]] = Field(default_factory=list)
stream: Optional[bool] = False
logprobs: Optional[int] = None
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you change the type to bool, which aligns with openAI spec. Also I think in the current codebase, it uses it as bool as well. Please correct me if I'm wrong.

Could you change the type definition of logprobs in CompletionRequest as well? thanks.

presence_penalty: Optional[float] = 0.0
frequency_penalty: Optional[float] = 0.0
logit_bias: Optional[Dict[str, float]] = None
Expand Down Expand Up @@ -155,6 +156,7 @@ class ChatMessage(BaseModel):
class ChatCompletionResponseChoice(BaseModel):
index: int
message: ChatMessage
logprobs: Optional[LogProbs] = None
finish_reason: Optional[Literal["stop", "length"]] = None


Expand All @@ -175,6 +177,7 @@ class DeltaMessage(BaseModel):
class ChatCompletionResponseStreamChoice(BaseModel):
index: int
delta: DeltaMessage
logprobs: Optional[LogProbs] = None
finish_reason: Optional[Literal["stop", "length"]] = None


Expand Down
Loading