Generation time + Input Sequence Length

nemo_skills/inference/generate.py

@@ -471,22 +471,10 @@ 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)
 
             for key in output:
@@ -520,8 +508,25 @@ async def process_single_datapoint(self, data_point, all_data):
             if self.cfg.override_max_code_executions and self.cfg.total_code_executions_in_prompt is not None:
                 generation_params["max_code_executions"] = data_point["total_code_executions"]
 
+        # Tracking the tokens and generation time
+        input_sequence_length = None
+        if self.prompt is not None:
+            if generation_params["prompt"] is not None:
+                input_sequence_length = self.prompt.get_token_count(generation_params["prompt"])
+
+        # Start to track the generation time
+        start_time = time.time()
+
         result = await self.llm.generate_async(**generation_params)
 
+        end_time = time.time()
+        # Add the generation time and input sequence length
+        if self.cfg.add_generation_stats:
+            result["generation_start_time"] = start_time
+            result["generation_end_time"] = end_time
+            result["generation_time"] = end_time - start_time
+
+        result["input_sequence_length"] = input_sequence_length
         return result
 
     async def apply_evaluation_hook(self, data_point):
@@ -536,9 +541,6 @@ 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
             output = await self.process_single_datapoint(data_point, all_data)
             # Apply evaluation hook if configured

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
@@ -302,6 +302,29 @@ def fill(
 
         return messages
 
+    def get_token_count(self, messages: Union[str, list[dict]]) -> int:
+        """
+        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 self.tokenizer is None:
+            return None
+
+        if isinstance(messages, str):
+            return len(self.tokenizer.encode(messages, add_special_tokens=False))
+        elif isinstance(messages, list):
+            try:
+                return len(self.tokenizer.apply_chat_template(messages, tokenize=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 __str__(self):
         return str(self.config)