fix: allow relaxing numeric ranges for highly numeric datasets

[mckornfield]

Summary

Pre-Review Checklist

Ensure that the following pass:

• mise run format && mise run check or via prek validation.
• mise run test passes locally
• mise run test:e2e passes locally
• mise run 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 #

Summary by CodeRabbit

• New Features

  • Added an option to enable or disable validation of generated floating-point values against observed numeric ranges.
  • Floating-point range validation is enabled by default.
  • When disabled, integer limits and enum constraints continue to be enforced.
  • Generation schemas now apply these settings consistently, including nullable numeric fields.

• Tests

  • Added coverage for numeric range handling, preserved constraints, nullable fields, and schema integrity.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 14987715-5231-458f-a237-c6348e06ecb8

📥 Commits

Reviewing files that changed from the base of the PR and between d73679a and b7f1eea.

📒 Files selected for processing (4)

• src/nemo_safe_synthesizer/config/generate.py
• src/nemo_safe_synthesizer/data_processing/dataset.py
• src/nemo_safe_synthesizer/generation/vllm_backend.py
• tests/data_processing/test_dataset.py

📜 Recent review details

⏰ Context from checks skipped due to timeout. (9)

• GitHub Check: End-user Wheel Install
• GitHub Check: Unit Tests (3.11)
• GitHub Check: Unit Tests (3.12)
• GitHub Check: Unit Tests (3.13)
• GitHub Check: Smoke Tests
• GitHub Check: Greptile Review
• GitHub Check: Typecheck
• GitHub Check: Analyze (Python)
• GitHub Check: Analyze (python)

🧰 Additional context used

📓 Path-based instructions (15)

**/*.{md,markdown,py}

📄 CodeRabbit inference engine (.cursor/rules/agent-markdown-style.mdc)

**/*.{md,markdown,py}: Avoid decorative bold (**text**) in list items, body text, and docstrings; use structural cues (headers, list markers, colons, backticks) for emphasis instead
Use backticks for code identifiers, paths, and CLI commands in markdown and docstrings

Files:

• src/nemo_safe_synthesizer/config/generate.py
• src/nemo_safe_synthesizer/data_processing/dataset.py
• src/nemo_safe_synthesizer/generation/vllm_backend.py
• tests/data_processing/test_dataset.py

**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

**/*.py: Place durable implementation guidance in function and class docstrings for public contracts and source comments for local invariants
Target Python 3.11–3.13 with modern syntax (X | Y, list[str], Self). Python 3.14+ is not supported

**/*.py: Use American English spelling in Python code, documentation, and messages.
Use Field(description=...) for every Pydantic model field.
Use assignment-style Field() by default; use Annotated only for additional metadata such as validators, constrained aliases, or discriminated unions.
Use @dataclass(frozen=True) for immutable value objects and validators; use mutable dataclasses only for builders, accumulators, and pipeline state.
Use field(default_factory=list) instead of mutable list defaults.
Use StrEnum for string-valued configuration or serialization enums and plain Enum for internal constants.
Obtain loggers with observability.get_logger(__name__); do not call logging.getLogger() or structlog.get_logger() directly.
Use .runtime, .user, and .system category loggers appropriately.
Do not use print() for operational library output; use the approved logger, click.echo(), or sys.stdout.write() where appropriate.
Use extra={} for machine-queryable logging data and f-strings only for human-readable context.
Raise errors from the custom Safe Synthesizer error hierarchy, using the documented dual inheritance for user and internal errors.
Keep shared package code compatible with Python 3.11; do not use Python 3.12-only syntax such as PEP 695 type statements or bracketed generic parameters.
Prefer X | Y, built-in collection generics, and Self over Optional, Union, and legacy typing collections.
Use collection ABCs for function arguments and concrete collection types for return values.
Use Protocol for structural subtyping and avoid Any when object, generics, or protocols are suitable.
Use TYPE_CHECKING guards for heavy imports such as pandas, torch, and transformers.
...

Files:

• src/nemo_safe_synthesizer/config/generate.py
• src/nemo_safe_synthesizer/data_processing/dataset.py
• src/nemo_safe_synthesizer/generation/vllm_backend.py
• tests/data_processing/test_dataset.py

src/**/config/**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

Use NSSBaseModel for user-facing configuration and parameter models.

Files:

• src/nemo_safe_synthesizer/config/generate.py

src/**/*.py

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

src/**/*.py: Use relative imports for package code under src/.
Do not use assert for validation in library code; raise an appropriate exception instead.

Write Google-style docstrings for public Python APIs because API reference pages are generated from source docstrings.

Files:

• src/nemo_safe_synthesizer/config/generate.py
• src/nemo_safe_synthesizer/data_processing/dataset.py
• src/nemo_safe_synthesizer/generation/vllm_backend.py

---

⚙️ CodeRabbit configuration file

Review library code against STYLE_GUIDE.md. Focus on behavior, API contracts, error handling, resource cleanup, typing, logging, and user-facing failures. Public APIs and nontrivial functions need Google-style docstrings.

Files:

• src/nemo_safe_synthesizer/config/generate.py
• src/nemo_safe_synthesizer/data_processing/dataset.py
• src/nemo_safe_synthesizer/generation/vllm_backend.py

**/*

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

**/*: Every source file requires the SPDX copyright and license header appropriate to its file format.
End files with a newline, remove trailing whitespace, use one space between sentences, and keep code, comments, and docstrings within 120 characters.

**/*: All contributions must use verified Git commits and DCO sign-off; unsigned or unsigned-off commits cannot be merged.
Branches other than main must follow <author>/<description>, optionally including an issue ID or type; branch names must use lowercase alphanumeric characters and hyphens.
Commits merged to main must follow Conventional Commits, using a valid lowercase type and a description of at most 100 characters.

Files:

• src/nemo_safe_synthesizer/config/generate.py
• src/nemo_safe_synthesizer/data_processing/dataset.py
• src/nemo_safe_synthesizer/generation/vllm_backend.py
• tests/data_processing/test_dataset.py

---

⚙️ CodeRabbit configuration file

**/*: Review as a senior maintainer for NeMo Safe Synthesizer. Prioritize issues that can change behavior, break user workflows, weaken privacy guarantees, hide failures, make tests unreliable, or create maintenance risk. Avoid generic style commentary unless it points to a concrete project convention that automated tools will not catch.
Comment only when the finding is actionable and tied to changed code. For each finding, state the impact, the condition that triggers it, and the smallest practical fix. Prefer one precise comment over broad advice. Do not ask for refactors outside the PR scope unless the changed code creates the problem.
Review type guidance: - Potential issue: use for correctness bugs, data loss, privacy leaks,
security risks, broken public APIs, invalid config behavior, missing
validation, hidden failures, nondeterministic tests, or CI breakage.

• Refactor suggestion: use for local maintainability problems introduced
  by the diff when they have clear future cost, such as duplicated setup,
  unclear boundaries, over-mocking, avoidable complexity, or opaque test
  helpers.
• Nitpick: avoid in chill mode. Do not emit formatting, import-order,
  wording, or style-only comments unless automated tools cannot catch the
  issue and it affects maintainability.

Severity guidance: - Critical: security/privacy leaks, data loss, training/test/holdout
contamination, or broken release/package/core pipeline execution.

• Major: incorrect generation/training/evaluation behavior, broken
  CLI/SDK public API, invalid config defaults or validators, or GPU/vLLM
  cleanup and process-isolation bugs likely to fail CI or production
  runs.
• Minor: localized bugs, missing focused tests for changed behavior, or
  bad test patterns that weaken regression coverage.
• Trivial: small cleanup with no behavior impact. Usually suppress in
  chill mode.
• Info: context only. Avoid unless it helps reviewers understand risk.
  Safe-Synthesizer-specific review focus: - Data ...

Files:

• src/nemo_safe_synthesizer/config/generate.py
• src/nemo_safe_synthesizer/data_processing/dataset.py
• src/nemo_safe_synthesizer/generation/vllm_backend.py
• tests/data_processing/test_dataset.py

**/*.{py,sh,yaml,yml,md}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

All Python, shell, YAML, YML, and Markdown source files require SPDX copyright headers.

Files:

• src/nemo_safe_synthesizer/config/generate.py
• src/nemo_safe_synthesizer/data_processing/dataset.py
• src/nemo_safe_synthesizer/generation/vllm_backend.py
• tests/data_processing/test_dataset.py

**/*.{py,pyi}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

Shared Python package code must remain compatible with Python 3.11 syntax; do not use Python 3.12-only syntax such as PEP 695 type statements or bracketed generic class/function parameters.

Files:

• src/nemo_safe_synthesizer/config/generate.py
• src/nemo_safe_synthesizer/data_processing/dataset.py
• src/nemo_safe_synthesizer/generation/vllm_backend.py
• tests/data_processing/test_dataset.py

**/*.{py,md}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

Use the repository's documented Python and Markdown style conventions and validate changes with the pinned mise formatting and checking tasks.

Files:

• src/nemo_safe_synthesizer/config/generate.py
• src/nemo_safe_synthesizer/data_processing/dataset.py
• src/nemo_safe_synthesizer/generation/vllm_backend.py
• tests/data_processing/test_dataset.py

**/*.{py,sh}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

Use the repository's pinned mise tasks for formatting, linting, type checking, and testing rather than invoking ruff or ty directly for project-wide checks.

Files:

• src/nemo_safe_synthesizer/config/generate.py
• src/nemo_safe_synthesizer/data_processing/dataset.py
• src/nemo_safe_synthesizer/generation/vllm_backend.py
• tests/data_processing/test_dataset.py

src/nemo_safe_synthesizer/config/**/*.py

⚙️ CodeRabbit configuration file

Treat config changes as user-facing API changes. Check Pydantic field descriptions, defaults, validators, aliases, override behavior, CLI help text impact, YAML compatibility, and documented parameter semantics.

Files:

• src/nemo_safe_synthesizer/config/generate.py

src/nemo_safe_synthesizer/data_processing/**/*.py

⚙️ CodeRabbit configuration file

Review for data-contract regressions. Check input/training/test/synthetic naming, group boundaries, token-budget math, record ordering, schema and column validation, nullable dtypes, and deterministic behavior.

Files:

• src/nemo_safe_synthesizer/data_processing/dataset.py

src/nemo_safe_synthesizer/generation/**/*.py

⚙️ CodeRabbit configuration file

Review generation changes for retry loops, stopping conditions, invalid record handling, regex/structured output contracts, backend teardown, memory cleanup, and vLLM assumptions.

Files:

• src/nemo_safe_synthesizer/generation/vllm_backend.py

**/test_*.py

📄 CodeRabbit inference engine (AGENTS.md)

Use the unit marker instead of the deprecated unit_test marker for test identification

Files:

• tests/data_processing/test_dataset.py

tests/**

📄 CodeRabbit inference engine (.cursor/rules/repo-navigation.mdc)

tests/**: Mirror src/ directory structure in tests/ directory for test organization
Auto-mark tests by directory: tests/e2e/ → e2e, tests/smoke/ → smoke, otherwise default to unit

Mirror source code directory structure in tests directory (e.g., tests/training/, tests/generation/ parallel to source structure)

Files:

• tests/data_processing/test_dataset.py

tests/**/*.py

📄 CodeRabbit inference engine (tests/TESTING.md)

tests/**/*.py: Auto-mark tests based on file path: tests under /e2e/ get e2e marker, tests under /smoke/ get smoke marker, all others get unit marker (only if no category marker already present)
Every test should have exactly one category marker: unit, smoke, or e2e
Use pytest.mark.requires_gpu modifier on tests that need CUDA hardware
Use pytest.mark.vllm on tests using vLLM generation backend and ensure each vLLM test file runs in its own process for GPU memory isolation
Use pytest.mark.slow on long-running tests
Use pytest.mark.smollm2 for SmolLM2 Hub download tests to enable process isolation
Use pytest.mark.noautouse to skip autouse fixtures for specific tests
Use load_test_dataset(filename) helper to load test datasets from tests/stub_datasets/ as HuggingFace Dataset objects
Use load_test_dataframe(filename) helper to load test data files from tests/stub_datasets/ as pandas DataFrames
Convert pandas columns to nullable dtypes (pd.Int64Dtype(), pd.BooleanDtype()) before assigning np.nan values
Use fake.seed_instance(seed) and random.seed(seed) together for Faker-based test data reproducibility
When sharing methods across multiple test files, define them in conftest.py and import them using relative imports (e.g., from .conftest import train_with_sdk); note that importing from other test files like tests/cli/helpers.py does not work
Use fixture_mock_processor or fixture_mock_processor_without_valid_records for mocking ParsedResponse objects with valid_records, invalid_records, errors, and prompt_number fields
Use pytest.importorskip to gate tests on optional dependencies that require specific extras (e.g., sentence_transformers, vllm)
Run vLLM tests with separate pytest invocations (one per file) using -n 0 (single process) for GPU memory isolation, or use staged mise tasks for CI visibility
Print statements are allowed in tests (ruff T201 is suppressed for tests/ directory) and should...

Files:

• tests/data_processing/test_dataset.py

---

⚙️ CodeRabbit configuration file

Review tests against tests/TESTING.md. Check marker usage, fixture naming, tmp_path usage, determinism, and GPU/vLLM process-isolation requirements. Flag slop tests that only check that code runs, assert result is not None when stronger invariants exist, over-mock internal implementation details, patch around the bug instead of reproducing it, or add broad snapshot/golden churn without a clear contract. Flag change detector tests that fail on harmless refactors, formatting, record ordering, incidental wording, or private implementation details without demonstrating a behavior regression. Prefer existing fixtures or focused new fixtures for repeated setup; keep tests DRY when reasonable without making the behavior under test opaque. print() is allowed in tests.

Files:

• tests/data_processing/test_dataset.py

🧠 Learnings (1)

📚 Learning: 2026-05-27T22:20:37.354Z

Learnt from: kendrickb-nvidia
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 520
File: tests/generation/test_vllm_backend.py:556-587
Timestamp: 2026-05-27T22:20:37.354Z
Learning: In NVIDIA-NeMo/Safe-Synthesizer, `tests/conftest.py`’s `pytest_collection_modifyitems` hook applies pytest category markers automatically based on each test file’s path: tests under `/e2e/` get `pytest.mark.e2e`, tests under `/smoke/` get `pytest.mark.smoke`, and all other tests get `pytest.mark.unit`. Therefore, when reviewing pytest tests outside `tests/e2e/` and `tests/smoke/`, do not flag missing explicit `pytest.mark.unit` decorators on test classes/functions as an issue (the hook will add them during collection). If a new test directory/category is introduced, ensure the hook is updated so it’s categorized correctly.

Applied to files:

• tests/data_processing/test_dataset.py

🔇 Additional comments (4)

src/nemo_safe_synthesizer/config/generate.py (1)

117-135: LGTM!

src/nemo_safe_synthesizer/data_processing/dataset.py (1)

170-210: LGTM!

src/nemo_safe_synthesizer/generation/vllm_backend.py (1)

34-34: LGTM!

Also applies to: 253-257

tests/data_processing/test_dataset.py (1)

1-77: LGTM!

---

Walkthrough

Adds configurable floating-point range enforcement. Disabled enforcement removes numeric bounds from generation schemas while retaining integer and enum constraints, with tests covering nullable types, immutability, and schemas generated from data.

Changes

Numeric range enforcement

Layer / File(s)	Summary
Schema relaxation and backend integration src/nemo_safe_synthesizer/config/generate.py, src/nemo_safe_synthesizer/data_processing/dataset.py, src/nemo_safe_synthesizer/generation/vllm_backend.py	Adds the enforce_numeric_range setting, implements removal of bounds from pure number properties, and applies it during VLLM backend initialization when disabled.
Numeric range transformation tests tests/data_processing/test_dataset.py	Covers float and nullable-float bound removal, preservation of integer and enum constraints, input immutability, and integration with generated schemas.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested labels: feature, test

Suggested reviewers: zywind

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: relaxing numeric ranges for highly numeric datasets.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch allow-numeric-range-relax/mck

---

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

[greptile-apps]

Greptile Summary

Adds an enforce_numeric_range flag to ValidationParameters (default True) that, when set to False, strips minimum/maximum bounds from pure-float JSON Schema properties before the schema is used for both structured-generation grammar constraints and post-generation jsonschema validation. This prevents wide float tables from compounding per-field range rejections that would otherwise invalidate nearly every well-formed record.

• New relax_numeric_bounds function in dataset.py: produces a non-mutating, "deep-ish" copy with float bounds removed; correctly preserves integer and enum constraints that the grammar enforces at token level.
• VllmBackend.__init__ wiring: relaxation is applied in-place on self.schema before create_processor, so both the grammar schema path (regex / json_schema / structural_tag) and jsonschema.validate in record_utils see the same relaxed schema. TimeseriesBackend inherits this for free.
• Five focused unit tests cover float relaxation, integer/enum preservation, nullable type lists, no-mutation guarantee, and a round-trip through make_json_schema.

Confidence Score: 5/5

Safe to merge — the change is opt-in (default keeps existing strict behavior), schema relaxation is applied consistently to both grammar and validation paths, and all callers receive the same schema object.

The feature is backward-compatible, the relax_numeric_bounds implementation is correct, the schema is relaxed before create_processor so both the grammar constraint and jsonschema.validate stay in sync, TimeseriesBackend inherits the fix automatically, and the test suite covers the key behavioral contracts.

No files require special attention. The only subtle point — shared non-float property dict references between the relaxed and original schemas — is safe because the original schema is immediately replaced by the relaxed one and goes unreferenced.

Important Files Changed

Filename	Overview
src/nemo_safe_synthesizer/config/generate.py	Adds enforce_numeric_range bool to ValidationParameters (default True); detailed docstring explains the float-grammar limitation and opt-out semantics. Backward-compatible.
src/nemo_safe_synthesizer/data_processing/dataset.py	New relax_numeric_bounds function correctly removes minimum/maximum from pure-number properties only; leaves integer and enum constraints untouched. Does not mutate input.
src/nemo_safe_synthesizer/generation/vllm_backend.py	Applies relax_numeric_bounds in init before create_processor, so both the grammar constraint and post-generation validation use the same relaxed schema. TimeseriesBackend inherits this for free.
tests/data_processing/test_dataset.py	Five unit tests cover float relaxation, integer/enum preservation, nullable-type lists, no-mutation guarantee, and integration with make_json_schema output.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant VB as VllmBackend.__init__
    participant DS as dataset.relax_numeric_bounds
    participant PC as create_processor
    participant PR as Processor
    participant RU as record_utils

    VB->>VB: load_json(schema_file)
    alt "enforce_numeric_range == False"
        VB->>DS: relax_numeric_bounds(self.schema)
        DS-->>VB: schema with float bounds removed
        VB->>VB: "self.schema = relaxed"
    end
    VB->>PC: create_processor(self.schema, ...)
    PC-->>VB: processor
    PR->>RU: extract_and_validate_records(text, self.schema)
    RU->>RU: jsonschema.validate(record, schema)

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant VB as VllmBackend.__init__
    participant DS as dataset.relax_numeric_bounds
    participant PC as create_processor
    participant PR as Processor
    participant RU as record_utils

    VB->>VB: load_json(schema_file)
    alt "enforce_numeric_range == False"
        VB->>DS: relax_numeric_bounds(self.schema)
        DS-->>VB: schema with float bounds removed
        VB->>VB: "self.schema = relaxed"
    end
    VB->>PC: create_processor(self.schema, ...)
    PC-->>VB: processor
    PR->>RU: extract_and_validate_records(text, self.schema)
    RU->>RU: jsonschema.validate(record, schema)

Loading

_{Reviews (1): Last reviewed commit: "fix: allow relaxing numeric ranges for h..." | Re-trigger Greptile}

[codecov]

Codecov Report

❌ Patch coverage is 96.22642% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...c/nemo_safe_synthesizer/data_processing/dataset.py	93.33%	1 Missing ⚠️
...c/nemo_safe_synthesizer/generation/vllm_backend.py	66.66%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!