[Frontend] Allow return_tokens_as_token_ids to be passed as a request param (#14066)

Signed-off-by: Benjamin Chislett <benjamin.chislett@centml.ai>
This commit is contained in:
Benjamin Chislett
2025-03-05 01:30:40 -05:00
committed by GitHub
parent dae9ec464c
commit 32985bed7c
4 changed files with 64 additions and 25 deletions

View File

@@ -316,6 +316,8 @@ class OpenAIServingCompletion(OpenAIServing):
num_output_top_logprobs=request.logprobs,
tokenizer=tokenizer,
initial_text_offset=previous_text_lens[i],
return_as_token_id=request.
return_tokens_as_token_ids,
)
else:
logprobs = None
@@ -436,6 +438,7 @@ class OpenAIServingCompletion(OpenAIServing):
top_logprobs=out_logprobs,
tokenizer=tokenizer,
num_output_top_logprobs=request.logprobs,
return_as_token_id=request.return_tokens_as_token_ids,
)
else:
logprobs = None
@@ -477,6 +480,7 @@ class OpenAIServingCompletion(OpenAIServing):
num_output_top_logprobs: int,
tokenizer: AnyTokenizer,
initial_text_offset: int = 0,
return_as_token_id: Optional[bool] = None,
) -> CompletionLogProbs:
"""Create logprobs for OpenAI Completion API."""
out_text_offset: list[int] = []
@@ -486,11 +490,13 @@ class OpenAIServingCompletion(OpenAIServing):
last_token_len = 0
should_return_as_token_id = return_as_token_id if \
return_as_token_id is not None else self.return_tokens_as_token_ids
for i, token_id in enumerate(token_ids):
step_top_logprobs = top_logprobs[i]
if step_top_logprobs is None:
token = tokenizer.decode(token_id)
if self.return_tokens_as_token_ids:
if should_return_as_token_id:
token = f"token_id:{token_id}"
out_tokens.append(token)
@@ -503,7 +509,7 @@ class OpenAIServingCompletion(OpenAIServing):
step_token,
token_id,
tokenizer,
return_as_token_id=self.return_tokens_as_token_ids,
return_as_token_id=should_return_as_token_id,
)
token_logprob = max(step_token.logprob, -9999.0)
@@ -520,7 +526,7 @@ class OpenAIServingCompletion(OpenAIServing):
self._get_decoded_token(top_lp[1],
top_lp[0],
tokenizer,
return_as_token_id=self.return_tokens_as_token_ids):
return_as_token_id=should_return_as_token_id):
max(top_lp[1].logprob, -9999.0)
for i, top_lp in enumerate(step_top_logprobs.items())
if num_output_top_logprobs >= i