diff --git a/tests/entrypoints/openai/completion/test_completion_text_offset.py b/tests/entrypoints/openai/completion/test_completion_text_offset.py new file mode 100644 index 0000000000..c8ba7df62f --- /dev/null +++ b/tests/entrypoints/openai/completion/test_completion_text_offset.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json + +import pytest + +from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata +from vllm.entrypoints.openai.completion.protocol import CompletionRequest +from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion +from vllm.logprobs import Logprob +from vllm.outputs import CompletionOutput, RequestOutput + +TOKENS = [(1, "abc"), (2, " de"), (3, "!")] +# Text deltas as the detokenizer streams them with stop=["\n\n"]: one char is +# held back until the request finishes, so the text lags the tokens. +TEXT_DELTAS = ["ab", "c d", "e!"] + + +def _logprobs(i: int) -> dict[int, Logprob]: + token_id, token = TOKENS[i] + return {token_id: Logprob(logprob=-1.0, rank=1, decoded_token=token)} + + +def _output(text: str, idxs: list[int], finished: bool) -> RequestOutput: + return RequestOutput( + request_id="r", + prompt="Hi", + prompt_token_ids=[9], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text=text, + token_ids=[TOKENS[i][0] for i in idxs], + cumulative_logprob=None, + logprobs=[_logprobs(i) for i in idxs], + finish_reason="stop" if finished else None, + ) + ], + finished=finished, + ) + + +def _serving() -> OpenAIServingCompletion: + serving = OpenAIServingCompletion.__new__(OpenAIServingCompletion) + serving.enable_prompt_tokens_details = False + serving.enable_force_include_usage = False + serving.enable_per_request_metrics = False + serving.system_fingerprint = None + serving.return_tokens_as_token_ids = False + return serving + + +def _request(stream: bool) -> CompletionRequest: + return CompletionRequest( + model="m", prompt="Hi", max_tokens=8, logprobs=0, stop=["\n\n"], stream=stream + ) + + +@pytest.mark.asyncio +async def test_stream_text_offset_matches_non_stream_with_stop_holdback(): + async def results(): + for i, text in enumerate(TEXT_DELTAS): + yield 0, _output(text, [i], finished=i == len(TEXT_DELTAS) - 1) + + serving = _serving() + stream_offsets = [] + async for line in serving.completion_stream_generator( + _request(stream=True), + [{"prompt": "Hi"}], + results(), + "r", + 0, + "m", + num_prompts=1, + tokenizer=None, + request_metadata=RequestResponseMetadata(request_id="r"), + ): + data = line.removeprefix("data: ").strip() + if data == "[DONE]": + continue + for choice in json.loads(data)["choices"]: + stream_offsets += choice["logprobs"]["text_offset"] + + response = serving.request_output_to_completion_response( + [_output("".join(TEXT_DELTAS), [0, 1, 2], finished=True)], + _request(stream=False), + "r", + 0, + "m", + None, + RequestResponseMetadata(request_id="r"), + ) + assert response.choices[0].logprobs is not None + assert response.choices[0].logprobs.text_offset == [0, 3, 6] + assert stream_offsets == [0, 3, 6] diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index 81dd00f0b3..7f7bdac075 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -289,7 +289,7 @@ class OpenAIServingCompletion(GenerateBaseServing): request_metadata: RequestResponseMetadata, ) -> AsyncGenerator[str, None]: num_choices = 1 if request.n is None else request.n - previous_text_lens = [0] * num_choices * num_prompts + next_text_offsets = [0] * num_choices * num_prompts previous_num_tokens = [0] * num_choices * num_prompts has_echoed = [False] * num_choices * num_prompts num_prompt_tokens = [0] * num_prompts @@ -386,13 +386,18 @@ class OpenAIServingCompletion(GenerateBaseServing): num_output_top_logprobs=request.logprobs, tokenizer=tokenizer, logprob_token_ids=request.logprob_token_ids, - initial_text_offset=previous_text_lens[i], + initial_text_offset=next_text_offsets[i], return_as_token_id=request.return_tokens_as_token_ids, ) + # Advance by the tokens just reported, not by delta_text: + # text held back for stop strings lags the tokens. + if logprobs.tokens: + next_text_offsets[i] = logprobs.text_offset[-1] + len( + logprobs.tokens[-1] + ) else: logprobs = None - previous_text_lens[i] += len(delta_text) previous_num_tokens[i] += len(output.token_ids) stop_reason = output.stop_reason