Fix chat endpoint support for Ray in-framework MBridge deployment (#572)

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

Author: athitten

nemo_deploy/llm/megatronllm_deployable.py

@@ -205,12 +205,12 @@ def apply_chat_template(self, messages, add_generation_prompt=True):
         Works when model's tokenizer has chat template (typically chat models).
         """
         try:
-            tokenizer_chat_template = self.mcore_tokenizer.tokenizer.tokenizer.chat_template
+            tokenizer_chat_template = self.mcore_tokenizer._tokenizer.chat_template
 
-            # Try to get bos_token, handle different tokenizer types
+            # Try to get bos_token
             bos_token = None
             try:
-                bos_token = self.mcore_tokenizer.tokenizer.tokenizer.bos_token
+                bos_token = self.mcore_tokenizer._tokenizer.bos_token
             except AttributeError:
                 # Some tokenizers might not have bos_token, use empty string as fallback
                 bos_token = ""
@@ -223,11 +223,11 @@ def apply_chat_template(self, messages, add_generation_prompt=True):
                 )
 
             template = Template(tokenizer_chat_template)
-        except AttributeError:
+        except AttributeError as e:
             # If the tokenizer does not have chat_template
             raise ValueError(
-                "The tokenizer does not have chat template, if you would like to evaluate chat model \
-                             ensure your model's tokenizer has a chat template"
+                f"The tokenizer does not have chat template ({e}), if you would like to evaluate chat model "
+                "ensure your model's tokenizer has a chat template"
             )
         # Render the template with the provided messages
         rendered_output = template.render(
@@ -240,14 +240,14 @@ def apply_chat_template(self, messages, add_generation_prompt=True):
 
     def remove_eos_token(self, text):
         """Removes eos token if it exists in the output, otherwise does nothing."""
-        # Handle different tokenizer types
+        # MBridge tokenizer: access underlying HF tokenizer via _tokenizer
         try:
-            eos_token = self.mcore_tokenizer.tokenizer.tokenizer.eos_token
+            eos_token = self.mcore_tokenizer._tokenizer.eos_token
         except AttributeError:
             # Fallback for TiktokenTokenizer and similar tokenizers
             try:
-                eos_id = self.mcore_tokenizer.tokenizer.tokenizer.eos_id
-                eos_token = self.mcore_tokenizer.tokenizer.tokenizer.special_tokens[eos_id]
+                eos_id = self.mcore_tokenizer._tokenizer.eos_id
+                eos_token = self.mcore_tokenizer._tokenizer.inv_vocab[eos_id]
             except AttributeError:
                 # If neither approach works, return text unchanged
                 return text

nemo_deploy/llm/megatronllm_deployable_ray.py

@@ -352,15 +352,25 @@ 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)
+            top_p = request.get("top_p", 0.0)
+            top_k = request.get("top_k", 0)
+
+            # Greedy sampling check: when both temperature and top_p are 0, set top_k to 1
+            if temperature == 0.0 and top_p == 0.0:
+                LOGGER.warning("Both temperature and top_p are 0. Setting top_k to 1 to ensure greedy sampling.")
+                top_k = 1
+
             # Prepare inference parameters
             # For chat templates, we need to pass the entire messages list as a single prompt
             # so that apply_chat_template receives the full conversation context
             inference_inputs = {
-                "prompts": [messages],  # Wrap messages in a list so apply_chat_template gets the full conversation
-                "max_length": request.get("max_tokens", 256),
-                "temperature": request.get("temperature", 1.0),
-                "top_k": request.get("top_k", 0),
-                "top_p": request.get("top_p", 0.0),
+                "prompts": [messages],
+                "max_length": request.get("max_tokens", 512),
+                "temperature": temperature,
+                "top_k": top_k,
+                "top_p": top_p,
                 "compute_logprob": True if request.get("logprobs") == 1 else False,
                 "apply_chat_template": request.get("apply_chat_template", True),
             }
@@ -370,7 +380,6 @@ async def chat_completions(self, request: Dict[Any, Any]):
 
             # Extract generated texts from results
             generated_texts = results["sentences"]
-
             # Calculate token counts
             prompt_tokens = sum(len(str(msg).split()) for msg in messages)
             completion_tokens = sum(len(r.split()) for r in generated_texts)

tests/unit_tests/deploy/test_megatronllm_deployable.py

@@ -28,10 +28,11 @@ def mock_engine_and_tokenizer():
     mock_engine = MagicMock()
     mock_model = MagicMock()
     mock_tokenizer = MagicMock()
-    mock_tokenizer.tokenizer.tokenizer = MagicMock()
-    mock_tokenizer.tokenizer.tokenizer.chat_template = "{{messages}}"
-    mock_tokenizer.tokenizer.tokenizer.bos_token = "<bos>"
-    mock_tokenizer.tokenizer.tokenizer.eos_token = "<eos>"
+    # MBridge tokenizer uses _tokenizer to access the underlying HF tokenizer
+    mock_tokenizer._tokenizer = MagicMock()
+    mock_tokenizer._tokenizer.chat_template = "{{messages}}"
+    mock_tokenizer._tokenizer.bos_token = "<bos>"
+    mock_tokenizer._tokenizer.eos_token = "<eos>"
 
     return mock_engine, mock_model, mock_tokenizer
 
@@ -70,7 +71,7 @@ def test_dict_to_str():
 @pytest.mark.run_only_on("GPU")
 def test_apply_chat_template_none_template(deployable):
     """Test chat template application when tokenizer has no chat template."""
-    deployable.mcore_tokenizer.tokenizer.tokenizer.chat_template = None
+    deployable.mcore_tokenizer._tokenizer.chat_template = None
     messages = [{"role": "user", "content": "Hello"}]
 
     with pytest.raises(ValueError, match="The tokenizer does not have a chat template defined"):
@@ -81,7 +82,7 @@ def test_apply_chat_template_none_template(deployable):
 def test_apply_chat_template_attribute_error(deployable):
     """Test chat template application when tokenizer raises AttributeError."""
     # Remove the chat_template attribute to trigger AttributeError
-    del deployable.mcore_tokenizer.tokenizer.tokenizer.chat_template
+    del deployable.mcore_tokenizer._tokenizer.chat_template
     messages = [{"role": "user", "content": "Hello"}]
 
     with pytest.raises(ValueError, match="The tokenizer does not have chat template"):
@@ -441,7 +442,7 @@ def test_apply_chat_template(deployable):
 @pytest.mark.run_only_on("GPU")
 def test_remove_eos_token(deployable):
     """Test EOS token removal covering all code paths."""
-    # Test normal case - eos_token accessible via self.mcore_tokenizer.tokenizer.tokenizer.eos_token
+    # Test normal case - eos_token accessible via self.mcore_tokenizer._tokenizer.eos_token
     texts = ["Hello<eos>", "World", "Test<eos>"]
     cleaned_texts = deployable.remove_eos_token(texts)
     assert cleaned_texts == ["Hello", "World", "Test"]
@@ -452,7 +453,7 @@ def test_remove_eos_token_comprehensive(deployable):
     """Test remove_eos_token method covering all code paths and edge cases."""
 
     # Test case 1: Normal case - eos_token accessible
-    deployable.mcore_tokenizer.tokenizer.tokenizer.eos_token = "<eos>"
+    deployable.mcore_tokenizer._tokenizer.eos_token = "<eos>"
     texts = ["Hello<eos>", "World", "Test<eos>"]
     cleaned_texts = deployable.remove_eos_token(texts)
     assert cleaned_texts == ["Hello", "World", "Test"]
@@ -477,10 +478,10 @@ def test_remove_eos_token_comprehensive(deployable):
     cleaned_texts = deployable.remove_eos_token(texts)
     assert cleaned_texts == ["", "Hello", ""]
 
-    # Test case 6: Fallback case 1 - AttributeError on eos_token, but eos_id and special_tokens work
-    del deployable.mcore_tokenizer.tokenizer.tokenizer.eos_token
-    deployable.mcore_tokenizer.tokenizer.tokenizer.eos_id = 2
-    deployable.mcore_tokenizer.tokenizer.tokenizer.special_tokens = {2: "</s>"}
+    # Test case 6: Fallback case 1 - AttributeError on eos_token, but eos_id and inv_vocab work
+    del deployable.mcore_tokenizer._tokenizer.eos_token
+    deployable.mcore_tokenizer._tokenizer.eos_id = 2
+    deployable.mcore_tokenizer._tokenizer.inv_vocab = {2: "</s>"}
 
     texts = ["Hello</s>", "World", "Test</s>"]
     cleaned_texts = deployable.remove_eos_token(texts)
@@ -492,16 +493,16 @@ def test_remove_eos_token_comprehensive(deployable):
     assert cleaned_texts == ["Hello</s>World"]
 
     # Test case 8: Fallback case 2 - Both AttributeErrors, return unchanged
-    del deployable.mcore_tokenizer.tokenizer.tokenizer.eos_id
+    del deployable.mcore_tokenizer._tokenizer.eos_id
 
     texts = ["Hello<eos>", "World</s>", "Test"]
     cleaned_texts = deployable.remove_eos_token(texts)
     assert cleaned_texts == ["Hello<eos>", "World</s>", "Test"]  # Should return unchanged
 
     # Test case 9: Fallback case 2 - AttributeError on eos_id but not on eos_token access
-    # First delete the special_tokens to trigger AttributeError in second try block
-    deployable.mcore_tokenizer.tokenizer.tokenizer.eos_id = 2
-    del deployable.mcore_tokenizer.tokenizer.tokenizer.special_tokens
+    # First delete the inv_vocab to trigger AttributeError in second try block
+    deployable.mcore_tokenizer._tokenizer.eos_id = 2
+    del deployable.mcore_tokenizer._tokenizer.inv_vocab
 
     texts = ["Hello<eos>", "World"]
     cleaned_texts = deployable.remove_eos_token(texts)