feat(generation): Track generation tokens and refactor with ParsedRecord

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)