Completion add raw_prediction/text_after_process (#3356)

This commit is contained in:
memoryCoderC
2025-08-12 23:06:45 +08:00
committed by GitHub
parent 2c0d853067
commit 2d1a4cacdf
7 changed files with 51 additions and 15 deletions

View File

@@ -100,6 +100,7 @@ class OpenAIServingCompletion:
api_server_logger.info(f"start inference for request {num_choices}")
prompt_batched_token_ids = []
text_after_process_list = []
try:
for idx, prompt in enumerate(request_prompts):
request_id_idx = f"{request_id}-{idx}"
@@ -109,6 +110,7 @@ class OpenAIServingCompletion:
prompt_token_ids = self.engine_client.format_and_add_data(current_req_dict)
if isinstance(prompt_token_ids, np.ndarray):
prompt_token_ids = prompt_token_ids.tolist()
text_after_process_list.append(current_req_dict.get("text_after_process"))
prompt_batched_token_ids.append(prompt_token_ids)
except Exception as e:
return ErrorResponse(message=str(e), code=400)
@@ -131,6 +133,7 @@ class OpenAIServingCompletion:
created_time=created_time,
model_name=request.model,
prompt_batched_token_ids=prompt_batched_token_ids,
text_after_process_list=text_after_process_list,
)
else:
try:
@@ -141,6 +144,7 @@ class OpenAIServingCompletion:
created_time=created_time,
model_name=request.model,
prompt_batched_token_ids=prompt_batched_token_ids,
text_after_process_list=text_after_process_list,
)
except Exception as e:
return ErrorResponse(code=400, message=str(e))
@@ -156,6 +160,7 @@ class OpenAIServingCompletion:
created_time: int,
model_name: str,
prompt_batched_token_ids: list(),
text_after_process_list: list(),
):
"""
Process the full completion request with multiple choices.
@@ -225,6 +230,7 @@ class OpenAIServingCompletion:
model_name=model_name,
prompt_batched_token_ids=prompt_batched_token_ids,
completion_batched_token_ids=completion_batched_token_ids,
text_after_process_list=text_after_process_list,
)
except Exception as e:
api_server_logger.error(f"Error in completion_full_generator: {e}", exc_info=True)
@@ -251,6 +257,7 @@ class OpenAIServingCompletion:
created_time: int,
model_name: str,
prompt_batched_token_ids: list(),
text_after_process_list: list(),
):
"""
Process the stream completion request.
@@ -309,6 +316,7 @@ class OpenAIServingCompletion:
index=idx,
text="",
prompt_token_ids=list(prompt_batched_token_ids[idx]),
text_after_process=text_after_process_list[idx],
completion_token_ids=None,
)
],
@@ -337,6 +345,7 @@ class OpenAIServingCompletion:
text=output["text"],
prompt_token_ids=None,
completion_token_ids=output.get("token_ids") if request.return_token_ids else None,
raw_prediction=output.get("raw_prediction") if request.return_token_ids else None,
tool_calls=output.get("tool_call_content"),
reasoning_content=output.get("reasoning_content"),
arrival_time=arrival_time,
@@ -398,6 +407,7 @@ class OpenAIServingCompletion:
model_name: str,
prompt_batched_token_ids: list(),
completion_batched_token_ids: list(),
text_after_process_list: list(),
) -> CompletionResponse:
choices: List[CompletionResponseChoice] = []
num_prompt_tokens = 0
@@ -444,6 +454,8 @@ class OpenAIServingCompletion:
text=output_text,
prompt_token_ids=prompt_token_ids if request.return_token_ids else None,
completion_token_ids=completion_token_ids if request.return_token_ids else None,
raw_prediction=output.get("raw_prediction") if request.return_token_ids else None,
text_after_process=text_after_process_list[idx] if request.return_token_ids else None,
reasoning_content=output.get("reasoning_content"),
tool_calls=output.get("tool_call_content"),
logprobs=aggregated_logprobs,