Add multimodal token usage (#1011)

Author: zhhangBian

lightllm/server/api_lightllm.py

@@ -52,6 +52,7 @@ async def lightllm_generate(request: Request, httpserver_manager: HttpServerMana
     prompt_tokens = 0
     prompt_token_ids = None
     is_first_metadata = True
+    usage = None
     async for sub_req_id, request_output, metadata, finish_status in results_generator:
         # when set "--return_all_prompt_logprobs", the first token metadata will contains
         # prompt_logprobs and prompt_token_ids
@@ -65,6 +66,8 @@ async def lightllm_generate(request: Request, httpserver_manager: HttpServerMana
                 del metadata["prompt_token_ids"]
             is_first_metadata = False
 
+        usage = metadata.get("usage", None)
+
         count_output_tokens_dict[sub_req_id] += 1
         final_output_dict[sub_req_id].append(request_output)
         if return_details:
@@ -95,6 +98,8 @@ async def lightllm_generate(request: Request, httpserver_manager: HttpServerMana
         ret["prompt_token_ids"] = prompt_token_ids
     if prompt_logprobs is not None:
         ret["prompt_logprobs"] = prompt_logprobs
+    if usage is not None:
+        ret["usage"] = usage
     return Response(content=json.dumps(ret, ensure_ascii=False).encode("utf-8"))
 
 
@@ -130,6 +135,7 @@ async def stream_results() -> AsyncGenerator[bytes, None]:
                 "finished": finish_status.is_finished(),
                 "finish_reason": finish_status.get_finish_reason(),
                 "details": None,
+                "usage": metadata.get("usage", None),
             }
 
             yield ("data:" + json.dumps(ret, ensure_ascii=False) + "\n\n").encode("utf-8")

lightllm/server/httpserver/manager.py

@@ -213,6 +213,21 @@ def tokens(self, prompt, multimodal_params, samping_params: SamplingParams, kwar
             audio_tokens += self.tokenizer.get_audio_token_length(audio)
         return len(prompt_ids) + image_tokens + img_count + audio_tokens + audio_count
 
+    def _calculate_multimodal_tokens(self, multimodal_params: MultimodalParams) -> Tuple[int, int]:
+        image_tokens = 0
+        audio_tokens = 0
+
+        if self.enable_multimodal and self.pd_mode.is_P_or_NORMAL():
+            for img in multimodal_params.images:
+                image_tokens += self.tokenizer.get_image_token_length(img)
+            for audio in multimodal_params.audios:
+                audio_tokens += self.tokenizer.get_audio_token_length(audio)
+        else:
+            image_tokens = len(multimodal_params.images)
+            audio_tokens = len(multimodal_params.audios)
+
+        return image_tokens, audio_tokens
+
     async def loop_for_request(self):
         assert self.args.node_rank > 0
         while True:
@@ -311,6 +326,16 @@ async def generate(
                 req_objs.append(req_obj)
 
             req_status = ReqStatus(group_request_id, multimodal_params, req_objs, start_time)
+
+            # 计算输入 token 使用量统计
+            text_tokens = len(prompt_ids)
+            image_tokens, audio_tokens = self._calculate_multimodal_tokens(multimodal_params)
+            input_usage = {
+                "input_text_tokens": text_tokens,
+                "input_audio_tokens": audio_tokens,
+                "input_image_tokens": image_tokens,
+            }
+
             self.req_id_to_out_inf[group_request_id] = req_status
 
             await self.transfer_to_next_module_or_node(
@@ -326,6 +351,7 @@ async def generate(
                 request,
             )
             async for sub_req_id, request_output, metadata, finish_status in results_generator:
+                metadata["usage"] = {**input_usage, **metadata.get("usage", {})}
                 yield sub_req_id, request_output, metadata, finish_status
 
         except Exception as e:
@@ -513,6 +539,8 @@ async def _wait_to_token_package(
                     if self.pd_mode == NodeRole.P and is_first_token:
                         metadata["prompt_ids"] = prompt_ids
 
+                    metadata["usage"] = {"output_tokens": out_token_counter}
+
                     prompt_cache_len = metadata.pop("prompt_cache_len", 0)
                     if is_first_token:
                         first_token_cost_ms = (time.time() - start_time) * 1000