explicitly give a chat template

Signed-off-by: yxia216 <[email protected]>
kserve · Sep 30, 2024 · 0299c07 · 0299c07
1 parent 6289d40
commit 0299c07
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 12 deletions.
diff --git a/python/huggingfaceserver/huggingfaceserver/generative_model.py b/python/huggingfaceserver/huggingfaceserver/generative_model.py
@@ -384,7 +384,7 @@ def build_generation_config(
         return GenerationConfig(**kwargs)
 
     def apply_chat_template(
-        self, messages: Iterable[ChatCompletionRequestMessage]
+        self, messages: Iterable[ChatCompletionRequestMessage], chat_template: Optional[str] = None,
     ) -> ChatPrompt:
         """
         Given a list of chat completion messages, convert them to a prompt.
@@ -394,6 +394,7 @@ def apply_chat_template(
                 str,
                 self._tokenizer.apply_chat_template(
                     [m.model_dump() for m in messages],
+                    chat_template=chat_template,
                     tokenize=False,
                     add_generation_prompt=True,
                 ),

diff --git a/python/huggingfaceserver/tests/test_model.py b/python/huggingfaceserver/tests/test_model.py
@@ -390,7 +390,8 @@ async def test_bloom_chat_completion(bloom_model: HuggingfaceGenerativeModel):
         max_tokens=20,
     )
     request = ChatCompletionRequest(params=params, context={})
-    response = await bloom_model.create_chat_completion(request)
+    chat_template = "{% for message in messages %}" "{{ message.content }}{{ eos_token }}" "{% endfor %}"
+    response = await bloom_model.create_chat_completion(request, chat_template)
     assert (
         response.choices[0].message.content
         == "The first thing you need to do is to get a good idea of what you are looking for."
@@ -417,7 +418,8 @@ async def test_bloom_chat_completion_streaming(bloom_model: HuggingfaceGenerativ
         max_tokens=20,
     )
     request = ChatCompletionRequest(params=params, context={})
-    response = await bloom_model.create_chat_completion(request)
+    chat_template = "{% for message in messages %}" "{{ message.content }}{{ eos_token }}" "{% endfor %}"
+    response = await bloom_model.create_chat_completion(request, chat_template)
     output = ""
     async for chunk in response:
         output += chunk.choices[0].delta.content

diff --git a/python/huggingfaceserver/tests/test_vllm_model.py b/python/huggingfaceserver/tests/test_vllm_model.py
@@ -162,7 +162,8 @@ async def mock_generate(*args, **kwargs) -> AsyncIterator[RequestOutput]:
             max_tokens=10,
         )
         request = ChatCompletionRequest(params=params, context={})
-        response = await opt_model.create_chat_completion(request)
+        chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content']}}{% if (loop.last and add_generation_prompt) or not loop.last %}{{ '<|im_end|>' + '\n'}}{% endif %}{% endfor %} {% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}{{ '<|im_start|>assistant\n' }}{% endif %}"
+        response = await opt_model.create_chat_completion(request, chat_template)
         expected = CreateChatCompletionResponse(
             id=request_id,
             choices=[
@@ -220,7 +221,8 @@ async def mock_generate(*args, **kwargs) -> AsyncIterator[RequestOutput]:
         request = ChatCompletionRequest(
             request_id=request_id, params=params, context={}
         )
-        response = await opt_model.create_chat_completion(request)
+        chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content']}}{% if (loop.last and add_generation_prompt) or not loop.last %}{{ '<|im_end|>' + '\n'}}{% endif %}{% endfor %} {% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}{{ '<|im_start|>assistant\n' }}{% endif %}"
+        response = await opt_model.create_chat_completion(request, chat_template)
         expected = CreateChatCompletionResponse(
             id=request_id,
             choices=[
@@ -279,7 +281,8 @@ async def mock_generate(*args, **kwargs) -> AsyncIterator[RequestOutput]:
         request = ChatCompletionRequest(
             request_id=request_id, params=params, context={}
         )
-        response_iterator = await opt_model.create_chat_completion(request)
+        chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content']}}{% if (loop.last and add_generation_prompt) or not loop.last %}{{ '<|im_end|>' + '\n'}}{% endif %}{% endfor %} {% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}{{ '<|im_start|>assistant\n' }}{% endif %}"
+        response_iterator = await opt_model.create_chat_completion(request, chat_template)
         completion = ""
         async for resp in response_iterator:
             assert len(resp.choices) == 1
@@ -329,7 +332,8 @@ async def mock_generate(*args, **kwargs) -> AsyncIterator[RequestOutput]:
         request = ChatCompletionRequest(
             request_id=request_id, params=params, context={}
         )
-        response = await opt_model.create_chat_completion(request)
+        chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content']}}{% if (loop.last and add_generation_prompt) or not loop.last %}{{ '<|im_end|>' + '\n'}}{% endif %}{% endfor %} {% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}{{ '<|im_start|>assistant\n' }}{% endif %}"
+        response = await opt_model.create_chat_completion(request, chat_template)
         expected = CreateChatCompletionResponse(
             id=request_id,
             choices=[
@@ -639,7 +643,8 @@ async def mock_generate(*args, **kwargs) -> AsyncIterator[RequestOutput]:
         request = ChatCompletionRequest(
             request_id=request_id, params=params, context={}
         )
-        response_iterator = await opt_model.create_chat_completion(request)
+        chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content']}}{% if (loop.last and add_generation_prompt) or not loop.last %}{{ '<|im_end|>' + '\n'}}{% endif %}{% endfor %} {% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}{{ '<|im_start|>assistant\n' }}{% endif %}"
+        response_iterator = await opt_model.create_chat_completion(request, chat_template)
         completion = ""
         log_probs = ChatCompletionChoiceLogprobs(
             content=[],
@@ -895,7 +900,8 @@ async def mock_generate(*args, **kwargs) -> AsyncIterator[RequestOutput]:
             request_id=request_id, params=params, context={}
         )
         with pytest.raises(OpenAIError):
-            await opt_model.create_chat_completion(request)
+            chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content']}}{% if (loop.last and add_generation_prompt) or not loop.last %}{{ '<|im_end|>' + '\n'}}{% endif %}{% endfor %} {% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}{{ '<|im_start|>assistant\n' }}{% endif %}"
+            await opt_model.create_chat_completion(request, chat_template)
 
     async def test_vllm_chat_completion_facebook_opt_model_with_logit_bias(
         self, vllm_opt_model
@@ -930,7 +936,8 @@ async def mock_generate(*args, **kwargs) -> AsyncIterator[RequestOutput]:
         request = ChatCompletionRequest(
             request_id=request_id, params=params, context={}
         )
-        response = await opt_model.create_chat_completion(request)
+        chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content']}}{% if (loop.last and add_generation_prompt) or not loop.last %}{{ '<|im_end|>' + '\n'}}{% endif %}{% endfor %} {% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}{{ '<|im_start|>assistant\n' }}{% endif %}"
+        response = await opt_model.create_chat_completion(request, chat_template)
         expected = CreateChatCompletionResponse(
             id=request_id,
             choices=[

diff --git a/python/kserve/kserve/protocol/rest/openai/openai_chat_adapter_model.py b/python/kserve/kserve/protocol/rest/openai/openai_chat_adapter_model.py
@@ -193,15 +193,15 @@ def completion_to_chat_completion_chunk(
         )
 
     async def create_chat_completion(
-        self, request: ChatCompletionRequest
+        self, request: ChatCompletionRequest, chat_template: str | None = None,
     ) -> Union[ChatCompletion, AsyncIterator[ChatCompletionChunk]]:
         params = request.params
 
         if params.n != 1:
             raise InvalidInput("n != 1 is not supported")
 
         # Convert the messages into a prompt
-        chat_prompt = self.apply_chat_template(params.messages)
+        chat_prompt = self.apply_chat_template(params.messages, chat_template)
         # Translate the chat completion request to a completion request
         completion_params = self.chat_completion_params_to_completion_params(
             params, chat_prompt.prompt