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

[seayang-nv]

Summary

Tracks token usage throughout the generation pipeline so operators can see how efficiently the LLM's output budget is being spent -- how many tokens land in valid records vs. invalid or non-record overhead.

• Data model: Introduces ParsedRecord (text, parsed dict, error, token count on one object) and reshapes ParsedResponse to hold list[ParsedRecord], replacing the prior parallel lists (valid_records / invalid_records / errors / *_token_counts) that previously threaded through extraction, validation, group-level rejection, batch aggregation, and data-action post-processing. ← new
• Each record is tokenized at extraction time (using the vLLM engine's tokenizer) and the count rides along on its ParsedRecord through validation, group-level rejection, and time-series cascade invalidation.
• New metrics are surfaced in batch progress logs, GenerateJobResults, and the final SafeSynthesizerSummary / W&B: completion token totals, valid/invalid/non-record breakdowns, tokens-per-second throughput (with and without startup), valid-tokens-per-second, and tokenization overhead.
• Training-eval callback path (InferenceEvalCallback) now passes the tokenizer through and forwards per-row completion-token counts so the same metrics apply during training-time evaluation.
• Fixes the rich table renderer so duration fields (_sec, _seconds) are displayed as plain floats instead of percentages.

Example summary for patient_events.csv

     Batch Generation Summary     
+--------------------------------+
| Metric                | Value  |
|-----------------------+--------|
| Num Prompts           | 100    |
| Num Valid Records     | 357    |
| Num Invalid Records   | 115    |
| Valid Record Fraction | 76.00% |
| Completion Tokens     | 72566  |
| Valid Record Tokens   | 51694  |
| Invalid Record Tokens | 15628  |
| Non Record Tokens     | 5244   |
+--------------------------------+
 
2026-04-15T18:03:19.083 | Nemo Safe Synthesizer |  user    |  info  |  batch.py: Batch.log_summary: 270 


           Error Statistics           
+------------------------------------+
| Metric                    | Value  |
|---------------------------+--------|
| Groupby Generation Failed | 81.00% |
| Invalid Field Value       | 13.00% |
| Invalid Json              | 6.00%  |
+------------------------------------+
 
2026-04-15T18:03:19.085 | Nemo Safe Synthesizer |  user    |  info  |  vllm_backend.py: VllmBackend._log_batch_timing_and_progress: 485 


           Batch Progress            
+-----------------------------------+
| Metric                  | Value   |
|-------------------------+---------|
| Records Per Second      | 8.31    |
| Duration Seconds        | 42.98   |
| Valid Records Generated | 753     |
| Target Records          | 1000    |
| Progress Fraction       | 75.30%  |
| Tokens Per Second       | 1688.30 |
| Valid Tokens Per Second | 1202.70 |
+-----------------------------------+

Pre-Review Checklist

Ensure that the following pass:

• make format && make check or via prek validation.
• make test passes locally
• make test-e2e passes locally
• make test-ci-container passes locally (recommended)
• GPU CI status check passes -- comment /sync on this PR to trigger a run (auto-triggers on ready-for-review)

Pre-Merge Checklist

• New or updated tests for any fix or new behavior
• Updated documentation for new features and behaviors, including docstrings for API docs.

Other Notes

• Closes Track generated and used tokens #364

[nina-xu]

Overall looks good!

Love that you have a test plan! I think some end-to-end testing should be added to the testing plan, like can we see what the changes look like in

• the log (you already have that)
• wandb
• the results summary json
• whatever else is changed

in the three modes

• non-group by
• group by
• time series?

[kendrickb-nvidia]

A few small comments, and then majority are around ExtractionResult and these parallel lists of valid, invalid, errors, valid tokens, invalid tokens that we're passing all over. I think the code would be simpler if we switch to lists of ParsedRecord where the ParsedRecord instance contains an is_valid, error messages, # tokens, etc. Instead of the ParsedResponse and ExtractionResult with several parallel lists. But that would also be a bigger refactor (and may be more complicated than I think).

Curious on other thoughts on the patterns @binaryaaron @mckornfield . @seayang-nv let's chat more tomorrow.

[binaryaaron]

Overall, looks good.

main issues

• Training path - should we do this? can we add it?
• src/nemo_safe_synthesizer/training/huggingface_backend.py:507: create_processor(...) is called without tokenizer=, so InferenceEvalCallback can't collect token metrics. self.tokenizer is in scope.
• src/nemo_safe_synthesizer/training/callbacks.py:156: batch.process(idx, text) doesn't pass completion_tokens=, so this path always records 0 tokens and the total_completion_tokens > 0 gates downstream drop the metrics.

Follow-ups

1. Batch-tokenize: one tokenizer(list_of_strings, add_special_tokens=False, return_length=True) call per extraction instead of N·K individual tokenizer.encode(str) calls. Gate on measured cost — instrument tokenization_overhead_sec / generation_time_sec in production first. <1% → defer, >5% → prioritize. An offset-mapping approach was rejected; it depends on tokenize(decode(ids)) == ids holding for the vLLM detokenizer.

2. Parallel-list bookkeeping: 5 aligned lists across 5 files, in-place surgery in _reject_group_records, and the zip(... or [0] * len(...)) workaround in _apply_data_actions_fn. Two refactors collapse it: a ValidationError NamedTuple for the (message, validator) pair, and list[ClassifiedRecord] with a RecordStatus enum so text, count, error, and status live on the same object. The _reject_group_records text/count mismatch flagged inline becomes unrepresentable under the second.

3. Aggregation: TokenStats with __add__ collapses the 10 sum-properties across Batch and GenerationBatches into one aggregable dataclass.

[binaryaaron]

Consolidating with @kendrickb-nvidia's review so @seayang-nv doesn't have to triangulate. we roughly agree on the ParsedRecord / ClassifiedRecord bit.

Direct overlap across reviews

Issue	Comments	Suggested resolution
W&B None-filter at external_results.py:190 / :195	nina, kendrick, my #10	Kendrick's: restore wandb.log(metrics) and add a comment that the None-logging is intentional.
Test helpers in conftest.py:19	kendrick, my #15	Either inline into test_processors.py (kendrick) or split out _helpers.py (me). Pick one.
Parallel valid/invalid/errors/tokens lists across ExtractionResult / ParsedResponse	both kendrick and i	Same refactor: list[ParsedRecord] with is_valid / error / tokens on the record. Kendrick's processors.py:405 flags a smaller intermediate step — drop ExtractionResult and have record_utils.py return ParsedResponse directly. Useful stepping stone toward the bigger move.

Points kendrick raised that i did not and that I endorse

• external_results.py:146: document the field invariants (num_non_record_tokens = num_completion_tokens - num_valid_record_tokens - num_invalid_record_tokens, tokens_per_completion = num_completion_tokens / num_prompts). Would have made the _reject_group_records text/count mismatch (my chore: sync 5737 from nmp #7) catchable at review time.
• processors.py:85: property that returns a callable is indeed odd. Either rename to _encode_func or collapse to def _count_tokens(self, s: str) -> int: returning 0 when the tokenizer is absent.
• record_utils.py:196: _timed_encode is a local helper, not a method — leading underscore is nonstandard here.
• processors.py:82: log tokenizer-present in the init debug line.

Stale comments

The 5 inline comments on generation_token_statistics_f26c553b.plan.md are anchored to 7ebe107c and the file was deleted in 5d084b79. Current status:

• "per-record token stats" (3081478331) — resolved, implementation is per-record.
• "time.monotonic() vs perf_counter()" (3081487882) — resolved, current code uses time.monotonic().
• "deferred tokenizer assignment" (3081497203) — still open; ties into the training-path gap in my summary (huggingface_backend.py:507, callbacks.py:156). The deferred slot exists but isn't fed on that path.
• The remaining two (total-tokens source, group-by stats semantics) should be confirmed against the current implementation.

Overall -

1. decide on training-path omission. it's not a blocker imo; generation stats are more important.
2. Restore wandb.log(metrics) without the None-filter, with a comment.
3. Decide whether to do the ParsedRecord refactor in this PR, or file it as a follow-up. doing Kendrick's smaller ExtractionResult → ParsedResponse change would be sufficient for now.

[codecov]

Codecov Report

❌ Patch coverage is 95.72954% with 24 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
..._safe_synthesizer/generation/timeseries_backend.py	11.11%	8 Missing ⚠️
src/nemo_safe_synthesizer/training/callbacks.py	0.00%	8 Missing ⚠️
...c/nemo_safe_synthesizer/generation/vllm_backend.py	50.00%	7 Missing ⚠️
src/nemo_safe_synthesizer/generation/results.py	97.72%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[binaryaaron]

Thanks for taking the ParsedRecord refactor all the way through, Sean. The new metrics surfaced through the summary / W&B / batch progress will be really helpful for tuning runs.

Ran an astnav sweep and left four suggestions inline — all optional for this PR, none of them change behavior. happy to file follow-up issues for any we don't land here.

One pre-existing observation (not inline because it's outside the diff): over in make_nss_results at src/nemo_safe_synthesizer/results.py:156, make_nss_summary(...) gets called before the None / empty-DataFrame guards on lines 157-160 run, so a None input would AttributeError inside make_nss_summary before the ValueError fires. One-line reorder to move the two guards above the call. Totally fine to skip or handle separately since this PR didn't touch it.