From 05135bbe299927fd1c971c0c61d52432a2ff94b3 Mon Sep 17 00:00:00 2001 From: Igor Gitman Date: Fri, 20 Feb 2026 12:15:30 -0800 Subject: [PATCH 1/5] Debugging Signed-off-by: Igor Gitman --- nemo_skills/inference/generate.py | 2 ++ nemo_skills/prompt/utils.py | 10 ++++++++++ 2 files changed, 12 insertions(+) diff --git a/nemo_skills/inference/generate.py b/nemo_skills/inference/generate.py index c2feefff8a..63f3548b16 100644 --- a/nemo_skills/inference/generate.py +++ b/nemo_skills/inference/generate.py @@ -701,7 +701,9 @@ async def process_single_datapoint(self, data_point, all_data, prompt_format=Non result = await self.generate_with_semaphore(**generation_params) + print("SDF", self.cfg.count_prompt_tokens) if self.cfg.count_prompt_tokens: + print("SD234F", generation_params["prompt"]) num_input_tokens = get_token_count(self.hf_tokenizer, generation_params["prompt"]) result["num_input_tokens"] = num_input_tokens diff --git a/nemo_skills/prompt/utils.py b/nemo_skills/prompt/utils.py index 332b051c73..7285aa3ef5 100644 --- a/nemo_skills/prompt/utils.py +++ b/nemo_skills/prompt/utils.py @@ -387,6 +387,7 @@ def message_to_dict(orig_message: Any) -> Dict[str, Any]: if messages is None: return None + print("!@$#!@$", messages) if isinstance(messages, str): return len(tokenizer.encode(messages, add_special_tokens=False)) elif isinstance(messages, list): @@ -395,9 +396,18 @@ def message_to_dict(orig_message: Any) -> Dict[str, Any]: message if isinstance(message, dict) else message_to_dict(copy.deepcopy(message)) for message in messages ] try: + print("ASDF", messages) + print( + "$#!@$!", + len(tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, tools=tools)), + tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, tools=tools), + ) + return len(tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, tools=tools)) + except Exception as e: raise ValueError(f"Invalid chat message format: {e}") + else: raise ValueError("messages must be a string or a list of dictionaries") From 56ae5e873e9501df88612504b99faae4a8ff966a Mon Sep 17 00:00:00 2001 From: Igor Gitman Date: Fri, 20 Feb 2026 12:29:03 -0800 Subject: [PATCH 2/5] Handle newer hf transformers chat template api Signed-off-by: Igor Gitman --- nemo_skills/inference/model/utils.py | 6 ++++- nemo_skills/inference/prover.py | 3 +++ nemo_skills/prompt/utils.py | 14 ++++------- tests/test_prompts.py | 35 +++++++++++++++++++++++++++- 4 files changed, 47 insertions(+), 11 deletions(-) diff --git a/nemo_skills/inference/model/utils.py b/nemo_skills/inference/model/utils.py index 27bf917e74..4d3f0683fa 100644 --- a/nemo_skills/inference/model/utils.py +++ b/nemo_skills/inference/model/utils.py @@ -95,7 +95,11 @@ def encode(self, prompt: str | list[dict], tools=None) -> list[int]: if isinstance(prompt, str): return self.tokenizer.encode(prompt) elif isinstance(prompt, list): - return self.tokenizer.apply_chat_template(prompt, add_generation_prompt=True, tools=tools) + result = self.tokenizer.apply_chat_template(prompt, add_generation_prompt=True, tools=tools) + # Handle newer HF tokenizer versions that return a dict instead of a list + if isinstance(result, dict): + result = result["input_ids"] + return result def decode(self, tokens: list[int]) -> str: """Decode a list of tokens using the tokenizer.""" diff --git a/nemo_skills/inference/prover.py b/nemo_skills/inference/prover.py index 8faebd63cd..37c2473e80 100644 --- a/nemo_skills/inference/prover.py +++ b/nemo_skills/inference/prover.py @@ -283,6 +283,9 @@ async def _single_data_point_generate(self, data_point, data): prefix_tokens = self.hf_tokenizer.apply_chat_template( prepared_conversation, tokenize=True, add_generation_prompt=True ) + # Handle newer HF tokenizer versions that return a dict instead of a list + if isinstance(prefix_tokens, dict): + prefix_tokens = prefix_tokens["input_ids"] num_tokens_prefix = len(prefix_tokens) prefix = self.hf_tokenizer.apply_chat_template( prepared_conversation, tokenize=False, add_generation_prompt=True diff --git a/nemo_skills/prompt/utils.py b/nemo_skills/prompt/utils.py index 7285aa3ef5..a34429cc2c 100644 --- a/nemo_skills/prompt/utils.py +++ b/nemo_skills/prompt/utils.py @@ -387,7 +387,6 @@ def message_to_dict(orig_message: Any) -> Dict[str, Any]: if messages is None: return None - print("!@$#!@$", messages) if isinstance(messages, str): return len(tokenizer.encode(messages, add_special_tokens=False)) elif isinstance(messages, list): @@ -396,14 +395,11 @@ def message_to_dict(orig_message: Any) -> Dict[str, Any]: message if isinstance(message, dict) else message_to_dict(copy.deepcopy(message)) for message in messages ] try: - print("ASDF", messages) - print( - "$#!@$!", - len(tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, tools=tools)), - tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, tools=tools), - ) - - return len(tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, tools=tools)) + result = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, tools=tools) + # Handle newer HF tokenizer versions that return a dict instead of a list + if isinstance(result, dict): + result = result["input_ids"] + return len(result) except Exception as e: raise ValueError(f"Invalid chat message format: {e}") diff --git a/tests/test_prompts.py b/tests/test_prompts.py index a63a7bfaf3..7577dc13c9 100644 --- a/tests/test_prompts.py +++ b/tests/test_prompts.py @@ -13,7 +13,40 @@ # limitations under the License. -from nemo_skills.prompt.utils import get_prompt +from unittest.mock import MagicMock + +from nemo_skills.prompt.utils import get_prompt, get_token_count + + +def test_get_token_count_string(): + tokenizer = MagicMock() + tokenizer.encode.return_value = [1, 2, 3, 4, 5] + assert get_token_count(tokenizer, "hello world") == 5 + tokenizer.encode.assert_called_once_with("hello world", add_special_tokens=False) + + +def test_get_token_count_messages_list_result(): + """Test with old HF tokenizer behavior that returns a list of token ids.""" + tokenizer = MagicMock() + tokenizer.apply_chat_template.return_value = [1, 2, 3, 4, 5, 6] + messages = [{"role": "user", "content": "hello"}] + assert get_token_count(tokenizer, messages) == 6 + + +def test_get_token_count_messages_dict_result(): + """Test with new HF tokenizer behavior that returns a dict with input_ids.""" + tokenizer = MagicMock() + tokenizer.apply_chat_template.return_value = { + "input_ids": [1, 2, 3, 4, 5, 6, 7], + "attention_mask": [1, 1, 1, 1, 1, 1, 1], + } + messages = [{"role": "user", "content": "hello"}] + assert get_token_count(tokenizer, messages) == 7 + + +def test_get_token_count_none(): + assert get_token_count(None, "hello") is None + assert get_token_count(MagicMock(), None) is None def test_generic_math_problem_augmentation_prompt(): From 30d6ca944c46f734c2f32bb9960322908850ef66 Mon Sep 17 00:00:00 2001 From: Igor Gitman Date: Fri, 20 Feb 2026 12:33:58 -0800 Subject: [PATCH 3/5] Update tests Signed-off-by: Igor Gitman --- nemo_skills/inference/generate.py | 2 -- tests/test_prompts.py | 47 ++++++++++++++----------------- 2 files changed, 21 insertions(+), 28 deletions(-) diff --git a/nemo_skills/inference/generate.py b/nemo_skills/inference/generate.py index 63f3548b16..c2feefff8a 100644 --- a/nemo_skills/inference/generate.py +++ b/nemo_skills/inference/generate.py @@ -701,9 +701,7 @@ async def process_single_datapoint(self, data_point, all_data, prompt_format=Non result = await self.generate_with_semaphore(**generation_params) - print("SDF", self.cfg.count_prompt_tokens) if self.cfg.count_prompt_tokens: - print("SD234F", generation_params["prompt"]) num_input_tokens = get_token_count(self.hf_tokenizer, generation_params["prompt"]) result["num_input_tokens"] = num_input_tokens diff --git a/tests/test_prompts.py b/tests/test_prompts.py index 7577dc13c9..67c0eba887 100644 --- a/tests/test_prompts.py +++ b/tests/test_prompts.py @@ -13,40 +13,35 @@ # limitations under the License. -from unittest.mock import MagicMock +from transformers import AutoTokenizer from nemo_skills.prompt.utils import get_prompt, get_token_count -def test_get_token_count_string(): - tokenizer = MagicMock() - tokenizer.encode.return_value = [1, 2, 3, 4, 5] - assert get_token_count(tokenizer, "hello world") == 5 - tokenizer.encode.assert_called_once_with("hello world", add_special_tokens=False) - - -def test_get_token_count_messages_list_result(): - """Test with old HF tokenizer behavior that returns a list of token ids.""" - tokenizer = MagicMock() - tokenizer.apply_chat_template.return_value = [1, 2, 3, 4, 5, 6] - messages = [{"role": "user", "content": "hello"}] - assert get_token_count(tokenizer, messages) == 6 - - -def test_get_token_count_messages_dict_result(): - """Test with new HF tokenizer behavior that returns a dict with input_ids.""" - tokenizer = MagicMock() - tokenizer.apply_chat_template.return_value = { - "input_ids": [1, 2, 3, 4, 5, 6, 7], - "attention_mask": [1, 1, 1, 1, 1, 1, 1], - } +def test_get_token_count(): + tokenizer = AutoTokenizer.from_pretrained("nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", trust_remote_code=True) messages = [{"role": "user", "content": "hello"}] - assert get_token_count(tokenizer, messages) == 7 + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + }, + } + ] -def test_get_token_count_none(): + assert get_token_count(tokenizer, "hello") == 1 + assert get_token_count(tokenizer, messages) == 17 + assert get_token_count(tokenizer, messages, tools=tools) == 266 assert get_token_count(None, "hello") is None - assert get_token_count(MagicMock(), None) is None + assert get_token_count(tokenizer, None) is None def test_generic_math_problem_augmentation_prompt(): From aa9e7fd3468fd56ae0951f3b65758f6a7567ccf0 Mon Sep 17 00:00:00 2001 From: Igor Gitman Date: Fri, 20 Feb 2026 14:13:33 -0800 Subject: [PATCH 4/5] Debugging Signed-off-by: Igor Gitman --- nemo_skills/prompt/utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/nemo_skills/prompt/utils.py b/nemo_skills/prompt/utils.py index a34429cc2c..c2d78903fd 100644 --- a/nemo_skills/prompt/utils.py +++ b/nemo_skills/prompt/utils.py @@ -397,8 +397,11 @@ def message_to_dict(orig_message: Any) -> Dict[str, Any]: try: result = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, tools=tools) # Handle newer HF tokenizer versions that return a dict instead of a list + print(result, type(result)) if isinstance(result, dict): result = result["input_ids"] + print("!@#") + print(len(result)) return len(result) except Exception as e: From 70ab73e9be13f7f7d178bf7dcd3f85f9bcb4ea20 Mon Sep 17 00:00:00 2001 From: Igor Gitman Date: Fri, 20 Feb 2026 14:16:04 -0800 Subject: [PATCH 5/5] Fix isinstance check Signed-off-by: Igor Gitman --- nemo_skills/inference/model/utils.py | 4 ++-- nemo_skills/inference/prover.py | 4 ++-- nemo_skills/prompt/utils.py | 7 ++----- 3 files changed, 6 insertions(+), 9 deletions(-) diff --git a/nemo_skills/inference/model/utils.py b/nemo_skills/inference/model/utils.py index 4d3f0683fa..e4e19dc437 100644 --- a/nemo_skills/inference/model/utils.py +++ b/nemo_skills/inference/model/utils.py @@ -96,8 +96,8 @@ def encode(self, prompt: str | list[dict], tools=None) -> list[int]: return self.tokenizer.encode(prompt) elif isinstance(prompt, list): result = self.tokenizer.apply_chat_template(prompt, add_generation_prompt=True, tools=tools) - # Handle newer HF tokenizer versions that return a dict instead of a list - if isinstance(result, dict): + # Handle newer HF tokenizer versions that return a BatchEncoding instead of a list + if not isinstance(result, list): result = result["input_ids"] return result diff --git a/nemo_skills/inference/prover.py b/nemo_skills/inference/prover.py index 37c2473e80..6371b5df72 100644 --- a/nemo_skills/inference/prover.py +++ b/nemo_skills/inference/prover.py @@ -283,8 +283,8 @@ async def _single_data_point_generate(self, data_point, data): prefix_tokens = self.hf_tokenizer.apply_chat_template( prepared_conversation, tokenize=True, add_generation_prompt=True ) - # Handle newer HF tokenizer versions that return a dict instead of a list - if isinstance(prefix_tokens, dict): + # Handle newer HF tokenizer versions that return a BatchEncoding instead of a list + if not isinstance(prefix_tokens, list): prefix_tokens = prefix_tokens["input_ids"] num_tokens_prefix = len(prefix_tokens) prefix = self.hf_tokenizer.apply_chat_template( diff --git a/nemo_skills/prompt/utils.py b/nemo_skills/prompt/utils.py index c2d78903fd..e258c9c19b 100644 --- a/nemo_skills/prompt/utils.py +++ b/nemo_skills/prompt/utils.py @@ -396,12 +396,9 @@ def message_to_dict(orig_message: Any) -> Dict[str, Any]: ] try: result = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, tools=tools) - # Handle newer HF tokenizer versions that return a dict instead of a list - print(result, type(result)) - if isinstance(result, dict): + # Handle newer HF tokenizer versions that return a BatchEncoding instead of a list + if not isinstance(result, list): result = result["input_ids"] - print("!@#") - print(len(result)) return len(result) except Exception as e: