From 7d48473f21a56ed7f400e47e202f4eda21fccecb Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Sun, 18 Feb 2024 17:50:26 +0800 Subject: [PATCH] add details info for streaming generation (#330) Co-authored-by: baishihao Co-authored-by: hiworldwzj <30762946+hiworldwzj@users.noreply.github.com> --- lightllm/server/api_server.py | 3 ++- lightllm/server/detokenization/manager.py | 3 +++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/lightllm/server/api_server.py b/lightllm/server/api_server.py index cedc0dc74..04a9545e8 100755 --- a/lightllm/server/api_server.py +++ b/lightllm/server/api_server.py @@ -171,7 +171,8 @@ async def stream_results() -> AsyncGenerator[bytes, None]: "id": metadata.get("id", None), "text": request_output, "logprob": metadata.get("logprob", None), - "special": False + "special": metadata.get("special", False), + "count_output_tokens": metadata.get("count_output_tokens", 0), }, "generated_text": None, "finished": finish_status.is_finished(), diff --git a/lightllm/server/detokenization/manager.py b/lightllm/server/detokenization/manager.py index 1620afdfe..2609012d1 100644 --- a/lightllm/server/detokenization/manager.py +++ b/lightllm/server/detokenization/manager.py @@ -33,6 +33,7 @@ def __init__( self.send_to_httpserver.connect(f"tcp://127.0.0.1:{httpserver_port}") self.tokenizer = get_tokenizer(model_weightdir, tokenizor_mode, trust_remote_code=trust_remote_code) + self.all_special_ids = set(self.tokenizer.all_special_ids) self.req_id_to_out = {} self.skip_special_tokens = skip_special_tokens @@ -58,6 +59,8 @@ async def handle_loop(self): continue req_out:ReqDetokenizationState = self.req_id_to_out[req_id] req_out.output_ids.append(new_token_id) + new_gen_metadata["special"] = new_token_id in self.all_special_ids + new_gen_metadata["count_output_tokens"] = len(req_out.output_ids) req_out.gen_metadata.update(new_gen_metadata) out_text = decode_token(