Fix chat endpoint support for HF deployment with Ray (#575)

Signed-off-by: Abhishree <abhishreetm@gmail.com>

Author: athitten

nemo_deploy/llm/hf_deployable.py

@@ -155,6 +155,45 @@ def _load(
         if self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
 
+    def apply_chat_template(self, messages, add_generation_prompt=True):
+        """Apply the chat template to messages using the tokenizer.
+
+        This method uses the HuggingFace tokenizer's built-in apply_chat_template method
+        to format messages according to the model's expected chat format.
+
+        Args:
+            messages: List of message dictionaries with 'role' and 'content' keys,
+                     or a JSON string representation of messages.
+            add_generation_prompt (bool): Whether to add the generation prompt. Defaults to True.
+
+        Returns:
+            str: The formatted prompt string.
+
+        Raises:
+            ValueError: If the tokenizer does not have a chat template.
+        """
+        import json
+
+        # Handle JSON string input
+        if isinstance(messages, str):
+            messages = json.loads(messages)
+
+        # Check if tokenizer has chat_template
+        if not hasattr(self.tokenizer, "chat_template") or self.tokenizer.chat_template is None:
+            raise ValueError(
+                "The tokenizer does not have a chat template defined. "
+                "If you would like to evaluate a chat model, ensure your model's tokenizer has a chat template."
+            )
+
+        # Use tokenizer's apply_chat_template method
+        formatted_prompt = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=add_generation_prompt,
+        )
+
+        return formatted_prompt
+
     def generate(
         self,
         **kwargs: Any,
@@ -529,6 +568,7 @@ def ray_infer_fn(self, inputs: Dict[Any, Any]):
                 - compute_logprob: Whether to compute log probabilities (optional)
                 - n_top_logprobs: Number of top log probabilities to return (optional)
                 - echo: Whether to echo the prompt in output (optional)
+                - apply_chat_template: Whether to apply chat template to prompts (optional)
 
         Returns:
             Dict[str, Any]: Dictionary containing:
@@ -549,6 +589,7 @@ def ray_infer_fn(self, inputs: Dict[Any, Any]):
             compute_logprob = inputs.pop("compute_logprob", False)
             n_top_logprobs = inputs.pop("n_top_logprobs", 0)
             echo = inputs.pop("echo", False)
+            apply_chat_template = inputs.pop("apply_chat_template", False)
 
             output_infer = self._infer_fn_ray(
                 prompts=prompts,
@@ -561,6 +602,7 @@ def ray_infer_fn(self, inputs: Dict[Any, Any]):
                 compute_logprob=compute_logprob,
                 n_top_logprobs=n_top_logprobs,
                 echo=echo,
+                apply_chat_template=apply_chat_template,
             )
             # Code to get logprobs (required in OAI API format for eval) from the scores in output_infer.
             if (compute_logprob or n_top_logprobs > 0) and "scores" in output_infer and output_infer["scores"]:
@@ -600,6 +642,7 @@ def _infer_fn_ray(
         compute_logprob=False,
         n_top_logprobs=0,
         echo=False,
+        apply_chat_template=False,
         cast_output_func=None,
     ):
         """Common internal function for inference operations.
@@ -615,11 +658,16 @@ def _infer_fn_ray(
             compute_logprob: Whether to compute log probabilities
             n_top_logprobs: Number of top log probabilities to return
             echo: Whether to echo the prompt in output
+            apply_chat_template: Whether to apply chat template to prompts
             cast_output_func: Optional function to cast output values
 
         Returns:
             Dict containing inference results with raw outputs
         """
+        # Apply chat template if requested (for Ray inference only)
+        if apply_chat_template:
+            prompts = [self.apply_chat_template(prompt) for prompt in prompts]
+
         # Enable return_dict if we need scores for logprobs or if output_logits/scores are requested
         return_dict_in_generate = output_logits or output_scores or compute_logprob or n_top_logprobs > 0
         # Enable output_scores if we need to compute logprobs. scores and logits from generate are both identical in

nemo_deploy/llm/hf_deployable_ray.py

@@ -294,27 +294,29 @@ async def chat_completions(self, request: Dict[Any, Any]):
             # Extract parameters from the request dictionary
             messages = request.get("messages", [])
 
+            # Get sampling parameters
+            temperature = request.get("temperature", 0.0)
             # vLLM requires top_p to be in (0, 1], so handle invalid values
             top_p = request.get("top_p", None)
             if top_p is not None and top_p <= 0.0:
                 LOGGER.warning(
                     f"top_p must be in (0, 1] for vLLM, got {top_p}. Setting to 0.1 for greedy-like sampling."
                 )
                 request["top_p"] = 0.1
-
-            # Convert messages to a single prompt
-            prompt = "\n".join([f"{msg.get('role', 'user')}: {msg.get('content', '')}" for msg in messages])
-            prompt += "\nassistant:"
+                if temperature == 0.0:
+                    LOGGER.warning("Both temperature and top_p are 0. Setting top_k to 1 to ensure greedy sampling.")
+                    request["top_k"] = 1
 
             # Prepare inference parameters using the formatted prompt
             inference_inputs = {
-                "prompts": [prompt],  # Use formatted prompt string instead of raw messages
+                "prompts": [messages],  # Pass messages list, will be formatted by apply_chat_template
                 "max_tokens": request.get("max_tokens", 256),
-                "temperature": request.get("temperature", 1.0),
+                "temperature": temperature,
                 "top_k": request.get("top_k", 0),
                 "top_p": request.get("top_p", 1.0),  # vLLM requires top_p in (0, 1], use 1.0 as default
                 "output_logits": request.get("output_logits", False),
                 "output_scores": request.get("output_scores", False),
+                "apply_chat_template": request.get("apply_chat_template", True),
             }
 
             # Run model inference in the thread pool
@@ -353,12 +355,7 @@ async def chat_completions(self, request: Dict[Any, Any]):
                                 "scores": scores,
                             }
                             if scores is not None
-                            else None,
-                            {
-                                "logits": logits,
-                            }
-                            if logits is not None
-                            else None,
+                            else None
                         ),
                         "finish_reason": (
                             "length"

tests/unit_tests/deploy/test_hf_deployable.py

@@ -378,6 +378,131 @@ def test_ray_infer_fn_max_tokens_edge_cases(self, mock_model, mock_tokenizer, mo
         assert "sentences" in output
 
 
+class TestHFDeployableApplyChatTemplate:
+    """Tests for apply_chat_template functionality added to HuggingFaceLLMDeploy."""
+
+    @pytest.fixture
+    def mock_tokenizer_with_chat_template(self):
+        """Create a mock tokenizer with chat template support."""
+        tokenizer = MagicMock(spec=AutoTokenizer)
+        tokenizer.pad_token = "[PAD]"
+        tokenizer.eos_token = "[EOS]"
+        tokenizer.chat_template = "{% for message in messages %}{{ message.role }}: {{ message.content }}\n{% endfor %}"
+        tokenizer.apply_chat_template = MagicMock(return_value="<|begin_of_text|>user: Hello\nassistant:")
+        return tokenizer
+
+    @pytest.fixture
+    def mock_tokenizer_without_chat_template(self):
+        """Create a mock tokenizer without chat template support."""
+        tokenizer = MagicMock(spec=AutoTokenizer)
+        tokenizer.pad_token = "[PAD]"
+        tokenizer.eos_token = "[EOS]"
+        tokenizer.chat_template = None
+        return tokenizer
+
+    @pytest.fixture
+    def hf_deployable_with_chat_template(self, mock_model, mock_tokenizer_with_chat_template):
+        """Create HuggingFaceLLMDeploy instance with chat template support."""
+        deployer = HuggingFaceLLMDeploy(
+            model=mock_model, tokenizer=mock_tokenizer_with_chat_template, task="text-generation"
+        )
+        return deployer
+
+    @pytest.fixture
+    def hf_deployable_without_chat_template(self, mock_model, mock_tokenizer_without_chat_template):
+        """Create HuggingFaceLLMDeploy instance without chat template support."""
+        deployer = HuggingFaceLLMDeploy(
+            model=mock_model, tokenizer=mock_tokenizer_without_chat_template, task="text-generation"
+        )
+        return deployer
+
+    def test_apply_chat_template_with_messages_list(self, hf_deployable_with_chat_template):
+        """Test apply_chat_template with a list of message dictionaries."""
+        messages = [{"role": "user", "content": "Hello, how are you?"}]
+
+        result = hf_deployable_with_chat_template.apply_chat_template(messages)
+
+        assert result == "<|begin_of_text|>user: Hello\nassistant:"
+        hf_deployable_with_chat_template.tokenizer.apply_chat_template.assert_called_once_with(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+
+    def test_apply_chat_template_with_json_string(self, hf_deployable_with_chat_template):
+        """Test apply_chat_template with JSON string input."""
+        import json
+
+        messages = [{"role": "user", "content": "Hello"}]
+        messages_json = json.dumps(messages)
+
+        result = hf_deployable_with_chat_template.apply_chat_template(messages_json)
+
+        assert result == "<|begin_of_text|>user: Hello\nassistant:"
+        # Verify it was called with the parsed list
+        hf_deployable_with_chat_template.tokenizer.apply_chat_template.assert_called_once()
+
+    def test_apply_chat_template_without_generation_prompt(self, hf_deployable_with_chat_template):
+        """Test apply_chat_template with add_generation_prompt=False."""
+        messages = [{"role": "user", "content": "Hello"}]
+
+        hf_deployable_with_chat_template.apply_chat_template(messages, add_generation_prompt=False)
+
+        hf_deployable_with_chat_template.tokenizer.apply_chat_template.assert_called_once_with(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False,
+        )
+
+    def test_apply_chat_template_raises_error_when_no_template(self, hf_deployable_without_chat_template):
+        """Test apply_chat_template raises ValueError when tokenizer has no chat template."""
+        messages = [{"role": "user", "content": "Hello"}]
+
+        with pytest.raises(ValueError, match="The tokenizer does not have a chat template defined"):
+            hf_deployable_without_chat_template.apply_chat_template(messages)
+
+    def test_ray_infer_fn_with_apply_chat_template(
+        self, mock_model, mock_tokenizer_with_chat_template, mock_torch_cuda
+    ):
+        """Test ray_infer_fn correctly applies chat template when requested."""
+        deployer = HuggingFaceLLMDeploy(
+            model=mock_model, tokenizer=mock_tokenizer_with_chat_template, task="text-generation"
+        )
+
+        messages = [{"role": "user", "content": "Hello"}]
+        inputs = {
+            "prompts": [messages],
+            "max_tokens": 100,
+            "apply_chat_template": True,
+        }
+
+        result = deployer.ray_infer_fn(inputs)
+
+        assert "sentences" in result
+        # Verify apply_chat_template was called
+        mock_tokenizer_with_chat_template.apply_chat_template.assert_called()
+
+    def test_ray_infer_fn_without_apply_chat_template(
+        self, mock_model, mock_tokenizer_with_chat_template, mock_torch_cuda
+    ):
+        """Test ray_infer_fn does not apply chat template when not requested."""
+        deployer = HuggingFaceLLMDeploy(
+            model=mock_model, tokenizer=mock_tokenizer_with_chat_template, task="text-generation"
+        )
+
+        inputs = {
+            "prompts": ["plain text prompt"],
+            "max_tokens": 100,
+            "apply_chat_template": False,
+        }
+
+        result = deployer.ray_infer_fn(inputs)
+
+        assert "sentences" in result
+        # Verify apply_chat_template was NOT called
+        mock_tokenizer_with_chat_template.apply_chat_template.assert_not_called()
+
+
 class TestHFDeployableEchoAndLogprobs:
     """Tests for echo parameter and logprobs functionality added to HuggingFaceLLMDeploy."""