feat(tokenization): refactor and improve tokenization support by binaryaaron · Pull Request #685 · NVIDIA-NeMo/Safe-Synthesizer

binaryaaron · 2026-07-30T17:50:58Z

mostly agent-cycles and review at this point with my driving and corrections; not really ready for human review just yet. agent review is fine.

draft:

Summary

Uses one native tokenizer authority for deterministic prompt and record IDs, capacity accounting, and training framing.
Routes training and generation JSONL through records_to_jsonl while preserving the current pandas compatibility bytes.
Pins production TinyLlama, Mistral, and SmolLM3 prompt, training, and persistence boundaries across tabular, grouped, and time-series modes.
Validates persisted vLLM tokenizer parity at production prompt and record seams, with explicit legacy text-prompt behavior.

Test plan

mise run check
Focused tokenizer and generation contract suite: 148 passed
A100 generation smoke with vLLM 0.26.0 and Torch 2.11.0+cu129: 2 passed
Full assembler module across eight workers: 32 passed
Decimal cache-state regression across separate workers: 2 passed

coderabbitai · 2026-07-30T17:51:05Z

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 60a1d6a5-5644-4338-a674-5425895e42bc

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

_{Comment @coderabbitai help to get the list of available commands.}

codecov · 2026-07-30T17:57:48Z

Codecov Report

❌ Patch coverage is 95.00916% with 109 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/nemo_safe_synthesizer/tokenization/core.py	85.81%	41 Missing ⚠️
tests/smoke/test_training_cpu.py	29.72%	26 Missing ⚠️
...c/nemo_safe_synthesizer/generation/vllm_backend.py	88.59%	13 Missing ⚠️
...nemo_safe_synthesizer/data_processing/assembler.py	92.42%	5 Missing ⚠️
..._safe_synthesizer/generation/timeseries_backend.py	79.16%	5 Missing ⚠️
.../nemo_safe_synthesizer/tokenization/persistence.py	94.52%	4 Missing ⚠️
src/nemo_safe_synthesizer/tokenization/cache.py	96.34%	3 Missing ⚠️
src/nemo_safe_synthesizer/tokenization/records.py	91.89%	3 Missing ⚠️
...o_safe_synthesizer/training/huggingface_backend.py	86.36%	3 Missing ⚠️
...rc/nemo_safe_synthesizer/data_processing/budget.py	91.30%	2 Missing ⚠️
... and 4 more

📢 Thoughts on this report? Let us know!

+    for count in range(10):
+        tokenization.capacity_for(prompt, context_limit=1000, sequence_count=count)
+
+    assert calls == 1


+
+class _PromptConfig(Protocol):
+    @property
+    def template(self) -> str: ...


+    def template(self) -> str: ...
+
+    @property
+    def add_bos_token_to_prompt(self) -> bool: ...


+    def add_bos_token_to_prompt(self) -> bool: ...
+
+    @property
+    def add_eos_token_to_prompt(self) -> bool: ...


+    def add_eos_token_to_prompt(self) -> bool: ...
+
+    @property
+    def bos_token(self) -> str: ...


+    def bos_token(self) -> str: ...
+
+    @property
+    def bos_token_id(self) -> int: ...


+    def bos_token_id(self) -> int: ...
+
+    @property
+    def eos_token(self) -> str: ...


+    def eos_token(self) -> str: ...
+
+    @property
+    def eos_token_id(self) -> int: ...


+
+class _Metadata(Protocol):
+    @property
+    def prompt_config(self) -> _PromptConfig: ...


+    def resolve_masks(
+        sequences: Sequence[Sequence[int]],
+        attention_masks: Sequence[Sequence[int]] | None,
+    ) -> tuple[tuple[int, ...], ...]:


Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

github-actions Bot added the area:tests label Jul 30, 2026

github-code-quality Bot found potential problems Jul 30, 2026

View reviewed changes

github-actions Bot added area:data-processing area:sdk-cli area:training area:llm area:generation labels Jul 30, 2026

binaryaaron force-pushed the feature/full-tokenizer branch from d551849 to 63c7362 Compare August 3, 2026 18:49

binaryaaron changed the base branch from main to agonzales/timeseries-prefill-serialization August 3, 2026 18:50

binaryaaron changed the title ~~feat(tokenization): add extensible tokenizer contracts~~ feat(tokenization): own native tokenizer lifecycle Aug 4, 2026

binaryaaron force-pushed the feature/full-tokenizer branch from e5f0823 to 0845a30 Compare August 4, 2026 16:52

github-actions Bot added area:ci area:config area:evaluation area:pii area:observability area:docs labels Aug 4, 2026

binaryaaron force-pushed the feature/full-tokenizer branch from 0845a30 to b6cffeb Compare August 4, 2026 17:01

binaryaaron changed the base branch from agonzales/timeseries-prefill-serialization to main August 4, 2026 17:01

binaryaaron force-pushed the feature/full-tokenizer branch from b6cffeb to 5711f66 Compare August 4, 2026 17:09

github-code-quality Bot found potential problems Aug 4, 2026

View reviewed changes

github-code-quality Bot found potential problems Aug 6, 2026

View reviewed changes

Comment thread src/nemo_safe_synthesizer/tokenization/core.py Fixed

binaryaaron force-pushed the feature/full-tokenizer branch from 0408f2d to 358bb2b Compare August 10, 2026 13:35

github-code-quality Bot found potential problems Aug 10, 2026

View reviewed changes

Comment thread src/nemo_safe_synthesizer/tokenization/core.py

Comment on lines +272 to +275

def resolve_masks(

sequences: Sequence[Sequence[int]],

attention_masks: Sequence[Sequence[int]] | None,

) -> tuple[tuple[int, ...], ...]:

binaryaaron added 4 commits August 10, 2026 21:05

feat(tokenization): add extensible tokenizer contracts

6f7cd76

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

feat(tokenization): migrate records and token cache

eeb1137

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

test(tokenization): allowlist revision fixture

4baec4a

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

test(smoke): provide tokenizer provenance

05f8407

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

binaryaaron added 14 commits August 10, 2026 21:05

feat(tokenization): own full tokenizer lifecycle

6050a47

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

style(tokenization): format persistence files

d55b400

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

feat(generation): dispatch canonical token prompts

6a4808b

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

test(tokenization): align prefill contract with generation

c153ca9

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

feat: gate vLLM generation on tokenizer parity

b126557

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

fix: fail closed on tokenizer initialization

05d2544

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

refactor: type tokenizer workload boundaries

9ef83a8

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

refactor(tokenization): replace tokenizer framework

43f1474

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

refactor(tokenization): clarify boundary parsing

e5a552d

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

test(data-processing): make shuffle cache state explicit

733faa6

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

refactor(tokenization): centralize mask validation

33fa702

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

ci: allowlist deterministic cache digest

99bd7a9

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

fix(tokenization): preserve cross-stage token identity

12746d9

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

fix(tokenization): enforce production boundary parity

9cee2aa

Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>

binaryaaron force-pushed the feature/full-tokenizer branch from 358bb2b to 9cee2aa Compare August 10, 2026 22:20

binaryaaron changed the title ~~feat(tokenization): own native tokenizer lifecycle~~ feat(tokenization): refactor and improve tokenization support Aug 12, 2026

Navigation Menu

Uh oh!

feat(tokenization): refactor and improve tokenization support - #685

feat(tokenization): refactor and improve tokenization support#685
binaryaaron wants to merge 18 commits into
mainfrom
feature/full-tokenizer

binaryaaron commented Jul 30, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Jul 30, 2026 •

edited

Loading

Review skipped

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

codecov Bot commented Jul 30, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

binaryaaron commented Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Uh oh!

coderabbitai Bot commented Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

codecov Bot commented Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

binaryaaron commented Jul 30, 2026 •

edited

Loading

coderabbitai Bot commented Jul 30, 2026 •

edited

Loading

codecov Bot commented Jul 30, 2026 •

edited

Loading