Add multimodal token usage (#1016)

Co-authored-by: pigKiller <zhhang.bian@gmail.com>

Author: hiworldwzj

lightllm/server/api_lightllm.py

@@ -52,17 +52,22 @@ async def lightllm_generate(request: Request, httpserver_manager: HttpServerMana
     prompt_tokens = 0
     prompt_token_ids = None
     is_first_metadata = True
+    input_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
         if is_first_metadata:
             prompt_logprobs = metadata.get("prompt_logprobs", None)
             prompt_token_ids = metadata.get("prompt_token_ids", None)
             prompt_tokens = metadata.get("prompt_tokens", 0)
+            input_usage = metadata.get("input_usage", None)
             if prompt_logprobs is not None:
                 del metadata["prompt_logprobs"]
             if prompt_token_ids is not None:
                 del metadata["prompt_token_ids"]
+            if input_usage is not None:
+                del metadata["input_usage"]
+
             is_first_metadata = False
 
         count_output_tokens_dict[sub_req_id] += 1
@@ -95,6 +100,9 @@ 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 input_usage is not None:
+        ret["input_usage"] = input_usage
+
     return Response(content=json.dumps(ret, ensure_ascii=False).encode("utf-8"))
 
 
@@ -116,7 +124,12 @@ async def lightllm_generate_stream(request: Request, httpserver_manager: HttpSer
 
     # Streaming case
     async def stream_results() -> AsyncGenerator[bytes, None]:
+        # input_usage 只会在第一个metadata中出现，所以需要保存下来
+        input_usage = None
         async for _, request_output, metadata, finish_status in results_generator:
+            if input_usage is None:
+                input_usage = metadata.get("input_usage", None)
+
             ret = {
                 "token": {
                     "id": metadata.get("id", None),
@@ -130,6 +143,7 @@ async def stream_results() -> AsyncGenerator[bytes, None]:
                 "finished": finish_status.is_finished(),
                 "finish_reason": finish_status.get_finish_reason(),
                 "details": None,
+                "input_usage": input_usage,
             }
 
             yield ("data:" + json.dumps(ret, ensure_ascii=False) + "\n\n").encode("utf-8")

lightllm/server/httpserver/manager.py

@@ -325,7 +325,23 @@ async def generate(
                 req_status,
                 request,
             )
+
+            # 计算输入 token 使用量统计
+            image_tokens, audio_tokens = self._count_multimodal_tokens(multimodal_params)
+            text_tokens = len(prompt_ids) - (image_tokens + audio_tokens)
+            input_usage = {
+                "input_text_tokens": text_tokens,
+                "input_audio_tokens": audio_tokens,
+                "input_image_tokens": image_tokens,
+            }
+
+            is_first_gen_token = True
             async for sub_req_id, request_output, metadata, finish_status in results_generator:
+                # 只有第一个生成的 token 的 metadata 中包含 input_usage
+                if is_first_gen_token:
+                    metadata["input_usage"] = input_usage
+                    is_first_gen_token = False
+
                 yield sub_req_id, request_output, metadata, finish_status
 
         except Exception as e:
@@ -340,6 +356,20 @@ async def generate(
             raise e
         return
 
+    def _count_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() and multimodal_params is not None:
+            for img in multimodal_params.images:
+                if img.token_num is not None:
+                    image_tokens += img.token_num
+            for audio in multimodal_params.audios:
+                if audio.token_num is not None:
+                    audio_tokens += audio.token_num
+
+        return image_tokens, audio_tokens
+
     async def _log_req_header(self, request_headers, group_request_id: int):
 
         x_request_id = request_headers.get("X-Request-Id", "")