diff --git a/src/nemo_safe_synthesizer/config/external_results.py b/src/nemo_safe_synthesizer/config/external_results.py index 345c16526..a24a141cb 100644 --- a/src/nemo_safe_synthesizer/config/external_results.py +++ b/src/nemo_safe_synthesizer/config/external_results.py @@ -57,7 +57,18 @@ def log_wandb(self, run: wandb.Run | None = None) -> None: class SafeSynthesizerSummary(NSSBaseModel): - """Aggregated quality, privacy, and record-count metrics for a pipeline run.""" + """Aggregated quality, privacy, and record-count metrics for a pipeline run. + + Token-field invariants (when all referenced fields are populated): + + * ``num_non_record_tokens == num_completion_tokens + - num_valid_record_tokens - num_invalid_record_tokens`` + (clamped to 0 if slight tokenizer-boundary drift makes the + subtraction negative). + * ``tokens_per_prompt == num_completion_tokens / num_prompts``. + * ``valid_record_token_fraction == num_valid_record_tokens + / num_completion_tokens``. + """ synthetic_data_quality_score: float | None = Field( default=None, @@ -116,6 +127,37 @@ class SafeSynthesizerSummary(NSSBaseModel): description="Ratio of valid records: ``num_valid_records / (num_valid_records + num_invalid_records)``.", ) + num_completion_tokens: int | None = Field( + default=None, description="Total tokens generated by the LLM across all completions." + ) + + num_valid_record_tokens: int | None = Field(default=None, description="Tokens in records that passed validation.") + + num_invalid_record_tokens: int | None = Field(default=None, description="Tokens in records that failed validation.") + + num_non_record_tokens: int | None = Field(default=None, description="Tokens not part of any recognized record.") + + valid_record_token_fraction: float | None = Field( + default=None, description="Fraction of total completion tokens in valid records." + ) + + tokens_per_prompt: float | None = Field( + default=None, + description="Average completion tokens per prompt: num_completion_tokens / num_prompts.", + ) + + tokens_per_second: float | None = Field( + default=None, description="Total completion tokens divided by generation wall-clock time." + ) + + valid_tokens_per_second: float | None = Field( + default=None, description="Valid record tokens divided by generation wall-clock time." + ) + + tokenization_overhead_sec: float | None = Field( + default=None, description="Wall-clock seconds spent on tokenization for statistics tracking." + ) + timing: SafeSynthesizerTiming = Field(description="Per-stage wall-clock durations.") def log_summary(self, logger: logging.Logger) -> None: @@ -130,30 +172,42 @@ def log_wandb(self) -> None: import wandb if wandb.run is not None: - wandb.log( - { - "gen/generation_time_sec": self.timing.generation_time_sec, - "gen/evaluation_time_sec": self.timing.evaluation_time_sec, - "eval/total_time_sec": self.timing.total_time_sec, - "train/pii_replacer_time_sec": self.timing.pii_replacer_time_sec, - "train/training_time_sec": self.timing.training_time_sec, - "gen/num_valid_records": self.num_valid_records, - "gen/num_invalid_records": self.num_invalid_records, - "gen/num_prompts": self.num_prompts, - "gen/valid_record_fraction": self.valid_record_fraction, - "eval/data_privacy_score": self.data_privacy_score, - "eval/membership_inference_protection_score": self.membership_inference_protection_score, - "eval/attribute_inference_protection_score": self.attribute_inference_protection_score, - "eval/synthetic_data_quality_score": self.synthetic_data_quality_score, - "eval/column_correlation_stability_score": self.column_correlation_stability_score, - "eval/deep_structure_stability_score": self.deep_structure_stability_score, - "eval/column_distribution_stability_score": self.column_distribution_stability_score, - "eval/text_semantic_similarity_score": self.text_semantic_similarity_score, - "eval/text_structure_similarity_score": self.text_structure_similarity_score, - "eval/success": 1 - if self.data_privacy_score is not None - and self.synthetic_data_quality_score is not None - and self.synthetic_data_quality_score > 0 - else 0, - } - ) + metrics: dict[str, float | int | None] = { + "gen/generation_time_sec": self.timing.generation_time_sec, + "gen/evaluation_time_sec": self.timing.evaluation_time_sec, + "eval/total_time_sec": self.timing.total_time_sec, + "train/pii_replacer_time_sec": self.timing.pii_replacer_time_sec, + "train/training_time_sec": self.timing.training_time_sec, + "gen/num_valid_records": self.num_valid_records, + "gen/num_invalid_records": self.num_invalid_records, + "gen/num_prompts": self.num_prompts, + "gen/valid_record_fraction": self.valid_record_fraction, + "gen/num_completion_tokens": self.num_completion_tokens, + "gen/num_valid_record_tokens": self.num_valid_record_tokens, + "gen/num_invalid_record_tokens": self.num_invalid_record_tokens, + "gen/num_non_record_tokens": self.num_non_record_tokens, + "gen/valid_record_token_fraction": self.valid_record_token_fraction, + "gen/tokens_per_prompt": self.tokens_per_prompt, + "gen/tokens_per_second": self.tokens_per_second, + "gen/valid_tokens_per_second": self.valid_tokens_per_second, + "gen/tokenization_overhead_sec": self.tokenization_overhead_sec, + "eval/data_privacy_score": self.data_privacy_score, + "eval/membership_inference_protection_score": self.membership_inference_protection_score, + "eval/attribute_inference_protection_score": self.attribute_inference_protection_score, + "eval/synthetic_data_quality_score": self.synthetic_data_quality_score, + "eval/column_correlation_stability_score": self.column_correlation_stability_score, + "eval/deep_structure_stability_score": self.deep_structure_stability_score, + "eval/column_distribution_stability_score": self.column_distribution_stability_score, + "eval/text_semantic_similarity_score": self.text_semantic_similarity_score, + "eval/text_structure_similarity_score": self.text_structure_similarity_score, + "eval/success": 1 + if self.data_privacy_score is not None + and self.synthetic_data_quality_score is not None + and self.synthetic_data_quality_score > 0 + else 0, + } + # Log ``None`` values too (rather than filtering them out) so that + # dashboard comparisons clearly show when a stage was skipped or a + # metric was not collected for a given run, instead of the metric + # silently carrying its last-logged value on the W&B chart. + wandb.log(metrics) diff --git a/src/nemo_safe_synthesizer/data_processing/record_utils.py b/src/nemo_safe_synthesizer/data_processing/record_utils.py index 4cc01e39d..3a62f8c23 100644 --- a/src/nemo_safe_synthesizer/data_processing/record_utils.py +++ b/src/nemo_safe_synthesizer/data_processing/record_utils.py @@ -12,7 +12,10 @@ import calendar import json import re +import time +from collections.abc import Callable from csv import QUOTE_NONNUMERIC +from dataclasses import dataclass, field from datetime import datetime from io import StringIO @@ -27,6 +30,91 @@ logger = get_logger() +@dataclass +class ParsedRecord: + """A single record extracted from an LLM completion. + + Validity is tracked by the invariant that exactly one of ``parsed`` + and ``error`` is non-``None``: a valid record has ``parsed`` set and + ``error`` as ``None``, an invalid record has ``error`` set and + ``parsed`` as ``None``. + [`is_valid`][nemo_safe_synthesizer.data_processing.record_utils.ParsedRecord.is_valid] + is the canonical accessor. + + ``text`` and ``token_count`` are captured at extraction time and + remain invariant even if the record is reclassified later (e.g. by + group-level checks or data-fidelity filters) via + [`invalidate`][nemo_safe_synthesizer.data_processing.record_utils.ParsedRecord.invalidate]. + """ + + text: str + """Original regex-matched JSON string (invariant under reclassification).""" + + parsed: dict | None = None + """Parsed dict when validation succeeded, ``None`` when invalid.""" + + error: tuple[str, str] | None = None + """``(detailed_msg, validator)`` when invalid, ``None`` when valid.""" + + token_count: int = 0 + """Number of tokens in ``text``; 0 when no tokenizer was provided.""" + + @property + def is_valid(self) -> bool: + """Return ``True`` when this record passed validation.""" + return self.error is None + + def invalidate(self, error: tuple[str, str]) -> None: + """Reclassify this record as invalid. + + ``text`` and ``token_count`` are kept intact; ``parsed`` is + cleared so downstream consumers don't accidentally use a stale + dict. + + Args: + error: ``(detailed_msg, validator)`` tuple describing the + reason for invalidation. + """ + self.error = error + self.parsed = None + + +@dataclass +class ParsedResponse: + """Parsed result of a single LLM prompt response. + + Holds a flat list of + [`ParsedRecord`][nemo_safe_synthesizer.data_processing.record_utils.ParsedRecord] + objects plus aggregated tokenization timing. The ``valid_records`` + / ``invalid_records`` / ``errors`` views are derived for backward + compatibility with callers that want the old parallel-list shape. + """ + + records: list[ParsedRecord] = field(default_factory=list) + """Per-record extraction + validation outcomes, in input order.""" + + tokenization_time_sec: float = 0.0 + """Wall-clock seconds spent tokenizing records in this response.""" + + prompt_number: int | None = None + """Index of the prompt within the batch (set by the processor call).""" + + @property + def valid_records(self) -> list[dict]: + """Parsed dicts for records that passed validation.""" + return [r.parsed for r in self.records if r.is_valid and r.parsed is not None] + + @property + def invalid_records(self) -> list[str]: + """Original text for records that failed validation.""" + return [r.text for r in self.records if not r.is_valid] + + @property + def errors(self) -> list[tuple[str, str]]: + """``(detailed_msg, validator)`` tuples for each invalid record.""" + return [r.error for r in self.records if r.error is not None] + + def is_safe_for_float_conversion(value: str | int | float | None | list | dict) -> bool: """Check if a value can be safely converted to float64 without overflow. @@ -111,35 +199,67 @@ def extract_groups_from_jsonl_string(jsonl_string: str, bos: str, eos: str) -> l return re.findall(rf"{bos_re}\s?(?:{RECORD_REGEX_PATTERN}\s?)+\s?{eos_re}", jsonl_string) +def timed_encode( + encode: Callable[[str], list[int]] | None, +) -> Callable[[str], tuple[int, float]]: + """Wrap an encode callable with timing, or return a no-op. + + Returns a function ``timed(text)`` that returns ``(n_tokens, + elapsed_seconds)``. When *encode* is ``None`` the returned + function always returns ``(0, 0.0)``. + """ + if encode is None: + + def _noop(_text: str) -> tuple[int, float]: + return 0, 0.0 + + return _noop + + def _timed(text: str) -> tuple[int, float]: + t0 = time.monotonic() + n = len(encode(text)) + return n, time.monotonic() - t0 + + return _timed + + def extract_and_validate_records( - jsonl_string: str, schema: dict -) -> tuple[list[dict], list[str], list[tuple[str, str]]]: + jsonl_string: str, + schema: dict, + encode: Callable[[str], list[int]] | None = None, +) -> ParsedResponse: """Extract and validate records from the given JSONL string. - The records are validated against the given schema using jsonschema. + Each regex-matched JSON string is tokenized (when *encode* is + provided) before validation so that exact token counts are + available for every record regardless of later reclassification. Args: jsonl_string: Single JSONL string containing tabular records. schema: JSON schema as a dictionary. + encode: Optional tokenizer encode callable. When provided, + each matched record string is tokenized and its token count + is stored on the corresponding + [`ParsedRecord`][nemo_safe_synthesizer.data_processing.record_utils.ParsedRecord]. Returns: - valid_records (list[dict]): List of valid records. - invalid_records (list[str]): List of invalid records. - invalid_record_errors (list[tuple[str, str]]): List of errors for invalid records, each a (message, validator) tuple. + A + [`ParsedResponse`][nemo_safe_synthesizer.data_processing.record_utils.ParsedResponse] + whose ``records`` list is in input order, with ``parsed`` set + for valid records and ``error`` set for invalid ones. """ - valid_records = [] - invalid_records = [] - invalid_record_errors = [] + records: list[ParsedRecord] = [] + tokenization_time = 0.0 + timed = timed_encode(encode) for matched_json in extract_records_from_jsonl_string(jsonl_string): - matched_dict, error = _parse_and_validate_json(matched_json, schema) - if error: - invalid_records.append(matched_json) - invalid_record_errors.append(error) - else: - valid_records.append(matched_dict) + n_tokens, dt = timed(matched_json) + tokenization_time += dt - return valid_records, invalid_records, invalid_record_errors + parsed, error = _parse_and_validate_json(matched_json, schema) + records.append(ParsedRecord(text=matched_json, parsed=parsed, error=error, token_count=n_tokens)) + + return ParsedResponse(records=records, tokenization_time_sec=tokenization_time) def _parse_timestamp_to_seconds(value: object, time_format: str) -> int: @@ -295,23 +415,37 @@ def extract_and_validate_timeseries_records( time_column: str, interval_seconds: int | None, time_format: str, -) -> tuple[list[dict], list[str], list[tuple[str, str]]]: - """Extract and validate sequential records with enforced time interval constraints. + encode: Callable[[str], list[int]] | None = None, +) -> ParsedResponse: + """Extract and validate sequential records with time-interval constraints. + + Each regex-matched JSON string is tokenized (when *encode* is + provided) before validation so that exact token counts are captured + for both validated and cascade-invalidated records. Args: jsonl_string: JSONL string containing series data. schema: JSON schema describing the records. - time_column: Column containing the timestamp used for interval validation. - interval_seconds: (Optional) Expected interval in seconds between consecutive - timestamps. If not provided, no time interval validation is performed. - time_format: Format of the timestamp column (required, should be set from config). + time_column: Column containing the timestamp used for interval + validation. + interval_seconds: Expected interval in seconds between + consecutive timestamps. When ``None``, no interval check + is performed. + time_format: Format of the timestamp column (required). + encode: Optional tokenizer encode callable. When provided, + each matched record string is tokenized and its token count + is stored on the corresponding + [`ParsedRecord`][nemo_safe_synthesizer.data_processing.record_utils.ParsedRecord]. Returns: - Tuple of valid records, invalid record strings, and their associated errors. + A + [`ParsedResponse`][nemo_safe_synthesizer.data_processing.record_utils.ParsedResponse] + in input order. Once a record fails, every subsequent record is + marked invalid with a cascade error. """ - valid_records: list[dict] = [] - invalid_records: list[str] = [] - invalid_record_errors: list[tuple[str, str]] = [] + records: list[ParsedRecord] = [] + tokenization_time = 0.0 + timed = timed_encode(encode) last_absolute_seconds: int | None = None day_offset = 0 @@ -326,25 +460,27 @@ def extract_and_validate_timeseries_records( allow_rollover = not has_date all_json_records = list(extract_records_from_jsonl_string(jsonl_string)) + cascade_error = ("Invalid due to previous record error", "TimeSeries") for idx, matched_json in enumerate(all_json_records): - # Step 1: Parse and validate JSON/schema - matched_dict, error = _parse_and_validate_json(matched_json, schema) - if error or matched_dict is None: - invalid_records.append(matched_json) - if error: - invalid_record_errors.append(error) + n_tokens, dt = timed(matched_json) + tokenization_time += dt + + # Step 1: Parse and validate JSON/schema. + parsed, error = _parse_and_validate_json(matched_json, schema) + if error or parsed is None: + records.append(ParsedRecord(text=matched_json, error=error, token_count=n_tokens)) + # Parse/schema errors stop validation without cascading to later records. break - # Step 2: Extract and parse timestamp - timestamp_seconds, error = _extract_timestamp_seconds(matched_dict, time_column, time_format) + # Step 2: Extract and parse timestamp. + timestamp_seconds, error = _extract_timestamp_seconds(parsed, time_column, time_format) if error or timestamp_seconds is None: - invalid_records.append(matched_json) - if error: - invalid_record_errors.append(error) + records.append(ParsedRecord(text=matched_json, error=error, token_count=n_tokens)) + # Missing timestamp stops validation without cascading to later records. break - # Step 3: Validate time interval (if interval_seconds is specified) + # Step 3: Validate time interval (if interval_seconds is specified). if interval_seconds is not None: absolute_seconds, day_offset, error = _validate_time_interval( timestamp_seconds, @@ -355,20 +491,19 @@ def extract_and_validate_timeseries_records( allow_rollover, ) if error: - # Mark current record with the specific error, and remaining records with cascade error - invalid_records.append(matched_json) - invalid_record_errors.append(error) - # Mark remaining records (after current) as invalid due to previous error - remaining_records = all_json_records[idx + 1 :] - cascade_error = ("Invalid due to previous record error", "TimeSeries") - invalid_records.extend(remaining_records) - invalid_record_errors.extend([cascade_error] * len(remaining_records)) + records.append(ParsedRecord(text=matched_json, error=error, token_count=n_tokens)) + # Interval errors cascade: mark all remaining records invalid so the + # caller can report how many were affected. + for remaining in all_json_records[idx + 1 :]: + rem_tokens, rem_dt = timed(remaining) + tokenization_time += rem_dt + records.append(ParsedRecord(text=remaining, error=cascade_error, token_count=rem_tokens)) break last_absolute_seconds = absolute_seconds - valid_records.append(matched_dict) + records.append(ParsedRecord(text=matched_json, parsed=parsed, token_count=n_tokens)) - return valid_records, invalid_records, invalid_record_errors + return ParsedResponse(records=records, tokenization_time_sec=tokenization_time) def normalize_dataframe(dataframe: pd.DataFrame) -> pd.DataFrame: diff --git a/src/nemo_safe_synthesizer/generation/batch.py b/src/nemo_safe_synthesizer/generation/batch.py index e7d614be7..8a7eb9480 100644 --- a/src/nemo_safe_synthesizer/generation/batch.py +++ b/src/nemo_safe_synthesizer/generation/batch.py @@ -93,6 +93,7 @@ class Batch: def __init__(self, processor: Processor): self._responses: list[ParsedResponse] = [] self._processor = processor + self._total_completion_tokens: int = 0 @property def num_prompts(self) -> int: @@ -125,6 +126,49 @@ def valid_record_fraction(self) -> float: total_records = self.num_valid_records + self.num_invalid_records return 0 if total_records == 0 else self.num_valid_records / total_records + @property + def total_completion_tokens(self) -> int: + """Total tokens across all completions in this batch (from vLLM output).""" + return self._total_completion_tokens + + @property + def total_valid_record_tokens(self) -> int: + """Sum of token counts for all valid records in this batch.""" + return sum(record.token_count for resp in self._responses for record in resp.records if record.is_valid) + + @property + def total_invalid_record_tokens(self) -> int: + """Sum of token counts for all invalid records in this batch.""" + return sum(record.token_count for resp in self._responses for record in resp.records if not record.is_valid) + + @property + def total_non_record_tokens(self) -> int: + """Tokens not part of any recognized record (completion - valid - invalid). + + Clamped to zero if negative (possible due to tokenizer boundary effects). + """ + value = self._total_completion_tokens - self.total_valid_record_tokens - self.total_invalid_record_tokens + if value < 0: + logger.warning( + "Non-record token count is negative; clamping to 0. " + "This can happen due to tokenizer boundary effects between individual records and full completions.", + extra={ + "ctx": { + "delta": value, + "completion_tokens": self._total_completion_tokens, + "valid_record_tokens": self.total_valid_record_tokens, + "invalid_record_tokens": self.total_invalid_record_tokens, + } + }, + ) + return 0 + return value + + @property + def total_tokenization_time_sec(self) -> float: + """Wall-clock seconds spent tokenizing records in this batch.""" + return sum(r.tokenization_time_sec for r in self._responses) + def error_statistics(self, detailed_errors: bool) -> pd.DataFrame: """Return count statistics on errors encountered during generation. @@ -201,7 +245,7 @@ def log_summary(self, detailed_errors: bool = False) -> None: err_stats: pd.DataFrame = self.error_statistics(detailed_errors=detailed_errors) # Build structured summary data - processor renders as table for console - summary_data = { + summary_data: dict[str, int | float] = { "num_prompts": self.num_prompts, "num_valid_records": self.num_valid_records, "num_invalid_records": self.num_invalid_records, @@ -209,6 +253,11 @@ def log_summary(self, detailed_errors: bool = False) -> None: } if self.num_data_config_rejected_records: summary_data["num_data_config_rejected_records"] = self.num_data_config_rejected_records + if self._total_completion_tokens > 0: + summary_data["completion_tokens"] = self._total_completion_tokens + summary_data["valid_record_tokens"] = self.total_valid_record_tokens + summary_data["invalid_record_tokens"] = self.total_invalid_record_tokens + summary_data["non_record_tokens"] = self.total_non_record_tokens # Pass structured data - processor renders for console, JSON keeps as-is logger.user.info( @@ -237,11 +286,15 @@ def log_summary(self, detailed_errors: bool = False) -> None: }, ) - def process(self, prompt_number: int, text: str) -> None: + def process(self, prompt_number: int, text: str, completion_tokens: int = 0) -> None: """Process text response from a single prompt in the current batch. Args: prompt_number: The prompt number in the current batch. text: Text generated by the fine-tuned model. + completion_tokens: Number of tokens in this completion + (from vLLM ``output.token_ids``). Defaults to 0 when + the backend does not provide token counts. """ self._responses.append(self._processor(prompt_number, text)) + self._total_completion_tokens += completion_tokens diff --git a/src/nemo_safe_synthesizer/generation/processors.py b/src/nemo_safe_synthesizer/generation/processors.py index 02adb410b..7ff18d368 100644 --- a/src/nemo_safe_synthesizer/generation/processors.py +++ b/src/nemo_safe_synthesizer/generation/processors.py @@ -6,12 +6,13 @@ from __future__ import annotations from abc import ABC, abstractmethod -from dataclasses import dataclass -from typing import Any +from typing import TYPE_CHECKING, Any from ..config import SafeSynthesizerParameters from ..config.generate import ValidationParameters from ..data_processing.record_utils import ( + ParsedRecord, + ParsedResponse, check_if_records_are_ordered, extract_and_validate_records, extract_and_validate_timeseries_records, @@ -20,25 +21,22 @@ from ..llm.metadata import ModelMetadata from ..observability import get_logger -logger = get_logger(__name__) - - -@dataclass -class ParsedResponse: - """Parsed result of a single LLM prompt response. +if TYPE_CHECKING: + from transformers import PreTrainedTokenizerBase - Attributes: - valid_records: Records that passed schema validation (as dicts). - invalid_records: Raw text of records that failed validation. - errors: ``(detailed_msg, summary_msg)`` tuples for each invalid - record. - prompt_number: Index of the prompt in the batch. - """ +logger = get_logger(__name__) - valid_records: list[dict] - invalid_records: list[str] - errors: list[tuple[str, str]] - prompt_number: int | None = None +# Re-export the parsed-record types so existing imports of +# ``ParsedResponse`` / ``ParsedRecord`` from this module keep working. +__all__ = [ + "GroupedDataProcessor", + "ParsedRecord", + "ParsedResponse", + "Processor", + "TabularDataProcessor", + "TimeSeriesDataProcessor", + "create_processor", +] class Processor(ABC): @@ -46,12 +44,25 @@ class Processor(ABC): Args: schema: JSON schema as a dictionary. + config: Validation parameters. + tokenizer: Optional tokenizer for exact per-record token + counting. ``None`` is supported for tests and the training + eval callback path when the tokenizer isn't readily + available; in that case ``ParsedRecord.token_count`` is 0. """ - def __init__(self, schema: dict[str, Any], config: ValidationParameters): + def __init__( + self, + schema: dict[str, Any], + config: ValidationParameters, + tokenizer: PreTrainedTokenizerBase | None = None, + ): self.schema = schema self.config = config - logger.debug(f"Initialized processor with schema={schema} and config={config}") + self._tokenizer = tokenizer + logger.debug( + f"Initialized processor with schema={schema}, config={config}, tokenizer_present={tokenizer is not None}" + ) @property def name(self) -> str: @@ -90,15 +101,27 @@ def __call__(self, prompt_number: int, text: str) -> ParsedResponse: Parsed response containing valid records, invalid records, and error tuples. """ - # Make sure the string composed of valid utf-8 characters. text = text.encode("utf-8", "ignore").decode("utf-8") - # Call the processor-specific method to process the text. response = self._process_text_generation(text) response.prompt_number = prompt_number return response + @property + def _encode(self): + """Encode callable for the extraction helpers, or ``None``. + + Returns a ``Callable[[str], list[int]]`` that calls the + tokenizer with ``add_special_tokens=False``, or ``None`` when no + tokenizer was provided so the extractors can skip tokenization + entirely. + """ + if self._tokenizer is None: + return None + tokenizer = self._tokenizer + return lambda text: tokenizer.encode(text, add_special_tokens=False) + class TabularDataProcessor(Processor): """Processor for standard (non-grouped, non-time-series) tabular data.""" @@ -112,8 +135,7 @@ def _process_text_generation(self, text: str) -> ParsedResponse: Returns: Parsed response with validated records and error details. """ - valid, invalid, errors = extract_and_validate_records(text, self.schema) - return ParsedResponse(valid_records=valid, invalid_records=invalid, errors=errors) + return extract_and_validate_records(text, self.schema, encode=self._encode) class TimeSeriesDataProcessor(Processor): @@ -130,6 +152,8 @@ class TimeSeriesDataProcessor(Processor): timestamps, or ``None`` if intervals vary. time_format: Timestamp format string (``strptime``), or ``"elapsed_seconds"`` for numeric elapsed time. + tokenizer: Optional tokenizer for exact per-record token + counting. Raises: ValueError: If ``time_column`` or ``time_format`` is ``None``. @@ -142,8 +166,9 @@ def __init__( time_column: str | None, interval_seconds: int | None, time_format: str | None, + tokenizer: PreTrainedTokenizerBase | None = None, ): - super().__init__(schema=schema, config=config) + super().__init__(schema=schema, config=config, tokenizer=tokenizer) if time_column is None: raise ValueError( "time_column is required for TimeSeriesDataProcessor but was None. " @@ -169,10 +194,14 @@ def _process_text_generation(self, text: str) -> ParsedResponse: Returns: Parsed response with validated records and error details. """ - valid, invalid, errors = extract_and_validate_timeseries_records( - text, self.schema, self.time_column, self.interval_seconds, self.time_format + return extract_and_validate_timeseries_records( + text, + self.schema, + self.time_column, + self.interval_seconds, + self.time_format, + encode=self._encode, ) - return ParsedResponse(valid_records=valid, invalid_records=invalid, errors=errors) class GroupedDataProcessor(Processor): @@ -191,6 +220,8 @@ class GroupedDataProcessor(Processor): group_by: Column name that defines groups. order_by: Column name to enforce ordering within a group, or ``None`` if ordering is not required. + tokenizer: Optional tokenizer for exact per-record token + counting. """ def __init__( @@ -201,8 +232,9 @@ def __init__( eos_token: str, group_by: str, order_by: str | None = None, + tokenizer: PreTrainedTokenizerBase | None = None, ): - super().__init__(schema=schema, config=config) + super().__init__(schema=schema, config=config, tokenizer=tokenizer) self.group_by: list[str] = [group_by] self.order_by = order_by self.bos_token = bos_token @@ -220,6 +252,11 @@ def _process_text_generation(self, text: str) -> ParsedResponse: These requirements may be relaxed and automatically fixed depending on the settings in self.config. + Per-record token counts (captured at extraction time) ride along + with each :class:`ParsedRecord` through group-level + reclassification, so aggregate token metrics stay consistent + with the final valid/invalid partition. + Args: text: Text generated by the fine-tuned model. @@ -234,77 +271,102 @@ def _process_text_generation(self, text: str) -> ParsedResponse: if len(groups) == 0: return ParsedResponse( - valid_records=[], - invalid_records=[text], - errors=[("Group BOS and/or EOS tokens missing", groupby_validator)], + records=[ + ParsedRecord( + text=text, + error=("Group BOS and/or EOS tokens missing", groupby_validator), + ) + ] ) - valid_groups, invalid_groups, errors_groups = [], [], [] + all_records: list[ParsedRecord] = [] + total_tokenization_time = 0.0 for group in groups: - valid, invalid, errors = extract_and_validate_records(group, self.schema) - valid_with_str_members = [str(item) for item in valid] + response = extract_and_validate_records(group, self.schema, encode=self._encode) + total_tokenization_time += response.tokenization_time_sec + group_records = response.records - if len(valid) == 0: - invalid_groups.extend(invalid) - errors_groups.extend(errors) - continue + has_valid = any(r.is_valid for r in group_records) + has_invalid = any(not r.is_valid for r in group_records) # Handle invalid records in the group (optionally ignore and proceed). - if len(invalid) > 0: + if has_valid and has_invalid: if self.config.group_by_ignore_invalid_records: - invalid = [] - errors = [] + group_records = [r for r in group_records if r.is_valid] else: - # If there are any invalid records, the entire group is invalid. - invalid = valid_with_str_members + invalid - errors = errors + [("Invalid JSON in other groupby records", groupby_validator)] * len(valid) - valid = [] - valid_groups.extend(valid) - invalid_groups.extend(invalid) - errors_groups.extend(errors) - continue + _reject_group_records( + group_records, + ("Invalid JSON in other groupby records", groupby_validator), + ) # Handle non-unique group_by values (optionally fix by using first record's values). - if len(set(tuple(record[gb] for gb in self.group_by) for record in valid)) != 1: + valid_parsed = [r.parsed for r in group_records if r.is_valid and r.parsed is not None] + if valid_parsed and len({tuple(record[gb] for gb in self.group_by) for record in valid_parsed}) != 1: if self.config.group_by_fix_non_unique_value: for group_by in self.group_by: - for record in valid[1:]: - record[group_by] = valid[0][group_by] + for record in valid_parsed[1:]: + record[group_by] = valid_parsed[0][group_by] else: - # The group is invalid if the set of group_by fields is not unique. - valid, invalid = [], valid_with_str_members + invalid - errors = [("Groupby value is not unique", groupby_validator)] * len(invalid) - valid_groups.extend(valid) - invalid_groups.extend(invalid) - errors_groups.extend(errors) - continue + _reject_group_records( + group_records, + ("Groupby value is not unique", groupby_validator), + ) # Handle unordered records when order_by is set (optionally fix by sorting). - if self.order_by is not None and not check_if_records_are_ordered(valid, self.order_by): + valid_parsed = [r.parsed for r in group_records if r.is_valid and r.parsed is not None] + if ( + valid_parsed + and self.order_by is not None + and not check_if_records_are_ordered(valid_parsed, self.order_by) + ): if self.config.group_by_fix_unordered_records: - valid.sort(key=lambda x: x[self.order_by]) + # ``valid_parsed`` aliases the underlying ParsedRecord.parsed dicts, + # so sorting a fresh list doesn't reorder the ParsedRecord objects. + # Instead, sort the valid subset of group_records in-place by order_by. + valid_indices = [i for i, r in enumerate(group_records) if r.is_valid] + valid_sorted = sorted( + (group_records[i] for i in valid_indices), + key=lambda r: r.parsed[self.order_by], # type: ignore[index] + ) + for i, record in zip(valid_indices, valid_sorted, strict=True): + group_records[i] = record else: - # If order_by is specified, the group is invalid if the records are not ordered. - valid, invalid = [], valid_with_str_members + invalid - errors = [("Group not ordered", groupby_validator)] * len(invalid) - valid_groups.extend(valid) - invalid_groups.extend(invalid) - errors_groups.extend(errors) - continue - - valid_groups.extend(valid) - invalid_groups.extend(invalid) - errors_groups.extend(errors) + _reject_group_records( + group_records, + ("Group not ordered", groupby_validator), + ) + + all_records.extend(group_records) return ParsedResponse( - valid_records=valid_groups, - invalid_records=invalid_groups, - errors=errors_groups, + records=all_records, + tokenization_time_sec=total_tokenization_time, ) -def create_processor(schema: dict[str, Any], metadata: ModelMetadata, config: SafeSynthesizerParameters) -> Processor: +def _reject_group_records(records: list[ParsedRecord], error: tuple[str, str]) -> None: + """Reclassify every record in *records* as invalid with the given group-level error. + + The group-level error overrides any per-record error that a record + may already carry, matching the legacy behavior where a failing + group invalidates all its records with a single groupby reason. + ``ParsedRecord.text`` and ``ParsedRecord.token_count`` are + preserved, so the original regex-matched JSON string stays + associated with the record (no ``str(dict)`` stringification) and + per-record token counts remain correct after reclassification. + """ + for record in records: + record.parsed = None + record.error = error + + +def create_processor( + schema: dict[str, Any], + metadata: ModelMetadata, + config: SafeSynthesizerParameters, + tokenizer: PreTrainedTokenizerBase | None = None, +) -> Processor: """Create the appropriate record processor for the current pipeline mode. Selects ``TimeSeriesDataProcessor``, ``GroupedDataProcessor``, or @@ -314,6 +376,8 @@ def create_processor(schema: dict[str, Any], metadata: ModelMetadata, config: Sa schema: JSON schema describing the expected record format. metadata: Model metadata (prompt template, BOS/EOS tokens, etc.). config: Pipeline configuration determining the generation mode. + tokenizer: Optional tokenizer for exact token counting during + record parsing. When ``None``, token counts are not tracked. Returns: Processor instance matching the configured generation mode. @@ -325,6 +389,7 @@ def create_processor(schema: dict[str, Any], metadata: ModelMetadata, config: Sa time_column=config.time_series.timestamp_column, interval_seconds=config.time_series.timestamp_interval_seconds, time_format=config.time_series.timestamp_format, + tokenizer=tokenizer, ) elif config.data.group_training_examples_by: processor = GroupedDataProcessor( @@ -334,9 +399,10 @@ def create_processor(schema: dict[str, Any], metadata: ModelMetadata, config: Sa order_by=config.data.order_training_examples_by, bos_token=metadata.prompt_config.bos_token, eos_token=metadata.prompt_config.eos_token, + tokenizer=tokenizer, ) else: - processor = TabularDataProcessor(schema, config=config.generation.validation) + processor = TabularDataProcessor(schema, config=config.generation.validation, tokenizer=tokenizer) logger.info(f"Initialized the {processor.name}") return processor diff --git a/src/nemo_safe_synthesizer/generation/results.py b/src/nemo_safe_synthesizer/generation/results.py index 6cffb9dda..08daccdad 100644 --- a/src/nemo_safe_synthesizer/generation/results.py +++ b/src/nemo_safe_synthesizer/generation/results.py @@ -67,6 +67,30 @@ class GenerateJobResults: elapsed_time: float | None = None """Wall-clock generation duration in seconds, or ``None`` if not yet set.""" + num_completion_tokens: int | None = None + """Total tokens generated by the LLM across all completions.""" + + num_valid_record_tokens: int | None = None + """Tokens in records that passed validation.""" + + num_invalid_record_tokens: int | None = None + """Tokens in records that failed validation.""" + + num_non_record_tokens: int | None = None + """Tokens not part of any recognized record.""" + + tokens_per_prompt: float | None = None + """Average completion tokens per prompt (``num_completion_tokens / num_prompts``).""" + + tokens_per_second: float | None = None + """Total completion tokens divided by generation wall-clock time.""" + + valid_tokens_per_second: float | None = None + """Valid record tokens divided by generation wall-clock time.""" + + tokenization_overhead_sec: float | None = None + """Wall-clock seconds spent tokenizing records for statistics.""" + @classmethod def from_batches( cls, batches: GenerationBatches, max_num_records: int | None, columns: list[str], elapsed_time: float @@ -94,6 +118,8 @@ def from_batches( else 0.0 ) batch_valid_record_fractions = [batch.valid_record_fraction for batch in batches._batches] + + has_token_data = batches.total_completion_tokens > 0 return cls( df=df, status=status, @@ -103,6 +129,20 @@ def from_batches( valid_record_fraction=valid_record_fraction, batch_valid_record_fractions=batch_valid_record_fractions, elapsed_time=elapsed_time, + num_completion_tokens=batches.total_completion_tokens if has_token_data else None, + num_valid_record_tokens=batches.total_valid_record_tokens if has_token_data else None, + num_invalid_record_tokens=batches.total_invalid_record_tokens if has_token_data else None, + num_non_record_tokens=batches.total_non_record_tokens if has_token_data else None, + tokens_per_prompt=( + batches.total_completion_tokens / num_prompts if has_token_data and num_prompts > 0 else None + ), + tokens_per_second=( + batches.total_completion_tokens / elapsed_time if has_token_data and elapsed_time > 0 else None + ), + valid_tokens_per_second=( + batches.total_valid_record_tokens / elapsed_time if has_token_data and elapsed_time > 0 else None + ), + tokenization_overhead_sec=batches.total_tokenization_time_sec if has_token_data else None, ) @@ -163,13 +203,15 @@ def __init__( def _apply_data_actions_fn(self, batch: Batch) -> None: """Post-process and validate a batch via ``data_actions_fn``. - Converts the batch's record lists into a DataFrame, runs the - configured post-processing (reversing training-time preprocessing) - and user-specified validation rules, then maps the results back - onto the batch's ``valid_records`` and ``invalid_records``. - - Each record is tagged with a temporary ID so that validated rows - can be mapped back to their originating response. + Converts the batch's valid records into a DataFrame, runs the + configured post-processing (reversing training-time + preprocessing) and user-specified validation rules, then maps + the results back onto the batch's :class:`ParsedRecord` objects + in place: each valid record either has its ``parsed`` dict + replaced with the transformed row, or is invalidated via + :meth:`ParsedRecord.invalidate` with a data-config error. + Because each ``ParsedRecord`` owns its own ``token_count``, + there is no separate bookkeeping to keep in sync. Args: batch: The batch whose records will be post-processed and @@ -178,16 +220,15 @@ def _apply_data_actions_fn(self, batch: Batch) -> None: if self.data_actions_fn is None: return - # Add a special key to each record, giving it a unique - # identifier. This will make it easier to map original records - # to the validated records, and ensure that we only allow through - # the records that passed validation. + # Tag every valid record with a temporary ID so transformed rows + # can be mapped back to their originating :class:`ParsedRecord`. record_id_key = MetadataColumns.INDEX.value record_id = 0 for response in batch._responses: - for record in response.valid_records: - record[record_id_key] = record_id - record_id += 1 + for record in response.records: + if record.is_valid and record.parsed is not None: + record.parsed[record_id_key] = record_id + record_id += 1 batch_df = batch.to_dataframe() if batch_df is None: @@ -202,28 +243,22 @@ def _apply_data_actions_fn(self, batch: Batch) -> None: # Incrementally add onto the `_batches_df` with each batch of valid records. self._batches_df = pd.concat([self._batches_df, valid_df]) - # map the records to their transformed records id_to_valid_records = dict(zip(valid_df[record_id_key], valid_df.to_dict("records"))) id_to_rejected_records = dict(zip(rejected_df[record_id_key], rejected_df.to_dict("records"))) for response in batch._responses: - new_valid_records = [] - new_rejected_records = [] - for record in response.valid_records: - if new_record := id_to_valid_records.get(record[record_id_key]): + for record in response.records: + if not record.is_valid or record.parsed is None: + continue + rid = record.parsed[record_id_key] + if new_record := id_to_valid_records.get(rid): del new_record[record_id_key] - new_valid_records.append(new_record) - elif new_record := id_to_rejected_records.get(record[record_id_key]): + record.parsed = new_record + elif new_record := id_to_rejected_records.get(rid): del new_record[record_id_key] - new_rejected_records.append(new_record) + record.invalidate(rejected_record_to_error(new_record)) else: raise AssertionError("Every record in response should map to either a valid or rejected row.") - response.valid_records = new_valid_records - for rejected_record in new_rejected_records: - error = rejected_record_to_error(rejected_record) - response.invalid_records.append(str(rejected_record)) - response.errors.append(error) - @property def num_batches(self) -> int: """The number of batches in the generation job.""" @@ -244,6 +279,31 @@ def num_valid_records(self) -> int: """The total number of valid records generated in the generation job.""" return sum([batch.num_valid_records for batch in self._batches]) + @property + def total_completion_tokens(self) -> int: + """Total tokens across all completions in all batches.""" + return sum(batch.total_completion_tokens for batch in self._batches) + + @property + def total_valid_record_tokens(self) -> int: + """Sum of token counts for valid records across all batches.""" + return sum(batch.total_valid_record_tokens for batch in self._batches) + + @property + def total_invalid_record_tokens(self) -> int: + """Sum of token counts for invalid records across all batches.""" + return sum(batch.total_invalid_record_tokens for batch in self._batches) + + @property + def total_non_record_tokens(self) -> int: + """Tokens not part of any recognized record across all batches.""" + return sum(batch.total_non_record_tokens for batch in self._batches) + + @property + def total_tokenization_time_sec(self) -> float: + """Wall-clock seconds spent tokenizing records across all batches.""" + return sum(batch.total_tokenization_time_sec for batch in self._batches) + def add_batch(self, batch: Batch) -> None: """Add a batch and update the generation status. diff --git a/src/nemo_safe_synthesizer/generation/timeseries_backend.py b/src/nemo_safe_synthesizer/generation/timeseries_backend.py index f2c8f9569..2e5aa825a 100644 --- a/src/nemo_safe_synthesizer/generation/timeseries_backend.py +++ b/src/nemo_safe_synthesizer/generation/timeseries_backend.py @@ -17,7 +17,11 @@ from .. import utils from ..config import SafeSynthesizerParameters -from ..data_processing.record_utils import _parse_timestamp_to_seconds, extract_records_from_jsonl_string +from ..data_processing.record_utils import ( + ParsedRecord, + _parse_timestamp_to_seconds, + extract_records_from_jsonl_string, +) from ..defaults import FIXED_RUNTIME_GENERATE_ARGS, LOG_DASHES, PSEUDO_GROUP_COLUMN from ..generation.batch import Batch from ..generation.results import GenerateJobResults, GenerationBatches, GenerationStatus @@ -477,14 +481,15 @@ def _check_chronological_for_group(self, batch: Batch, group_state: GroupState) group_state: Current state of the group (provides the last known timestamp). """ + error = ("Out-of-order time step", "TimeSeries") for response in batch._responses: if not response.valid_records: continue if self._is_chronological_for_group(response.valid_records, group_state): continue - response.invalid_records.extend([json.dumps(record) for record in response.valid_records]) - response.errors.extend([("Out-of-order time step", "TimeSeries")] * len(response.valid_records)) - response.valid_records = [] + for record in response.records: + if record.is_valid: + record.invalidate(error) def _update_group_state(self, group_state: GroupState, records: list[dict]) -> None: """Update a group's state with new records. @@ -767,7 +772,7 @@ def _generate_parallel_groups( group_state = active_states[prompt_idx] batch = group_batches[group_state.group_id] for completion_idx, completion in enumerate(output.outputs): - batch.process(completion_idx, completion.text) + batch.process(completion_idx, completion.text, completion_tokens=len(completion.token_ids)) duration = time.perf_counter() - start_time @@ -839,9 +844,9 @@ def _retain_single_valid_response(self, batch: Batch) -> list[dict]: Returns: List of valid records from the retained response. """ - final_records = [] + final_records: list[dict] = [] - # Find the index of the response with the most valid records + # Find the index of the response with the most valid records. max_valid_idx = None max_valid_count = -1 for idx, response in enumerate(batch._responses): @@ -849,11 +854,14 @@ def _retain_single_valid_response(self, batch: Batch) -> list[dict]: if count > max_valid_count: max_valid_count = count max_valid_idx = idx + + trim_error = ("Extra response trimmed for sliding window", "TimeSeries") for idx, response in enumerate(batch._responses): if idx != max_valid_idx: - response.invalid_records = [] - response.errors.extend([("Extra response trimmed for sliding window", "TimeSeries")]) - response.valid_records = [] + # Drop all records from the trimmed response; replace with a + # single synthetic marker record so the trim is visible in + # error statistics without carrying stale text/token counts. + response.records = [ParsedRecord(text="", error=trim_error)] else: final_records.extend(response.valid_records) diff --git a/src/nemo_safe_synthesizer/generation/vllm_backend.py b/src/nemo_safe_synthesizer/generation/vllm_backend.py index 20388d3cd..df26a1b45 100644 --- a/src/nemo_safe_synthesizer/generation/vllm_backend.py +++ b/src/nemo_safe_synthesizer/generation/vllm_backend.py @@ -9,9 +9,10 @@ import os import time from functools import partial -from typing import TYPE_CHECKING, Any +from typing import Any, cast import torch +from transformers import PreTrainedTokenizerBase from vllm import LLM as vLLM from vllm import RequestOutput from vllm.config import StructuredOutputsConfig @@ -23,9 +24,10 @@ from ..cli.artifact_structure import Workdir from ..config import SafeSynthesizerParameters from ..defaults import DEFAULT_SAMPLING_PARAMETERS, FIXED_RUNTIME_GENERATE_ARGS +from ..errors import InternalError from ..generation.backend import GeneratorBackend from ..generation.batch import Batch -from ..generation.processors import TabularDataProcessor, create_processor +from ..generation.processors import Processor, TabularDataProcessor, create_processor from ..generation.regex_manager import build_json_based_regex from ..generation.results import GenerateJobResults, GenerationBatches, GenerationStatus from ..llm.metadata import ModelMetadata @@ -145,7 +147,11 @@ def __init__( self.use_detailed_logs = kwargs.pop("use_detailed_logs", False) self.gen_method: partial | None = None self._gen_method: partial | None = None - self.processor = create_processor(self.schema, self.model_metadata, self.config) + # Initial processor without a tokenizer; replaced in ``initialize()`` with a + # tokenizer-aware processor once the vLLM engine (and its tokenizer) exists. + # This lets callers introspect the processor type before ``initialize()``, + # at the cost of token counts being zero until the tokenizer is attached. + self.processor: Processor = create_processor(self.schema, self.model_metadata, self.config) adapter_path = self.workdir.adapter_path if self.workdir.adapter_path else self.model_metadata.adapter_path self.lora_req = LoRARequest("lora", 1, str(adapter_path)) if adapter_path else None self._torn_down = False @@ -178,7 +184,12 @@ def __del__(self) -> None: pass def initialize(self, **kwargs) -> None: - """Initialize and load the model into memory.""" + """Initialize and load the model into memory. + + Creates the vLLM engine and then builds the record processor + with the engine's tokenizer so that exact token counts are + available during generation. + """ self._torn_down = False # vLLM 0.12+ accepts attention_config as a constructor arg (replaces the @@ -205,6 +216,16 @@ def initialize(self, **kwargs) -> None: attention_config=attention_config, ) + # vLLM's get_tokenizer() returns a wider union than HF's PreTrainedTokenizerBase; + # in practice it's always a HF tokenizer subclass, so cast for the processor. + tokenizer = cast(PreTrainedTokenizerBase, self.llm.get_tokenizer()) + self.processor = create_processor( + self.schema, + self.model_metadata, + self.config, + tokenizer=tokenizer, + ) + def _build_structured_output_params(self) -> StructuredOutputsParams | None: """Build structured output parameters based on generation config. @@ -337,8 +358,10 @@ def prepare_params(self, **kwargs) -> None: # Create a partially parametrized version of the underlying vllm.LLM.generate # method that is immediately callable downstream. - if TYPE_CHECKING: - assert self.llm is not None + if self.llm is None: + raise InternalError( + "VllmBackend._configure_sampling_params() called before initialize() -- self.llm is None." + ) self._gen_method = partial( self.llm.generate, sampling_params=real_params, @@ -437,7 +460,7 @@ def _generate_batch( f"finish_reason={out.finish_reason}, " f"stop_reason={out.stop_reason}" ) - batch.process(idx, out.text) + batch.process(idx, out.text, completion_tokens=len(out.token_ids)) return batch @@ -458,13 +481,16 @@ def _log_batch_timing_and_progress( records_per_second = 0 if duration == 0 else batch.num_valid_records / duration # Build structured data - processor renders as table for console - progress_data = { + progress_data: dict[str, int | float] = { "records_per_second": round(records_per_second, 2), "duration_seconds": round(duration, 2), "valid_records_generated": batches.num_valid_records, "target_records": self.config.generation.num_records, "progress_fraction": round(batches.num_valid_records / self.config.generation.num_records, 4), } + if batch.total_completion_tokens > 0 and duration > 0: + progress_data["tokens_per_second"] = round(batch.total_completion_tokens / duration, 1) + progress_data["valid_tokens_per_second"] = round(batch.total_valid_record_tokens / duration, 1) # Pass structured data - processor renders for console, JSON keeps as-is logger.user.info( diff --git a/src/nemo_safe_synthesizer/observability.py b/src/nemo_safe_synthesizer/observability.py index c85203dad..e07d091fe 100644 --- a/src/nemo_safe_synthesizer/observability.py +++ b/src/nemo_safe_synthesizer/observability.py @@ -220,14 +220,11 @@ def _render_rich_table(data: dict, title: str | None = None) -> str: table.add_column("Value") for key, value in data.items(): display_key = key.replace("_", " ").title() - # TODO: Refactor this formatting logic to be more generic and maintainable. - # Currently requires updating this file whenever new metrics are added that - # need special formatting (e.g., the "loss", "eval_loss" exclusion list). if isinstance(value, float): - if key not in ("loss", "eval_loss") and value < 1 and value > 0: - display_value = f"{value:.2%}" - else: - display_value = f"{value:.2f}" + is_fraction = ( + 0 < value < 1 and key not in ("loss", "eval_loss") and not key.endswith(("_sec", "_seconds")) + ) + display_value = f"{value:.2%}" if is_fraction else f"{value:.2f}" else: display_value = str(value) table.add_row(display_key, display_value) diff --git a/src/nemo_safe_synthesizer/results.py b/src/nemo_safe_synthesizer/results.py index eeae055bb..c53176c78 100644 --- a/src/nemo_safe_synthesizer/results.py +++ b/src/nemo_safe_synthesizer/results.py @@ -64,11 +64,31 @@ def make_nss_summary( num_invalid_records = None num_prompts = None valid_record_fraction = None + num_completion_tokens = None + num_valid_record_tokens = None + num_invalid_record_tokens = None + num_non_record_tokens = None + valid_record_token_fraction = None + tokens_per_prompt = None + tokens_per_second = None + valid_tokens_per_second = None + tokenization_overhead_sec = None + if isinstance(results, GenerateJobResults): num_valid_records = results.num_valid_records num_invalid_records = results.num_invalid_records num_prompts = results.num_prompts valid_record_fraction = results.valid_record_fraction + num_completion_tokens = results.num_completion_tokens + num_valid_record_tokens = results.num_valid_record_tokens + num_invalid_record_tokens = results.num_invalid_record_tokens + num_non_record_tokens = results.num_non_record_tokens + tokens_per_prompt = results.tokens_per_prompt + tokens_per_second = results.tokens_per_second + valid_tokens_per_second = results.valid_tokens_per_second + tokenization_overhead_sec = results.tokenization_overhead_sec + if num_completion_tokens is not None and num_completion_tokens > 0 and num_valid_record_tokens is not None: + valid_record_token_fraction = num_valid_record_tokens / num_completion_tokens return SafeSynthesizerSummary( timing=timing, @@ -76,6 +96,15 @@ def make_nss_summary( num_invalid_records=num_invalid_records, num_prompts=num_prompts, valid_record_fraction=valid_record_fraction, + num_completion_tokens=num_completion_tokens, + num_valid_record_tokens=num_valid_record_tokens, + num_invalid_record_tokens=num_invalid_record_tokens, + num_non_record_tokens=num_non_record_tokens, + valid_record_token_fraction=valid_record_token_fraction, + tokens_per_prompt=tokens_per_prompt, + tokens_per_second=tokens_per_second, + valid_tokens_per_second=valid_tokens_per_second, + tokenization_overhead_sec=tokenization_overhead_sec, synthetic_data_quality_score=synthetic_data_quality_score, column_correlation_stability_score=column_correlation_stability_score, deep_structure_stability_score=deep_structure_stability_score, diff --git a/src/nemo_safe_synthesizer/training/callbacks.py b/src/nemo_safe_synthesizer/training/callbacks.py index a190bba28..bb122845c 100644 --- a/src/nemo_safe_synthesizer/training/callbacks.py +++ b/src/nemo_safe_synthesizer/training/callbacks.py @@ -150,10 +150,24 @@ def on_evaluate( skip_special_tokens=self.is_tabular_processor, ) + # Per-row completion token counts feed Batch's token + # statistics alongside vLLM-backed generation. When the + # tokenizer has a pad_token_id, trailing pads are + # excluded; otherwise we fall back to the full + # completion width (slight over-count when generation + # ends early, acceptable for an eval metric). + prompt_len = input_ids.shape[1] + completion_ids = outputs[:, prompt_len:] + pad_id = tokenizer.pad_token_id + if pad_id is not None: + completion_tokens_per_row = (completion_ids != pad_id).sum(dim=1).tolist() + else: + completion_tokens_per_row = [completion_ids.shape[1]] * outputs.shape[0] + start_time = time.perf_counter() batch = Batch(processor=self.processor) - for idx, text in enumerate(decoded): - batch.process(idx, text) + for idx, (text, n_tokens) in enumerate(zip(decoded, completion_tokens_per_row, strict=True)): + batch.process(idx, text, completion_tokens=n_tokens) duration = time.perf_counter() - start_time self.generation.add_batch(batch) diff --git a/src/nemo_safe_synthesizer/training/huggingface_backend.py b/src/nemo_safe_synthesizer/training/huggingface_backend.py index 806ded9cf..701107aeb 100644 --- a/src/nemo_safe_synthesizer/training/huggingface_backend.py +++ b/src/nemo_safe_synthesizer/training/huggingface_backend.py @@ -509,6 +509,7 @@ def _add_inference_eval_callback(self, trainer: Trainer, training_args: dict) -> config=self.params, schema=self.dataset_schema, metadata=self.model_metadata, + tokenizer=self.tokenizer, ), num_prompts_per_batch=DEFAULT_VALID_RECORD_EVAL_BATCH_SIZE, **training_args["inference_eval_kwargs"], diff --git a/tests/conftest.py b/tests/conftest.py index 3ec0ff099..6a969ee70 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -269,13 +269,17 @@ def fixture_mock_processor(): ] from unittest.mock import MagicMock - from nemo_safe_synthesizer.generation.processors import ParsedResponse + from nemo_safe_synthesizer.generation.processors import ParsedRecord, ParsedResponse mock_processor = MagicMock() mock_processor.return_value = ParsedResponse( - valid_records=stub_valid_records, - invalid_records=["invalidjson"], - errors=[("some error msg", "some error msg")], + records=[ + *(ParsedRecord(text=str(r), parsed=r) for r in stub_valid_records), + ParsedRecord( + text="invalidjson", + error=("some error msg", "some error msg"), + ), + ], prompt_number=1, ) return mock_processor @@ -286,13 +290,13 @@ def fixture_mock_processor_without_valid_records(): """Mock processor returning a ``ParsedResponse`` with zero valid records.""" from unittest.mock import MagicMock - from nemo_safe_synthesizer.generation.processors import ParsedResponse + from nemo_safe_synthesizer.generation.processors import ParsedRecord, ParsedResponse mock_processor = MagicMock() mock_processor.return_value = ParsedResponse( - valid_records=[], - invalid_records=["invalidjson"], - errors=[("some error msg", "some error msg")], + records=[ + ParsedRecord(text="invalidjson", error=("some error msg", "some error msg")), + ], prompt_number=1, ) return mock_processor diff --git a/tests/data_processing/test_records.py b/tests/data_processing/test_records.py index e10998d12..d1670a36f 100644 --- a/tests/data_processing/test_records.py +++ b/tests/data_processing/test_records.py @@ -20,6 +20,11 @@ ) +def _mock_encode(text: str) -> list[int]: + """Deterministic mock tokenizer: one token per character.""" + return list(range(len(text))) + + def test_is_safe_for_float_conversion(): # Test with safe values assert is_safe_for_float_conversion(100) is True @@ -101,21 +106,31 @@ def test_extract_and_validate_records_with_invalid_records_with_large_numbers( jsonl_schema["properties"]["sepal.length"]["maximum"] = 10**309 corrupt_record_str = json.dumps(corrupt_record) valid_jsonl_str += f"""{corrupt_record_str}""" - valid, invalid, errors = extract_and_validate_records(valid_jsonl_str, schema=jsonl_schema) - assert len(valid) == 5 - assert len(invalid) == 1 - assert len(errors) == 1 - assert len(errors[0]) == 2 - assert "is too large to convert to float64" in errors[0][0] - assert errors[0][1] == "Float Conversion" + result = extract_and_validate_records(valid_jsonl_str, schema=jsonl_schema) + assert len(result.valid_records) == 5 + assert len(result.invalid_records) == 1 + assert len(result.errors) == 1 + assert len(result.errors[0]) == 2 + assert "is too large to convert to float64" in result.errors[0][0] + assert result.errors[0][1] == "Float Conversion" def test_extract_and_validate_records(fixture_valid_iris_dataset_jsonl_and_schema): valid_jsonl_str, jsonl_schema = fixture_valid_iris_dataset_jsonl_and_schema - valid, invalid, errors = extract_and_validate_records(valid_jsonl_str, schema=jsonl_schema) - assert len(valid) == 5 - assert len(invalid) == 0 - assert len(errors) == 0 + result = extract_and_validate_records(valid_jsonl_str, schema=jsonl_schema) + assert len(result.valid_records) == 5 + assert len(result.invalid_records) == 0 + assert len(result.errors) == 0 + + # Without a tokenizer, per-record counts default to 0 and no tokenization time is spent. + assert all(r.token_count == 0 for r in result.records) + assert result.tokenization_time_sec == 0.0 + + # With a tokenizer, each record gets a positive count and tokenization time is recorded. + result_with_tokens = extract_and_validate_records(valid_jsonl_str, schema=jsonl_schema, encode=_mock_encode) + assert len(result_with_tokens.records) == 5 + assert all(r.is_valid and r.token_count > 0 for r in result_with_tokens.records) + assert result_with_tokens.tokenization_time_sec >= 0 def test_extract_and_validate_records_with_invalid_records( @@ -125,13 +140,13 @@ def test_extract_and_validate_records_with_invalid_records( # intentionally corrupt the last record with invalid text that would lead to a json decode error valid_jsonl_str = valid_jsonl_str[:-2] + "invalid entry }" - valid, invalid, errors = extract_and_validate_records(valid_jsonl_str, schema=jsonl_schema) - assert len(valid) == 4 - assert len(invalid) == 1 - assert len(errors) == 1 - assert len(errors[0]) == 2 - assert errors[0][0] == "Invalid JSON: Expecting ',' delimiter" - assert errors[0][1] == "Invalid JSON" + result = extract_and_validate_records(valid_jsonl_str, schema=jsonl_schema) + assert len(result.valid_records) == 4 + assert len(result.invalid_records) == 1 + assert len(result.errors) == 1 + assert len(result.errors[0]) == 2 + assert result.errors[0][0] == "Invalid JSON: Expecting ',' delimiter" + assert result.errors[0][1] == "Invalid JSON" def test_extract_and_validate_records_with_invalid_schema( @@ -140,14 +155,19 @@ def test_extract_and_validate_records_with_invalid_schema( valid_jsonl_str, jsonl_schema = fixture_valid_iris_dataset_jsonl_and_schema # tack on a new json obj with a different schema than in the spec - valid_jsonl_str = valid_jsonl_str + f"{json.dumps(dict(a=1, b=2))}" - valid, invalid, errors = extract_and_validate_records(valid_jsonl_str, schema=jsonl_schema) - assert len(valid) == 5 - assert len(invalid) == 1 - assert len(errors) == 1 - assert len(errors[0]) == 2 - assert errors[0][0] == "'sepal.length' is a required property" - assert errors[0][1] == "required" + invalid_text = json.dumps(dict(a=1, b=2)) + result = extract_and_validate_records(valid_jsonl_str + invalid_text, schema=jsonl_schema, encode=_mock_encode) + assert len(result.valid_records) == 5 + assert len(result.invalid_records) == 1 + assert len(result.errors) == 1 + assert len(result.errors[0]) == 2 + assert result.errors[0][0] == "'sepal.length' is a required property" + assert result.errors[0][1] == "required" + + # Invalid records keep their original regex-matched text and still get a token count. + assert not result.records[-1].is_valid + assert result.records[-1].text == invalid_text + assert result.records[-1].token_count > 0 def test_extract_and_validate_records_with_invalid_utf_characters_embedded( @@ -163,25 +183,28 @@ def test_extract_and_validate_records_with_invalid_utf_characters_embedded( valid_jsonl_str = valid_jsonl_str.replace(existing_petal_variety, new_petal_variety_with_invalid_utf8) jsonl_schema["properties"]["variety"]["enum"].append(new_petal_variety_with_invalid_utf8) - valid, invalid, errors = extract_and_validate_records(valid_jsonl_str, schema=jsonl_schema) - assert len(valid) == 5 - assert len(invalid) == 0 - assert len(errors) == 0 + result = extract_and_validate_records(valid_jsonl_str, schema=jsonl_schema) + assert len(result.valid_records) == 5 + assert len(result.invalid_records) == 0 + assert len(result.errors) == 0 def test_extract_and_validate_records_with_non_english_characters(fixture_lmsys_dataset_jsonl_and_schema): valid_jsonl_str, jsonl_schema = fixture_lmsys_dataset_jsonl_and_schema - valid, invalid, errors = extract_and_validate_records(valid_jsonl_str, schema=jsonl_schema) - assert len(valid) == 5 - assert len(invalid) == 0 - assert len(errors) == 0 + result = extract_and_validate_records(valid_jsonl_str, schema=jsonl_schema) + assert len(result.valid_records) == 5 + assert len(result.invalid_records) == 0 + assert len(result.errors) == 0 # Check that the non-English content is preserved correctly - assert valid[0]["conversation"][0]["content"] == "ПРИВЕТ" - assert valid[1]["conversation"][0]["content"] == "一家4000人的化工企业需要配备几名安全员" - assert valid[2]["conversation"][0]["content"][:10] == "경찰 관계자는 “범" - assert valid[3]["conversation"][0]["content"] == "نعم انا خبير في ادارة تطبيقات وحسابات التواصل الاجتماعي بحترافيه" - assert valid[4]["conversation"][0]["content"][41:74] == "\"Ім'я і національність\n@Ukrostap\n" + assert result.valid_records[0]["conversation"][0]["content"] == "ПРИВЕТ" + assert result.valid_records[1]["conversation"][0]["content"] == "一家4000人的化工企业需要配备几名安全员" + assert result.valid_records[2]["conversation"][0]["content"][:10] == "경찰 관계자는 “범" + assert ( + result.valid_records[3]["conversation"][0]["content"] + == "نعم انا خبير في ادارة تطبيقات وحسابات التواصل الاجتماعي بحترافيه" + ) + assert result.valid_records[4]["conversation"][0]["content"][41:74] == "\"Ім'я і національність\n@Ukrostap\n" def _run_csv_writer(df: pd.DataFrame): @@ -389,13 +412,19 @@ def test_extract_and_validate_timeseries_records_valid(): '{"timestamp": "2024-01-01 02:00:00", "value": 3}\n' ) - valid, invalid, errors = extract_and_validate_timeseries_records( - jsonl, schema, time_column="timestamp", interval_seconds=3600, time_format="%Y-%m-%d %H:%M:%S" + result = extract_and_validate_timeseries_records( + jsonl, + schema, + time_column="timestamp", + interval_seconds=3600, + time_format="%Y-%m-%d %H:%M:%S", + encode=_mock_encode, ) - assert len(valid) == 3 - assert len(invalid) == 0 - assert len(errors) == 0 + assert len(result.valid_records) == 3 + assert len(result.invalid_records) == 0 + assert len(result.errors) == 0 + assert all(r.is_valid and r.token_count > 0 for r in result.records) def test_extract_and_validate_timeseries_records_no_interval_validation(): @@ -412,13 +441,13 @@ def test_extract_and_validate_timeseries_records_no_interval_validation(): '{"timestamp": "2024-01-01 05:00:00", "value": 3}\n' ) - valid, invalid, errors = extract_and_validate_timeseries_records( + result = extract_and_validate_timeseries_records( jsonl, schema, time_column="timestamp", interval_seconds=None, time_format="%Y-%m-%d %H:%M:%S" ) - assert len(valid) == 3 - assert len(invalid) == 0 - assert len(errors) == 0 + assert len(result.valid_records) == 3 + assert len(result.invalid_records) == 0 + assert len(result.errors) == 0 def test_extract_and_validate_timeseries_records_invalid_interval(): @@ -435,13 +464,22 @@ def test_extract_and_validate_timeseries_records_invalid_interval(): '{"timestamp": "2024-01-01 03:00:00", "value": 3}\n' ) - valid, invalid, errors = extract_and_validate_timeseries_records( - jsonl, schema, time_column="timestamp", interval_seconds=3600, time_format="%Y-%m-%d %H:%M:%S" + result = extract_and_validate_timeseries_records( + jsonl, + schema, + time_column="timestamp", + interval_seconds=3600, + time_format="%Y-%m-%d %H:%M:%S", + encode=_mock_encode, ) - assert len(valid) == 1 # Only first record is valid - assert len(invalid) == 2 # Second and third are invalid - assert len(errors) == 2 + assert len(result.valid_records) == 1 # Only first record is valid + assert len(result.invalid_records) == 2 # Second and third are invalid + assert len(result.errors) == 2 + + # Cascade-invalidated records still keep their original text and token count. + assert all(r.token_count > 0 for r in result.records) + assert result.records[2].text.startswith('{"timestamp":') def test_extract_and_validate_timeseries_records_invalid_json(): @@ -456,14 +494,14 @@ def test_extract_and_validate_timeseries_records_invalid_json(): '{"timestamp": "2024-01-01 02:00:00", "value": 3}\n' ) - valid, invalid, errors = extract_and_validate_timeseries_records( + result = extract_and_validate_timeseries_records( jsonl, schema, time_column="timestamp", interval_seconds=3600, time_format="%Y-%m-%d %H:%M:%S" ) - assert len(valid) == 1 - assert len(invalid) == 1 - assert len(errors) == 1 - assert "Invalid JSON" in errors[0][0] + assert len(result.valid_records) == 1 + assert len(result.invalid_records) == 1 + assert len(result.errors) == 1 + assert "Invalid JSON" in result.errors[0][0] def test_extract_and_validate_timeseries_records_elapsed_time(): @@ -479,13 +517,13 @@ def test_extract_and_validate_timeseries_records_elapsed_time(): '{"elapsed_seconds": 7200, "value": 3}\n' ) - valid, invalid, errors = extract_and_validate_timeseries_records( + result = extract_and_validate_timeseries_records( jsonl, schema, time_column="elapsed_seconds", interval_seconds=3600, time_format="elapsed_seconds" ) - assert len(valid) == 3 - assert len(invalid) == 0 - assert len(errors) == 0 + assert len(result.valid_records) == 3 + assert len(result.invalid_records) == 0 + assert len(result.errors) == 0 def test_extract_and_validate_timeseries_records_missing_timestamp(): @@ -497,11 +535,11 @@ def test_extract_and_validate_timeseries_records_missing_timestamp(): } jsonl = '{"value": 1}\n{"value": 2}\n' - valid, invalid, errors = extract_and_validate_timeseries_records( + result = extract_and_validate_timeseries_records( jsonl, schema, time_column="timestamp", interval_seconds=3600, time_format="%Y-%m-%d %H:%M:%S" ) - assert len(valid) == 0 - assert len(invalid) == 1 - assert len(errors) == 1 - assert "Missing 'timestamp'" in errors[0][0] + assert len(result.valid_records) == 0 + assert len(result.invalid_records) == 1 + assert len(result.errors) == 1 + assert "Missing 'timestamp'" in result.errors[0][0] diff --git a/tests/generation/test_batch.py b/tests/generation/test_batch.py index 07107b1e4..1a774d9d7 100644 --- a/tests/generation/test_batch.py +++ b/tests/generation/test_batch.py @@ -13,7 +13,7 @@ import nemo_safe_synthesizer.generation.batch from nemo_safe_synthesizer.data_processing.actions.utils import MetadataColumns from nemo_safe_synthesizer.generation.batch import Batch -from nemo_safe_synthesizer.generation.processors import ParsedResponse +from nemo_safe_synthesizer.generation.processors import ParsedRecord, ParsedResponse from nemo_safe_synthesizer.generation.results import rejected_record_to_error @@ -26,37 +26,39 @@ def fixture_mock_processor_errors(): dict(some="value0", other=1), dict(some="value1", other=2), ] + errors_list = [ + ("err2", "const"), + ("err3", "Invalid JSON"), + ("err4", "groupby"), + ("err4", "groupby"), + ("err4", "groupby"), + ("err4", "groupby"), + ("err4", "groupby"), + ("'col1' is a required property", "required"), + ("'col2' is a required property", "required"), + ("'col3' is a required property", "required"), + ( + "0 is not one of [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32 33, 34, 35, 36, 37, 38, 39, 40]", + "enum", + ), + ( + "0 is not one of [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32 33, 34, 35, 36, 37, 38, 39, 40]", + "enum", + ), + ("0 is not one of [4, 5, 6]", "enum"), + ("1 is not one of [4, 5, 6]", "enum"), + ("2 is not one of [4, 5, 6]", "enum"), + ("3 is not one of [4, 5, 6]", "enum"), + ("7 is not one of [4, 5, 6]", "enum"), + ("8 is not one of [4, 5, 6]", "enum"), + ("9 is not one of [4, 5, 6]", "enum"), + ("10 is not one of [4, 5, 6]", "enum"), + ] mock_processor = MagicMock() mock_processor.return_value = ParsedResponse( - valid_records=stub_valid_records, - invalid_records=["invalidjson"] * 20, - errors=[ - ("err2", "const"), - ("err3", "Invalid JSON"), - ("err4", "groupby"), - ("err4", "groupby"), - ("err4", "groupby"), - ("err4", "groupby"), - ("err4", "groupby"), - ("'col1' is a required property", "required"), - ("'col2' is a required property", "required"), - ("'col3' is a required property", "required"), - ( - "0 is not one of [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32 33, 34, 35, 36, 37, 38, 39, 40]", - "enum", - ), - ( - "0 is not one of [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32 33, 34, 35, 36, 37, 38, 39, 40]", - "enum", - ), - ("0 is not one of [4, 5, 6]", "enum"), - ("1 is not one of [4, 5, 6]", "enum"), - ("2 is not one of [4, 5, 6]", "enum"), - ("3 is not one of [4, 5, 6]", "enum"), - ("7 is not one of [4, 5, 6]", "enum"), - ("8 is not one of [4, 5, 6]", "enum"), - ("9 is not one of [4, 5, 6]", "enum"), - ("10 is not one of [4, 5, 6]", "enum"), + records=[ + *(ParsedRecord(text=str(r), parsed=r) for r in stub_valid_records), + *(ParsedRecord(text="invalidjson", error=err) for err in errors_list), ], prompt_number=1, ) @@ -75,11 +77,9 @@ def fixture_mock_processor_rejected_records(fixture_mock_processor_errors: Magic {"a": "a4", reject_reason: "reason2"}, {"a": "a5", reject_reason: "reason2"}, ] - errors = [rejected_record_to_error(r) for r in rejected_records] - invalid_records = [str(r) for r in rejected_records] - - fixture_mock_processor_errors.return_value.invalid_records.extend(invalid_records) - fixture_mock_processor_errors.return_value.errors.extend(errors) + fixture_mock_processor_errors.return_value.records.extend( + ParsedRecord(text=str(r), error=rejected_record_to_error(r)) for r in rejected_records + ) return fixture_mock_processor_errors @@ -220,3 +220,76 @@ def test_log_summary_data_config(caplog, fixture_mock_processor_rejected_records assert error_ctx is not None error_data = error_ctx["tabular_data"] assert any("Failed data_config validation due to [reason1]" in key for key in error_data) + + +# --------------------------------------------------------------------------- +# Token-counting tests +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def fixture_mock_processor_with_tokens(): + """Mock processor returning a ``ParsedResponse`` with per-record token counts.""" + mock_processor = MagicMock() + mock_processor.return_value = ParsedResponse( + records=[ + ParsedRecord(text='{"a": 1}', parsed={"a": 1}, token_count=10), + ParsedRecord(text='{"a": 2}', parsed={"a": 2}, token_count=20), + ParsedRecord(text='{"a": 3}', parsed={"a": 3}, token_count=30), + ParsedRecord(text="bad1", error=("err", "err"), token_count=15), + ], + prompt_number=1, + tokenization_time_sec=0.01, + ) + return mock_processor + + +def test_batch_completion_tokens(fixture_mock_processor_with_tokens): + batch = Batch(processor=fixture_mock_processor_with_tokens) + batch.process(prompt_number=1, text="stub", completion_tokens=100) + batch.process(prompt_number=2, text="stub", completion_tokens=150) + assert batch.total_completion_tokens == 250 + + +def test_batch_aggregate_token_properties(fixture_mock_processor_with_tokens): + batch = Batch(processor=fixture_mock_processor_with_tokens) + batch.process(prompt_number=1, text="stub", completion_tokens=200) + + assert batch.total_valid_record_tokens == 60 # 10+20+30 + assert batch.total_invalid_record_tokens == 15 + assert batch.total_non_record_tokens == 200 - 60 - 15 # 125 + assert batch.total_tokenization_time_sec == pytest.approx(0.01) + + +def test_batch_non_record_tokens_clamped_to_zero(): + """If completion_tokens < valid + invalid, clamp to 0.""" + mock_processor = MagicMock() + mock_processor.return_value = ParsedResponse( + records=[ParsedRecord(text='{"a": 1}', parsed={"a": 1}, token_count=100)], + tokenization_time_sec=0.0, + ) + batch = Batch(processor=mock_processor) + batch.process(prompt_number=1, text="stub", completion_tokens=50) + assert batch.total_non_record_tokens == 0 + + +def test_batch_no_completion_tokens_defaults_zero(fixture_mock_processor): + """Default completion_tokens=0 when not provided.""" + batch = Batch(processor=fixture_mock_processor) + batch.process(prompt_number=1, text="stub") + assert batch.total_completion_tokens == 0 + + +def test_batch_log_summary_includes_token_data(caplog, fixture_mock_processor_with_tokens): + caplog.set_level(INFO) + batch = Batch(processor=fixture_mock_processor_with_tokens) + batch.process(prompt_number=1, text="stub", completion_tokens=200) + batch.log_summary() + + ctx_data = [getattr(record, "ctx", None) for record in caplog.records if hasattr(record, "ctx")] + summary_ctx = ctx_data[0] + assert summary_ctx is not None + assert summary_ctx["tabular_data"]["completion_tokens"] == 200 + assert summary_ctx["tabular_data"]["valid_record_tokens"] == 60 + assert summary_ctx["tabular_data"]["invalid_record_tokens"] == 15 + assert summary_ctx["tabular_data"]["non_record_tokens"] == 125 diff --git a/tests/generation/test_generation.py b/tests/generation/test_generation.py index edd23b265..246b711a4 100644 --- a/tests/generation/test_generation.py +++ b/tests/generation/test_generation.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import logging +from unittest.mock import MagicMock import pandas as pd import pytest @@ -14,8 +15,13 @@ ) from nemo_safe_synthesizer.errors import GenerationError from nemo_safe_synthesizer.generation.batch import Batch -from nemo_safe_synthesizer.generation.processors import ParsedResponse -from nemo_safe_synthesizer.generation.results import NUM_PROMPT_BUFFER, GenerationBatches, GenerationStatus +from nemo_safe_synthesizer.generation.processors import ParsedRecord, ParsedResponse +from nemo_safe_synthesizer.generation.results import ( + NUM_PROMPT_BUFFER, + GenerateJobResults, + GenerationBatches, + GenerationStatus, +) # Purpose: Builds reusable good/bad Batch sets for generation tests. @@ -224,15 +230,11 @@ def test_apply_data_actions(fixture_mock_processor, caplog): batch = Batch(fixture_mock_processor) batch._responses = [ ParsedResponse( - valid_records=data.iloc[:3].to_dict("records"), - invalid_records=[], - errors=[], + records=[ParsedRecord(text=str(r), parsed=r) for r in data.iloc[:3].to_dict("records")], prompt_number=1, ), ParsedResponse( - valid_records=data.iloc[3:].to_dict("records"), - invalid_records=[], - errors=[], + records=[ParsedRecord(text=str(r), parsed=r) for r in data.iloc[3:].to_dict("records")], prompt_number=2, ), ] @@ -289,3 +291,91 @@ def test_log_status_raises_generation_error_on_no_records(fixture_stub_batches): assert generation.status == GenerationStatus.STOP_NO_RECORDS with pytest.raises(GenerationError): generation.log_status() + + +# --------------------------------------------------------------------------- +# Token-aggregation tests +# --------------------------------------------------------------------------- + + +def _make_batch_with_tokens( + valid_counts: list[int], + invalid_counts: list[int], + completion_tokens: int, + tok_time: float = 0.01, +) -> Batch: + """Build a Batch with a mock processor that returns specific token counts.""" + records = [ + ParsedRecord(text=f'{{"a": {i}}}', parsed={"a": i}, token_count=tc) for i, tc in enumerate(valid_counts) + ] + [ParsedRecord(text=f"bad{i}", error=("err", "err"), token_count=tc) for i, tc in enumerate(invalid_counts)] + mock_proc = MagicMock() + mock_proc.return_value = ParsedResponse( + records=records, + tokenization_time_sec=tok_time, + ) + batch = Batch(processor=mock_proc) + batch.process(0, "stub", completion_tokens=completion_tokens) + return batch + + +class TestGenerationBatchesTokenAggregation: + """GenerationBatches aggregates token properties across batches.""" + + def test_total_completion_tokens(self): + b1 = _make_batch_with_tokens([10, 20], [5], completion_tokens=100) + b2 = _make_batch_with_tokens([30], [10, 15], completion_tokens=200) + gen = GenerationBatches(batches=[b1, b2]) + assert gen.total_completion_tokens == 300 + + def test_total_valid_record_tokens(self): + b1 = _make_batch_with_tokens([10, 20], [5], completion_tokens=100) + b2 = _make_batch_with_tokens([30], [10], completion_tokens=200) + gen = GenerationBatches(batches=[b1, b2]) + assert gen.total_valid_record_tokens == 60 # 10+20+30 + + def test_total_invalid_record_tokens(self): + b1 = _make_batch_with_tokens([10], [5, 15], completion_tokens=100) + gen = GenerationBatches(batches=[b1]) + assert gen.total_invalid_record_tokens == 20 + + def test_total_non_record_tokens(self): + b1 = _make_batch_with_tokens([10], [5], completion_tokens=100) + gen = GenerationBatches(batches=[b1]) + assert gen.total_non_record_tokens == 85 # 100 - 10 - 5 + + def test_total_tokenization_time_sec(self): + b1 = _make_batch_with_tokens([10], [5], completion_tokens=100, tok_time=0.01) + b2 = _make_batch_with_tokens([20], [10], completion_tokens=200, tok_time=0.02) + gen = GenerationBatches(batches=[b1, b2]) + assert gen.total_tokenization_time_sec == pytest.approx(0.03) + + +class TestGenerateJobResultsFromBatches: + """GenerateJobResults.from_batches populates token fields correctly.""" + + def test_with_token_data(self): + b = _make_batch_with_tokens([10, 20], [5], completion_tokens=100, tok_time=0.05) + gen = GenerationBatches(batches=[b]) + gen.job_complete() + results = GenerateJobResults.from_batches(gen, max_num_records=None, columns=["a"], elapsed_time=2.0) + + assert results.num_completion_tokens == 100 + assert results.num_valid_record_tokens == 30 + assert results.num_invalid_record_tokens == 5 + assert results.num_non_record_tokens == 65 + assert results.tokens_per_prompt == pytest.approx(100.0) # 100/1 prompt + assert results.tokens_per_second == pytest.approx(50.0) # 100/2.0 + assert results.valid_tokens_per_second == pytest.approx(15.0) # 30/2.0 + assert results.tokenization_overhead_sec == pytest.approx(0.05) + + def test_without_token_data(self): + """When completion_tokens=0, all token fields stay None.""" + b = _make_batch_with_tokens([], [], completion_tokens=0) + gen = GenerationBatches(batches=[b]) + gen.job_complete() + results = GenerateJobResults.from_batches(gen, max_num_records=None, columns=["a"], elapsed_time=2.0) + + assert results.num_completion_tokens is None + assert results.tokens_per_second is None + assert results.valid_tokens_per_second is None + assert results.tokenization_overhead_sec is None diff --git a/tests/generation/test_processors.py b/tests/generation/test_processors.py index b7e2de3d9..6c78c65f5 100644 --- a/tests/generation/test_processors.py +++ b/tests/generation/test_processors.py @@ -3,6 +3,7 @@ import copy from pathlib import Path +from unittest.mock import MagicMock import pandas as pd import pytest @@ -564,3 +565,96 @@ def test_assembler_to_processor_non_english( tokenizer=tokenizer, cache_file_path=fixture_session_cache_dir, ) + + +# --------------------------------------------------------------------------- +# Token-counting tests for processors +# --------------------------------------------------------------------------- + + +def _make_mock_tokenizer(chars_per_token: int = 1) -> MagicMock: + """Create a mock tokenizer whose encode returns one token per `chars_per_token` chars.""" + tokenizer = MagicMock() + + def _encode(text: str, add_special_tokens: bool = True) -> list[int]: + n = max(1, len(text) // chars_per_token) + return list(range(n)) + + tokenizer.encode = _encode + return tokenizer + + +class TestTabularProcessorTokenCounts: + def test_populates_token_counts(self, fixture_valid_iris_dataset_jsonl_and_schema, fixture_validation_config): + jsonl_str, jsonl_schema = fixture_valid_iris_dataset_jsonl_and_schema + tokenizer = _make_mock_tokenizer() + proc = TabularDataProcessor(schema=jsonl_schema, config=fixture_validation_config, tokenizer=tokenizer) + response = proc(1, jsonl_str) + + assert all(r.is_valid and r.token_count > 0 for r in response.records) + assert response.tokenization_time_sec >= 0 + + def test_no_tokenizer_gives_zero_counts( + self, fixture_valid_iris_dataset_jsonl_and_schema, fixture_validation_config + ): + jsonl_str, jsonl_schema = fixture_valid_iris_dataset_jsonl_and_schema + proc = TabularDataProcessor(schema=jsonl_schema, config=fixture_validation_config, tokenizer=None) + response = proc(1, jsonl_str) + + assert all(r.is_valid for r in response.records) + assert all(r.token_count == 0 for r in response.records) + assert response.tokenization_time_sec == 0.0 + + +class TestGroupedProcessorTokenShift: + """Token counts stay attached to each record through reclassification.""" + + def test_non_unique_group_preserves_text_and_tokens( + self, + fixture_valid_iris_dataset_jsonl_and_schema, + fixture_validation_config, + ): + jsonl_str, jsonl_schema = fixture_valid_iris_dataset_jsonl_and_schema + + jsonl_schema_copy = copy.deepcopy(jsonl_schema) + jsonl_schema_copy["properties"]["variety"]["enum"].append("NewSetosa") + groups_jsonl_str = ( + BOS + + jsonl_str + + '{"sepal.length":5.1,"sepal.width":3.5,"petal.length":1.4,"petal.width":0.2,"variety":"NewSetosa"}\n' + + EOS + ) + tokenizer = _make_mock_tokenizer() + proc = GroupedDataProcessor( + schema=jsonl_schema_copy, + config=fixture_validation_config, + group_by="variety", + bos_token=BOS, + eos_token=EOS, + tokenizer=tokenizer, + ) + response = proc(1, groups_jsonl_str) + + assert len(response.records) == 6 + assert all(not r.is_valid for r in response.records) + assert all(r.token_count > 0 for r in response.records) + # After reclassification ``text`` is the original JSON string, not + # str(dict), so users can recover the original completion slice. + assert all(r.text.startswith("{") and r.text.endswith("}") for r in response.records) + + def test_valid_group_keeps_tokens(self, fixture_valid_iris_dataset_jsonl_and_schema, fixture_validation_config): + jsonl_str, jsonl_schema = fixture_valid_iris_dataset_jsonl_and_schema + groups_jsonl_str = BOS + jsonl_str + EOS + tokenizer = _make_mock_tokenizer() + proc = GroupedDataProcessor( + schema=jsonl_schema, + config=fixture_validation_config, + group_by="variety", + bos_token=BOS, + eos_token=EOS, + tokenizer=tokenizer, + ) + response = proc(1, groups_jsonl_str) + + assert len(response.records) == 5 + assert all(r.is_valid and r.token_count > 0 for r in response.records) diff --git a/tests/observability/test_observability.py b/tests/observability/test_observability.py index 42c98977a..12f36d2cf 100644 --- a/tests/observability/test_observability.py +++ b/tests/observability/test_observability.py @@ -263,6 +263,15 @@ def test_renders_flat_dict(self): assert "Rate" in result assert "95.00%" in result # Formatted as percentage + def test_sec_suffix_not_formatted_as_percentage(self): + """Keys ending in _sec or _seconds are rendered as plain floats, not percentages.""" + data = {"tokenization_overhead_sec": 0.35, "total_seconds": 0.72, "rate": 0.95} + result = _render_rich_table(data) + + assert "0.35" in result + assert "0.72" in result + assert "95.00%" in result + def test_renders_nested_dict(self): """Test rendering a nested statistics dictionary.""" data = { diff --git a/tests/test_results.py b/tests/test_results.py new file mode 100644 index 000000000..76030e8f1 --- /dev/null +++ b/tests/test_results.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import sys +from unittest.mock import MagicMock, patch + +import pandas as pd + +from nemo_safe_synthesizer.config import SafeSynthesizerTiming +from nemo_safe_synthesizer.generation.results import GenerateJobResults +from nemo_safe_synthesizer.generation.utils import GenerationStatus +from nemo_safe_synthesizer.results import make_nss_summary + + +def _stub_generate_results(*, with_token_stats: bool = True) -> GenerateJobResults: + """Build a minimal GenerateJobResults, optionally with token stats.""" + if with_token_stats: + return GenerateJobResults( + df=pd.DataFrame({"col": [1, 2]}), + status=GenerationStatus.COMPLETE, + num_valid_records=2, + num_invalid_records=1, + num_prompts=4, + valid_record_fraction=2 / 3, + batch_valid_record_fractions=[0.67], + elapsed_time=2.0, + num_completion_tokens=1000, + num_valid_record_tokens=700, + num_invalid_record_tokens=200, + num_non_record_tokens=100, + tokens_per_prompt=250.0, + tokens_per_second=500.0, + valid_tokens_per_second=350.0, + tokenization_overhead_sec=0.05, + ) + + return GenerateJobResults( + df=pd.DataFrame({"col": [1, 2]}), + status=GenerationStatus.COMPLETE, + num_valid_records=2, + num_invalid_records=1, + num_prompts=4, + valid_record_fraction=2 / 3, + batch_valid_record_fractions=[0.67], + elapsed_time=2.0, + ) + + +class TestMakeNssSummaryTokenStats: + """Verify make_nss_summary propagates token stats from GenerateJobResults.""" + + def test_propagates_token_stats(self): + timing = SafeSynthesizerTiming(generation_time_sec=2.0) + results = _stub_generate_results(with_token_stats=True) + summary = make_nss_summary(timing, results) + + assert summary.num_completion_tokens == 1000 + assert summary.num_valid_record_tokens == 700 + assert summary.num_invalid_record_tokens == 200 + assert summary.num_non_record_tokens == 100 + assert summary.valid_record_token_fraction == 700 / 1000 + assert summary.tokens_per_prompt == 250.0 + assert summary.tokens_per_second == 500.0 + assert summary.valid_tokens_per_second == 350.0 + assert summary.tokenization_overhead_sec == 0.05 + + def test_no_token_stats_leaves_fields_none(self): + timing = SafeSynthesizerTiming(generation_time_sec=2.0) + results = _stub_generate_results(with_token_stats=False) + summary = make_nss_summary(timing, results) + + assert summary.num_completion_tokens is None + assert summary.tokens_per_second is None + assert summary.valid_tokens_per_second is None + + def test_dataframe_input_leaves_token_fields_none(self): + timing = SafeSynthesizerTiming(generation_time_sec=2.0) + summary = make_nss_summary(timing, pd.DataFrame({"col": [1]})) + + assert summary.num_completion_tokens is None + assert summary.tokens_per_second is None + + def test_none_input_leaves_token_fields_none(self): + timing = SafeSynthesizerTiming(generation_time_sec=2.0) + summary = make_nss_summary(timing, None) + + assert summary.num_completion_tokens is None + + +class TestSafeSynthesizerSummaryLogWandb: + """Verify log_wandb includes token metrics in the W&B payload.""" + + def test_log_wandb_includes_token_metrics(self): + timing = SafeSynthesizerTiming(generation_time_sec=2.0) + results = _stub_generate_results(with_token_stats=True) + summary = make_nss_summary(timing, results) + + mock_wandb = MagicMock() + mock_wandb.run = MagicMock() + with patch.dict(sys.modules, {"wandb": mock_wandb}): + summary.log_wandb() + + assert mock_wandb.log.call_count == 1 + logged = mock_wandb.log.call_args[0][0] + + assert logged["gen/num_completion_tokens"] == 1000 + assert logged["gen/num_valid_record_tokens"] == 700 + assert logged["gen/num_invalid_record_tokens"] == 200 + assert logged["gen/num_non_record_tokens"] == 100 + assert logged["gen/valid_record_token_fraction"] == 700 / 1000 + assert logged["gen/tokens_per_prompt"] == 250.0 + assert logged["gen/tokens_per_second"] == 500.0 + assert logged["gen/valid_tokens_per_second"] == 350.0 + assert logged["gen/tokenization_overhead_sec"] == 0.05 + + def test_log_wandb_logs_none_token_metrics(self): + """None-valued token metrics are logged (not filtered) to keep dashboards consistent.""" + timing = SafeSynthesizerTiming(generation_time_sec=2.0) + results = _stub_generate_results(with_token_stats=False) + summary = make_nss_summary(timing, results) + + mock_wandb = MagicMock() + mock_wandb.run = MagicMock() + with patch.dict(sys.modules, {"wandb": mock_wandb}): + summary.log_wandb() + + assert mock_wandb.log.call_count == 1 + logged = mock_wandb.log.call_args[0][0] + assert logged["gen/num_completion_tokens"] is None + assert logged["gen/tokens_per_second"] is None + + def test_log_wandb_noop_when_no_run(self): + timing = SafeSynthesizerTiming(generation_time_sec=2.0) + summary = make_nss_summary(timing, None) + + mock_wandb = MagicMock() + mock_wandb.run = None + with patch.dict(sys.modules, {"wandb": mock_wandb}): + summary.log_wandb() + + mock_wandb.log.assert_not_called()