[feat] completion api supports passing input token ids in either prompt or prompt_token_ids (#3311)

* [feat] completion api supports passing input token ids in either `prompt` or `prompt_token_ids`

* [fix] update comment

* [fix] fix type error

* [test] add a unittest file for serving api test

* [test] try to fix ci error

* [chore] rename test function names

* [test] try to fix ci error

* [test] try to fix ci error

* [test] add tests for qwen

Author: liyonghua0910

fastdeploy/entrypoints/openai/protocol.py

@@ -438,7 +438,7 @@ class CompletionRequest(BaseModel):
 
     max_streaming_response_tokens: Optional[int] = None
     return_token_ids: Optional[bool] = None
-    prompt_token_ids: Optional[List[int]] = None
+    prompt_token_ids: Optional[Union[List[int], List[List[int]]]] = None
     # doc: end-completion-extra-params
 
     def to_dict_for_infer(self, request_id=None, prompt=None):
@@ -463,11 +463,11 @@ def to_dict_for_infer(self, request_id=None, prompt=None):
         if prompt is not None:
             req_dict["prompt"] = prompt
 
-        if "prompt_token_ids" in req_dict:
-            if "prompt" in req_dict:
-                del req_dict["prompt"]
-        else:
-            assert len(prompt) > 0
+        # if "prompt_token_ids" in req_dict:
+        #     if "prompt" in req_dict:
+        #         del req_dict["prompt"]
+        # else:
+        #     assert len(prompt) > 0
 
         guided_json_object = None
         if self.response_format is not None:
@@ -572,7 +572,7 @@ class ChatCompletionRequest(BaseModel):
     stop_token_ids: Optional[List[int]] = Field(default_factory=list)
     # doc: end-chat-completion-sampling-params
 
-    # doc: start-completion-extra-params
+    # doc: start-chat-completion-extra-params
     chat_template_kwargs: Optional[dict] = None
     chat_template: Optional[str] = None
     reasoning_max_tokens: Optional[int] = None

fastdeploy/entrypoints/openai/serving_completion.py

@@ -81,35 +81,50 @@ async def create_completion(self, request: CompletionRequest):
             request_id = f"cmpl-{request.user}-{uuid.uuid4()}"
         else:
             request_id = f"cmpl-{uuid.uuid4()}"
-        api_server_logger.info(f"initialize request {request_id}")
+        api_server_logger.info(f"Initialize request {request_id}: {request}")
         request_prompt_ids = None
         request_prompts = None
+
+        # Handle prompt and prompt_token_ids
         try:
-            if isinstance(request.prompt, str):
-                request_prompts = [request.prompt]
-            elif isinstance(request.prompt, list) and all(isinstance(item, int) for item in request.prompt):
-                request_prompt_ids = [request.prompt]
-            elif isinstance(request.prompt, list) and all(isinstance(item, str) for item in request.prompt):
-                request_prompts = request.prompt
-            elif isinstance(request.prompt, list):
-                for item in request.prompt:
-                    if isinstance(item, list) and all(isinstance(x, int) for x in item):
-                        continue
-                    else:
-                        raise ValueError("Prompt must be a string, a list of strings or a list of integers.")
-                request_prompt_ids = request.prompt
+            if request.prompt_token_ids is not None:  # let `prompt_token_ids` support batch inference
+                assert len(request.prompt_token_ids) > 0, "prompt_token_ids should not be an empty list"
+                if isinstance(request.prompt_token_ids[0], list):
+                    request_prompt_ids = request.prompt_token_ids
+                elif isinstance(request.prompt_token_ids[0], int):
+                    request_prompt_ids = [request.prompt_token_ids]
+                else:
+                    raise ValueError(
+                        "If prompt_token_ids is provided, its type should be one of: list[int], list[list[int]]"
+                    )
+                # reset `prompt_token_ids` to avoid data processor directly using it; let data processor fill it
+                request.prompt_token_ids = None
             else:
-                raise ValueError("Prompt must be a string, a list of strings or a list of integers.")
+                if isinstance(request.prompt, str):
+                    request_prompts = [request.prompt]
+                elif isinstance(request.prompt, list) and all(isinstance(item, int) for item in request.prompt):
+                    request_prompt_ids = [request.prompt]
+                elif isinstance(request.prompt, list) and all(isinstance(item, str) for item in request.prompt):
+                    request_prompts = request.prompt
+                elif isinstance(request.prompt, list):
+                    for item in request.prompt:
+                        if isinstance(item, list) and all(isinstance(x, int) for x in item):
+                            continue
+                        else:
+                            raise ValueError("If prompt is a list, each item type must be one of: str, list[int]")
+                    request_prompt_ids = request.prompt
+                else:
+                    raise ValueError("Prompt type must be one of: str, list[str], list[int], list[list[int]]")
         except Exception as e:
             error_msg = f"OpenAIServingCompletion create_completion: {e}, {str(traceback.format_exc())}"
             api_server_logger.error(error_msg)
             return ErrorResponse(message=error_msg, code=400)
 
         if request_prompt_ids is not None:
             request_prompts = request_prompt_ids
-        num_choices = len(request_prompts)
 
-        api_server_logger.info(f"start inference for request {num_choices}")
+        num_choices = len(request_prompts)
+        api_server_logger.info(f"Start preprocessing request: req_id={request_id}), num_choices={num_choices}")
         prompt_batched_token_ids = []
         text_after_process_list = []
         try:
@@ -131,7 +146,7 @@ async def create_completion(self, request: CompletionRequest):
                     request_id_idx = f"{request_id}-{idx}"
                     current_req_dict = request.to_dict_for_infer(request_id_idx, prompt)
                     current_req_dict["arrival_time"] = time.time()
-                    prompt_token_ids = self.engine_client.format_and_add_data(current_req_dict)
+                    prompt_token_ids = self.engine_client.format_and_add_data(current_req_dict)  # tokenize
                     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"))

fastdeploy/input/ernie4_5_processor.py

@@ -87,33 +87,45 @@ def process_request(self, request, max_model_len=None, **kwargs):
             bool: Whether preprocessing is successful
             str: error message
         """
+        data_processor_logger.info(f"Start processing request: {request}")
         request.chat_template = kwargs.get("chat_template")
         request = self._apply_default_parameters(request)
         if request.get("eos_token_ids") is None or len(request.eos_token_ids) == 0:
             request.eos_token_ids = self.eos_token_ids
+
+        # processing stop_sequences
         stop_sequences = request.get("stop", [])
         if stop_sequences is not None and len(stop_sequences) != 0:
             stop_seqs, stop_seqs_len = self.update_stop_seq(stop_sequences)
             request.set("stop_token_ids", stop_seqs)
             request.set("stop_seqs_len", stop_seqs_len)
 
+        # processing bad_words
         bad_words = request.get("bad_words")
         bad_words_token_ids = request.get("bad_words_token_ids")
         if bad_words:
             bad_words_token_ids = self.update_bad_words(bad_words, bad_words_token_ids)
             request["bad_words_token_ids"] = bad_words_token_ids
 
+        # processing prompt_token_ids
         if request.prompt_token_ids is None or len(request.prompt_token_ids) == 0:
-            if request.prompt is None and request.messages is None:
-                raise ValueError(f"The request should have `prompt_token_ids`, `prompt` or `messages`: {request}.")
             if request.prompt is not None:
-                prompt = request.prompt if request.prompt is not None else request.messages[0]
-                prompt = prompt[0] if isinstance(prompt, list) else prompt
-                tokens = self.tokenizer.tokenize(prompt)
-                token_ids = self.tokenizer.convert_tokens_to_ids(tokens)
-                request.prompt_token_ids = token_ids
-                data_processor_logger.info(f"req_id:{request.request_id}, tokens:{tokens}, token_ids: {token_ids}")
-            else:
+                # prompt = request.prompt if request.prompt is not None else request.messages[0]
+                prompt = request.prompt
+                assert isinstance(prompt, str) or (
+                    isinstance(prompt, list) and all([isinstance(t, int) for t in prompt])
+                ), f"prompt must be a string or a list of integers, but got {type(prompt)}"
+
+                if isinstance(prompt, list):  # if prompt is a token id list
+                    request.prompt_token_ids = prompt
+                else:
+                    tokens = self.tokenizer.tokenize(prompt)
+                    token_ids = self.tokenizer.convert_tokens_to_ids(tokens)
+                    request.prompt_token_ids = token_ids
+                    data_processor_logger.debug(
+                        f"request_ids: {request.request_id}, prompt: {prompt}, tokens: {tokens}, token_ids: {token_ids}"
+                    )
+            elif request.messages is not None:
                 task = request.to_dict()
                 chat_template_kwargs = kwargs.get("chat_template_kwargs")
                 if chat_template_kwargs:
@@ -124,24 +136,26 @@ def process_request(self, request, max_model_len=None, **kwargs):
                     else:
                         raise ValueError("Invalid input: chat_template_kwargs must be a dict")
                 request.prompt_token_ids = self.messages2ids(task)
+            else:
+                raise ValueError(f"The request should have `prompt_token_ids`, `prompt` or `messages`: {request}.")
 
         if len(request.prompt_token_ids) == 0:
             raise ValueError("Invalid input: prompt_token_ids must be a non-empty sequence of token IDs")
+
+        # truncate prompts that exceed the length limit
         if max_model_len is not None and len(request.prompt_token_ids) > max_model_len:
             request.prompt_token_ids = request.prompt_token_ids[: max_model_len - 1]
         if request.get("max_tokens") is None:
-            request.set(
-                "max_tokens",
-                max(1, max_model_len - len(request.prompt_token_ids)),
-            )
+            request.set("max_tokens", max(1, max_model_len - len(request.prompt_token_ids)))
         if request.get("temperature") < _SAMPLING_EPS:
             # zero temperature is equivalent to greedy sampling
             request.set("temperature", 1)
         if request.get("top_p") < _SAMPLING_EPS:
             request.set("top_p", _SAMPLING_EPS)
         if self.reasoning_parser and self.reasoning_parser.__class__.__name__ == "ErnieX1ReasoningParser":
             request.enable_thinking = True
-        data_processor_logger.info(f"Processed request {request}")
+
+        data_processor_logger.info(f"Processed request: {request}")
         return request
 
     def process_request_dict(self, request, max_model_len=None):
@@ -155,6 +169,7 @@ def process_request_dict(self, request, max_model_len=None):
             bool: Whether preprocessing is successful
             str: error message
         """
+        data_processor_logger.info(f"Start processing request dict: {request}")
         request = self._apply_default_parameters(request)
         if not request.get("eos_token_ids"):
             request["eos_token_ids"] = self.eos_token_ids
@@ -175,18 +190,21 @@ def process_request_dict(self, request, max_model_len=None):
 
         # processing prompt_token_ids
         if not request.get("prompt_token_ids"):
-            if request.get("prompt") is None and request.get("messages") is None:
-                raise ValueError(f"Request must contain 'prompt_token_ids', 'prompt', or 'messages': {request}")
             if request.get("prompt"):
                 prompt = request.get("prompt")
-                prompt = prompt[0] if isinstance(prompt, list) else prompt
-                request["text_after_process"] = prompt
-                tokens = self.tokenizer.tokenize(prompt)
-                token_ids = self.tokenizer.convert_tokens_to_ids(tokens)
-                request["prompt_token_ids"] = token_ids
-                req_id = request.get("request_id", None)
-                data_processor_logger.info(f"req_id:{req_id}, tokens:{tokens}, token_ids: {token_ids}")
-            else:
+                assert isinstance(prompt, str) or (
+                    isinstance(prompt, list) and all([isinstance(t, int) for t in prompt])
+                ), f"prompt must be a string or a list of integers, but got {type(prompt)}"
+                if isinstance(prompt, list):  # if prompt is a token id list
+                    request["prompt_token_ids"] = prompt
+                else:
+                    request["text_after_process"] = prompt
+                    tokens = self.tokenizer.tokenize(prompt)
+                    token_ids = self.tokenizer.convert_tokens_to_ids(tokens)
+                    request["prompt_token_ids"] = token_ids
+                    req_id = request.get("request_id", None)
+                    data_processor_logger.info(f"req_id:{req_id}, tokens:{tokens}, token_ids: {token_ids}")
+            elif request.get("messages"):
                 chat_template_kwargs = request.get("chat_template_kwargs")
                 if chat_template_kwargs:
                     if isinstance(chat_template_kwargs, dict):
@@ -196,6 +214,9 @@ def process_request_dict(self, request, max_model_len=None):
                     else:
                         raise ValueError("Invalid input: chat_template_kwargs must be a dict")
                 request["prompt_token_ids"] = self.messages2ids(request)
+            else:
+                raise ValueError(f"Request must contain 'prompt_token_ids', 'prompt', or 'messages': {request}")
+
         if len(request["prompt_token_ids"]) == 0:
             raise ValueError("Invalid input: prompt_token_ids must be a non-empty sequence of token IDs")
 
@@ -211,8 +232,8 @@ def process_request_dict(self, request, max_model_len=None):
             request["top_p"] = _SAMPLING_EPS
         if self.reasoning_parser and self.reasoning_parser.__class__.__name__ == "ErnieX1ReasoningParser":
             request["enable_thinking"] = True
-        data_processor_logger.info(f"Processed request {request}")
 
+        data_processor_logger.info(f"Processed request dict: {request}")
         return request
 
     def process_response(self, response_dict, **kwargs):