diff --git a/nemo_skills/inference/generate.py b/nemo_skills/inference/generate.py index 638289e821..957d93edfe 100644 --- a/nemo_skills/inference/generate.py +++ b/nemo_skills/inference/generate.py @@ -29,6 +29,7 @@ import litellm from omegaconf import ListConfig from tqdm import tqdm +from transformers import AutoTokenizer from nemo_skills.code_execution.sandbox import get_sandbox, sandbox_params from nemo_skills.inference.model import ( @@ -39,7 +40,7 @@ get_tool_calling_model, server_params, ) -from nemo_skills.prompt.utils import get_prompt +from nemo_skills.prompt.utils import get_prompt, get_token_count from nemo_skills.utils import ( chunk_data, get_help_message, @@ -113,6 +114,9 @@ class GenerateSolutionsConfig: # Useful when running judge jobs to keep the original generation statistics add_generation_stats: bool = True + # Count the number of tokens in the prompt + count_prompt_tokens: bool = False + generation_key: str = "generation" async_position_key: str = "_async_position" # key to use for preserving position in async loop in data dict @@ -268,7 +272,11 @@ def __init__(self, cfg: GenerateSolutionsConfig): self.cfg.chat_template_kwargs = None # Setup tokenizer - if self.cfg.use_completions_api or self.cfg.server.get("enable_soft_fail", False): + if ( + self.cfg.use_completions_api + or self.cfg.server.get("enable_soft_fail", False) + or self.cfg.count_prompt_tokens + ): # These are the only cases where we need a tokenizer self.tokenizer = self.cfg.tokenizer or self.cfg.server["model"] else: @@ -284,6 +292,17 @@ def __init__(self, cfg: GenerateSolutionsConfig): self.prompt = self.setup_prompt() self.llm = self.setup_llm() + # Setup hf_tokenizer for counting prompt tokens + self.hf_tokenizer = None + if self.cfg.count_prompt_tokens: + if self.prompt is not None: + self.hf_tokenizer = self.prompt.tokenizer + else: + self.hf_tokenizer = AutoTokenizer.from_pretrained(self.tokenizer) + + if self.hf_tokenizer is None: + raise ValueError("Tokenizer could not be initialized. Needed for counting prompt tokens.") + if self.cfg.code_execution: self.extra_generate_params = self.prompt.get_code_execution_args() else: @@ -331,6 +350,7 @@ def setup_prompt(self): examples_type=self.cfg.examples_type, system_message=self.cfg.system_message, ) + LOG.info("Prompt used: %s", prompt) return prompt @@ -471,23 +491,12 @@ def dump_outputs(self, outputs, data_points, fout): # all of the ground-truth data to the output file alongside the generated solutions output[self.cfg.generation_key] = output.pop("generation") - # calculating total generation time - if self.cfg.add_generation_stats: - output["generation_end_time"] = time.time() - # TODO: start time is saved in data_point, not output, need to fix that - output["generation_time"] = ( - output["generation_end_time"] - original_data_point["generation_start_time"] - ) - else: - # generation_start_time was overriden, so restoring it from end and total - # TODO: this is a bit hacky, need a rewrite - if "generation_end_time" in original_data_point and "generation_time" in original_data_point: - output["generation_start_time"] = ( - original_data_point["generation_end_time"] - original_data_point["generation_time"] - ) - else: - output.pop("generation_start_time", None) + if not self.cfg.add_generation_stats: + output.pop("generation_start_time", None) + output.pop("generation_end_time", None) + output.pop("generation_time", None) output.pop("num_generated_tokens", None) + output.pop("input_sequence_length", None) for key in output: original_data_point.pop(key, None) @@ -522,6 +531,9 @@ async def process_single_datapoint(self, data_point, all_data): result = await self.llm.generate_async(**generation_params) + if self.cfg.count_prompt_tokens: + input_sequence_length = get_token_count(self.hf_tokenizer, generation_params["prompt"]) + result["input_sequence_length"] = input_sequence_length return result async def apply_evaluation_hook(self, data_point): @@ -536,11 +548,16 @@ async def apply_evaluation_hook(self, data_point): async def _process_single_datapoint_with_semaphore(self, data_point, all_data, fout, pbar): """Process a single data point with semaphore control.""" async with self.semaphore: - # registering current time to calculate total generation time - data_point["generation_start_time"] = time.time() - # Generate output for this single data point + start_time = time.time() output = await self.process_single_datapoint(data_point, all_data) + end_time = time.time() + + if self.cfg.add_generation_stats: + output["generation_start_time"] = start_time + output["generation_end_time"] = end_time + output["generation_time"] = end_time - start_time + # Apply evaluation hook if configured # TODO: note that this currently only evaluates independently--if there # is any post-processing that needs to be done on the full set of diff --git a/nemo_skills/prompt/utils.py b/nemo_skills/prompt/utils.py index ca89ddeec9..6c7bb955d9 100644 --- a/nemo_skills/prompt/utils.py +++ b/nemo_skills/prompt/utils.py @@ -18,7 +18,7 @@ import re from dataclasses import asdict, field from pathlib import Path -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Union import yaml from transformers import AutoTokenizer @@ -306,6 +306,33 @@ def __str__(self): return str(self.config) +def get_token_count(tokenizer, messages: Union[str, list[dict]]) -> int | None: + """ + Count the number of tokens in a string or chat message list. + + Args: + messages (str | list[dict]): Input text or chat messages. + + Returns: + int | None: Token count, or None if no tokenizer is set. + """ + if tokenizer is None: + return None + + if messages is None: + return None + + if isinstance(messages, str): + return len(tokenizer.encode(messages, add_special_tokens=False)) + elif isinstance(messages, list): + try: + return len(tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True)) + 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") + + def get_config_path(config: str, config_dir: str | None = None, config_extension: str = "yaml") -> Path: if config_dir is None: config_dir = str(Path(__file__).parent.absolute() / "config") diff --git a/tests/test_generation.py b/tests/test_generation.py index bc1e0eba8d..0d0f42c586 100644 --- a/tests/test_generation.py +++ b/tests/test_generation.py @@ -26,7 +26,7 @@ def test_eval_gsm8k_api(tmp_path): cmd = ( f"ns eval " f" --server_type=openai " - f" --model=meta/llama-3.1-8b-instruct " + f" --model=nvidia/nvidia-nemotron-nano-9b-v2 " f" --server_address=https://integrate.api.nvidia.com/v1 " f" --benchmarks=gsm8k " f" --output_dir={tmp_path} " @@ -53,11 +53,11 @@ def test_eval_judge_api(tmp_path): cmd = ( f"ns eval " f" --server_type=openai " - f" --model=meta/llama-3.1-8b-instruct " + f" --model=nvidia/nvidia-nemotron-nano-9b-v2 " f" --server_address=https://integrate.api.nvidia.com/v1 " f" --benchmarks=math-500 " f" --output_dir={tmp_path} " - f" --judge_model=meta/llama-3.1-8b-instruct " + f" --judge_model=nvidia/nvidia-nemotron-nano-9b-v2 " f" --judge_server_address=https://integrate.api.nvidia.com/v1 " f" --judge_server_type=openai " f" --judge_generation_type=math_judge " @@ -85,7 +85,7 @@ def test_fail_on_api_key_env_var(tmp_path): cmd = ( f"ns eval " f" --server_type=openai " - f" --model=meta/llama-3.1-8b-instruct " + f" --model=nvidia/nvidia-nemotron-nano-9b-v2 " f" --server_address=https://integrate.api.nvidia.com/v1 " f" --benchmarks=gsm8k " f" --output_dir={tmp_path} " @@ -106,7 +106,7 @@ def test_succeed_on_api_key_env_var(tmp_path): f"unset NVIDIA_API_KEY && " f"ns eval " f" --server_type=openai " - f" --model=meta/llama-3.1-8b-instruct " + f" --model=nvidia/nvidia-nemotron-nano-9b-v2 " f" --server_address=https://integrate.api.nvidia.com/v1 " f" --benchmarks=gsm8k " f" --output_dir={tmp_path} " @@ -135,7 +135,7 @@ def test_generate_openai_format(tmp_path, format): cmd = ( f"ns generate " f" --server_type=openai " - f" --model=meta/llama-3.1-8b-instruct " + f" --model=nvidia/nvidia-nemotron-nano-9b-v2 " f" --server_address=https://integrate.api.nvidia.com/v1 " f" --input_file=/nemo_run/code/tests/data/openai-input-{format}.test " f" --output_dir={tmp_path} "