add details info for streaming generation (#330)

Co-authored-by: baishihao <[email protected]> Co-authored-by: hiworldwzj <[email protected]>
ModelTC · Feb 18, 2024 · 7d48473 · 7d48473
1 parent b4541bd
commit 7d48473
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git 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
@@ -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(