diff --git a/docs/tutorials/differential-privacy.ipynb b/docs/tutorials/differential-privacy.ipynb
index 90b1b09a1..ad9d9d7a6 100644
--- a/docs/tutorials/differential-privacy.ipynb
+++ b/docs/tutorials/differential-privacy.ipynb
@@ -159,7 +159,7 @@
" .with_data_source(df) # .with_replace_pii(enable=False) to disable PII replacement\n",
" .with_differential_privacy(dp_enabled=True, delta=recommended_delta)\n",
" .with_train(batch_size=16) # Override the default batch size of 1, which is designed for non-DP training\n",
- " .with_generate(use_structured_generation=True) # Improves the percentage of valid records when DP is enabled\n",
+ " .with_generate(structured_generation={'enabled': True}) # Improves the percentage of valid records when DP is enabled\n",
")\n",
"builder.run()\n",
"results = builder.results"
diff --git a/docs/user-guide/configuration.md b/docs/user-guide/configuration.md
index a2b6283c8..0f5229354 100644
--- a/docs/user-guide/configuration.md
+++ b/docs/user-guide/configuration.md
@@ -56,6 +56,40 @@ synthesizer.run()
See [Using YAML Config Files with the CLI and SDK](running.md#using-yaml-config-files)
for more detail on combining config files with runtime overrides.
+### Python Parameter Construction
+
+The Python SDK accepts both fully nested config objects and compatibility
+shortcuts for fields on top-level parameter sections:
+
+```python
+from nemo_safe_synthesizer.config import SafeSynthesizerParameters
+
+config = SafeSynthesizerParameters.from_params(
+ num_records=2000, # generation.num_records
+ dp_enabled=True, # privacy.dp_enabled
+ structured_generation={"enabled": True}, # generation.structured_generation.enabled
+)
+```
+
+Flat keyword arguments are matched by field name against the top-level parameter
+sections. Use the nested shape for fields inside nested subobjects, especially
+when a generic field name could appear in multiple places:
+
+```python
+# Preferred: unambiguous nested form.
+SafeSynthesizerParameters.from_params(
+ structured_generation={"enabled": True},
+)
+
+# Also valid: fully nested generation section.
+SafeSynthesizerParameters.from_params(
+ generation={"structured_generation": {"enabled": True}},
+)
+
+# Avoid: this configures evaluation.enabled, not structured generation.
+SafeSynthesizerParameters.from_params(enabled=True)
+```
+
---
## Training
@@ -158,10 +192,10 @@ for the full API reference.
| `generation.repetition_penalty` | `1.0` | Penalty for repeated tokens; increase slightly if generation produces repetitive output | 1.0--1.15 typical; start at 1.05 if repetition is a problem |
| `generation.patience` | `3` | Consecutive bad batches before stopping | Leave at default |
| `generation.invalid_fraction_threshold` | `0.8` | Invalid record fraction that triggers the patience counter | Leave at default |
-| `generation.use_structured_generation` | `false` | Enable structured output to constrain record format (typically at the cost of reducing the quality of generated records and increasing generation time; use when the pipeline struggles to produce valid records) | Leave off unless the pipeline cannot produce valid records |
-| `generation.structured_generation_backend` | `"auto"` | vLLM guided-decoding backend | Leave at `"auto"` |
-| `generation.structured_generation_schema_method` | `"auto"` | Schema method (`"auto"`, `"structural_tag"`, `"json_schema"`, or `"regex"`) | Leave at `"auto"`; it picks `"structural_tag"` on xgrammar-capable backends and `"regex"` otherwise |
-| `generation.structured_generation_use_single_sequence` | `false` | Match exactly one sequence when `max_sequences_per_example` is 1 | Leave at default |
+| `generation.structured_generation.enabled` | `false` | Enable structured output to constrain record format (typically at the cost of reducing the quality of generated records and increasing generation time; use when the pipeline struggles to produce valid records) | Leave off unless the pipeline cannot produce valid records |
+| `generation.structured_generation.backend` | `"auto"` | vLLM guided-decoding backend | Leave at `"auto"` |
+| `generation.structured_generation.schema_method` | `"auto"` | Schema method (`"auto"`, `"structural_tag"`, `"json_schema"`, or `"regex"`) | Leave at `"auto"`; it picks `"structural_tag"` on xgrammar-capable backends and `"regex"` otherwise |
+| `generation.structured_generation.use_single_sequence` | `false` | Match exactly one sequence when `max_sequences_per_example` is 1 | Leave at default |
| `generation.enforce_timeseries_fidelity` | `false` | Enforce time series order, intervals, and timestamps | Enable for time series data |
| `generation.attention_backend` | `"auto"` | vLLM attention backend | Leave at `"auto"` |
diff --git a/docs/user-guide/running.md b/docs/user-guide/running.md
index 07f84c42e..4018c30a9 100644
--- a/docs/user-guide/running.md
+++ b/docs/user-guide/running.md
@@ -864,16 +864,16 @@ flowchart TD
### Structured Generation
-Set `generation.use_structured_generation` to `true` to constrain the model's
-output so every record matches the dataset schema. This reduces the fraction of
-invalid records, typically at the cost of reducing the quality of the generated
-records. Use it when the pipeline struggles to produce valid records.
+Set `generation.structured_generation.enabled` to `true` to constrain the model's
+output toward the dataset schema. This usually reduces the fraction of invalid
+records, typically at the cost of reducing the quality of the generated records.
+Use it when the pipeline struggles to produce valid records.
=== "CLI"
```bash
safe-synthesizer run \
- --generation__use_structured_generation true \
+ --generation__structured_generation__enabled true \
--data-source data.csv
```
@@ -885,19 +885,24 @@ records. Use it when the pipeline struggles to produce valid records.
synthesizer = (
SafeSynthesizer()
.with_data_source("data.csv")
- .with_generate(use_structured_generation=True)
+ .with_generate(structured_generation={"enabled": True})
)
```
+ Pass `structured_generation` as a nested dict or
+ `StructuredGenerationParameters` object. Do not use a bare `enabled=True`
+ shortcut for structured generation.
+
=== "Config reference"
```yaml
generation:
- use_structured_generation: true
- structured_generation_schema_method: "auto"
+ structured_generation:
+ enabled: true
+ schema_method: "auto"
```
-- `"auto"`: picks `"structural_tag"` when `structured_generation_backend` is `"auto"` or `"xgrammar"`, otherwise `"regex"`.
+- `"auto"`: picks `"structural_tag"` when `structured_generation.backend` is `"auto"` or `"xgrammar"`, otherwise `"regex"`.
- `"structural_tag"`: uses XGrammar Structural Tag to compose schema-constrained JSONL output.
- `"regex"`: constructs a custom regex from the dataset schema. More comprehensive but slower.
- `"json_schema"`: passes a JSON Schema to the backend. Faster, but may miss edge cases.
@@ -941,8 +946,9 @@ fraction per batch.
!!! tip "Early stopping"
If the pipeline stops early due to patience, try enabling
- `use_structured_generation: true` to constrain outputs to the dataset
- schema, or lower `temperature` to reduce the chance of malformed records.
+ `generation.structured_generation.enabled: true` to constrain outputs to
+ the dataset schema, or lower `temperature` to reduce the chance of malformed
+ records.
See [Configuration Reference -- Generation](configuration.md#generation) for the full parameter table.
diff --git a/docs/user-guide/troubleshooting.md b/docs/user-guide/troubleshooting.md
index 6f6f34d2d..194b3e9ce 100644
--- a/docs/user-guide/troubleshooting.md
+++ b/docs/user-guide/troubleshooting.md
@@ -379,7 +379,7 @@ Generation stopped prematurely because the average fraction of invalid records w
```
: Too many invalid records across `generation.patience` consecutive batches.
- Consider retraining with more records, adjusting `training.number_of_input_records_to_sample`, or setting `use_structured_generation=True`.
+ Consider retraining with more records, adjusting `training.num_input_records_to_sample`, or setting `generation.structured_generation.enabled=true`.
For context-length errors during data assembly (`"The number of tokens in an
example exceeds the available context length"`), see
diff --git a/src/nemo_safe_synthesizer/config/__init__.py b/src/nemo_safe_synthesizer/config/__init__.py
index 50a6063be..bcd17089f 100644
--- a/src/nemo_safe_synthesizer/config/__init__.py
+++ b/src/nemo_safe_synthesizer/config/__init__.py
@@ -9,7 +9,7 @@
from .differential_privacy import DifferentialPrivacyHyperparams
from .evaluate import EvaluationParameters
from .external_results import SafeSynthesizerSummary, SafeSynthesizerTiming
-from .generate import GenerateParameters
+from .generate import GenerateParameters, StructuredGenerationParameters
from .internal_results import SafeSynthesizerResults
from .job import SafeSynthesizerJobConfig
from .parameters import SafeSynthesizerParameters
@@ -31,6 +31,7 @@
"SafeSynthesizerResults",
"SafeSynthesizerSummary",
"SafeSynthesizerTiming",
+ "StructuredGenerationParameters",
"TimeSeriesParameters",
"TrainingHyperparams",
]
diff --git a/src/nemo_safe_synthesizer/config/generate.py b/src/nemo_safe_synthesizer/config/generate.py
index 8eca34108..ff02495fe 100644
--- a/src/nemo_safe_synthesizer/config/generate.py
+++ b/src/nemo_safe_synthesizer/config/generate.py
@@ -3,7 +3,8 @@
from __future__ import annotations
-from typing import Annotated, Literal, Self
+import warnings
+from typing import Annotated, Any, ClassVar, Literal, Self
from pydantic import (
BaseModel,
@@ -29,6 +30,7 @@
__all__ = [
"GenerateParameters",
"ResolvedStructuredGenerationSchemaMethod",
+ "StructuredGenerationParameters",
"StructuredGenerationBackend",
"StructuredGenerationSchemaMethod",
"STRUCTURAL_TAG_COMPATIBLE_BACKENDS",
@@ -66,8 +68,8 @@ def structural_tag_backend_error_message(backend: str) -> str | None:
return None
return (
"Invalid structured generation configuration: "
- "`structured_generation_schema_method='structural_tag'` requires "
- f"`structured_generation_backend` to be 'xgrammar' or 'auto', got {backend!r}."
+ "`schema_method='structural_tag'` requires "
+ f"`backend` to be 'xgrammar' or 'auto', got {backend!r}."
)
@@ -111,6 +113,67 @@ class ValidationParameters(Parameters, BaseModel):
] = False
+class StructuredGenerationParameters(Parameters, BaseModel):
+ """Configuration for vLLM structured generation.
+
+ These parameters control whether generation is constrained to schema-shaped
+ output, which backend enforces the constraint, and how the constraint schema
+ is built.
+ """
+
+ enabled: Annotated[
+ bool,
+ Field(
+ title="enabled",
+ description="Whether to use structured generation for better format control.",
+ ),
+ ] = False
+
+ backend: Annotated[
+ StructuredGenerationBackend,
+ Field(
+ title="backend",
+ description=(
+ "The backend used by vLLM when structured generation is enabled. "
+ "Supported backends: 'outlines', 'guidance', 'xgrammar', 'lm-format-enforcer'. "
+ "'auto' will allow vLLM to choose the backend."
+ ),
+ ),
+ ] = "auto"
+
+ schema_method: Annotated[
+ StructuredGenerationSchemaMethod,
+ Field(
+ title="schema_method",
+ description=(
+ "The method used to generate the schema from your dataset and pass it to the generation backend. "
+ "'auto' picks 'structural_tag' on xgrammar-capable backends and 'regex' otherwise. "
+ "'regex' uses a custom regex construction method that tends to be more comprehensive "
+ "than 'json_schema' at the cost of speed. 'structural_tag' uses XGrammar Structural Tag "
+ "to compose schema-constrained JSONL output."
+ ),
+ ),
+ ] = "auto"
+
+ use_single_sequence: Annotated[
+ bool,
+ Field(
+ title="use_single_sequence",
+ description="Whether to use a regex that matches exactly one sequence or record if ``max_sequences_per_example`` is 1.",
+ ),
+ ] = False
+
+ @model_validator(mode="after")
+ def _validate_structural_tag_backend(self) -> Self:
+ if not self.enabled:
+ return self
+ if self.schema_method != "structural_tag":
+ return self
+ if message := structural_tag_backend_error_message(self.backend):
+ raise ParameterError(message)
+ return self
+
+
class GenerateParameters(Parameters, BaseModel):
"""Configuration parameters for synthetic data generation.
@@ -176,47 +239,10 @@ class GenerateParameters(Parameters, BaseModel):
),
] = 0.8
- use_structured_generation: Annotated[
- bool,
- Field(
- title="use_structured_generation",
- description="Whether to use structured generation for better format control.",
- ),
- ] = False
-
- structured_generation_backend: Annotated[
- Literal["auto", "xgrammar", "guidance", "outlines", "lm-format-enforcer"],
- Field(
- title="structured_generation_backend",
- description=(
- "The backend used by vLLM when ``use_structured_generation`` is ``True``. "
- "Supported backends: 'outlines', 'guidance', 'xgrammar', 'lm-format-enforcer'. "
- "'auto' will allow vLLM to choose the backend."
- ),
- ),
- ] = "auto"
-
- structured_generation_schema_method: Annotated[
- StructuredGenerationSchemaMethod,
- Field(
- title="structured_generation_schema_method",
- description=(
- "The method used to generate the schema from your dataset and pass it to the generation backend. "
- "'auto' picks 'structural_tag' on xgrammar-capable backends and 'regex' otherwise. "
- "'regex' uses a custom regex construction method that tends to be more comprehensive "
- "than 'json_schema' at the cost of speed. 'structural_tag' uses XGrammar Structural Tag "
- "to compose schema-constrained JSONL output."
- ),
- ),
- ] = "auto"
-
- structured_generation_use_single_sequence: Annotated[
- bool,
- Field(
- title="structured_generation_use_single_sequence",
- description="Whether to use a regex that matches exactly one sequence or record if ``max_sequences_per_example`` is 1.",
- ),
- ] = False
+ structured_generation: StructuredGenerationParameters = Field(
+ description="Structured generation parameters controlling schema-constrained output.",
+ default_factory=StructuredGenerationParameters,
+ )
# TODO: We will merge this with `timestamp_validation_mode` described in the MR !5153
enforce_timeseries_fidelity: Annotated[
@@ -244,12 +270,118 @@ class GenerateParameters(Parameters, BaseModel):
),
] = "auto"
- @model_validator(mode="after")
- def _validate_structural_tag_backend(self) -> Self:
- if not self.use_structured_generation:
- return self
- if self.structured_generation_schema_method != "structural_tag":
- return self
- if message := structural_tag_backend_error_message(self.structured_generation_backend):
- raise ParameterError(message)
- return self
+ _STRUCTURED_GENERATION_LEGACY_FIELDS: ClassVar[dict[str, str]] = {
+ "use_structured_generation": "enabled",
+ "structured_generation_backend": "backend",
+ "structured_generation_schema_method": "schema_method",
+ "structured_generation_use_single_sequence": "use_single_sequence",
+ }
+
+ @model_validator(mode="before")
+ @classmethod
+ def _migrate_legacy_structured_generation_fields(cls, data: Any) -> Any:
+ if not isinstance(data, dict):
+ return data
+
+ values = dict(data)
+ legacy = {
+ new_name: values.pop(old_name)
+ for old_name, new_name in cls._STRUCTURED_GENERATION_LEGACY_FIELDS.items()
+ if old_name in values
+ }
+ if not legacy:
+ return values
+
+ structured_generation = values.get("structured_generation")
+ match structured_generation:
+ case StructuredGenerationParameters() as params:
+ structured_values = params.model_dump()
+ case BaseModel() as model:
+ structured_values = model.model_dump()
+ case dict() as mapping:
+ structured_values = dict(mapping)
+ case None:
+ structured_values = {}
+ case _:
+ return values
+
+ # Legacy flat keys are treated as explicit overrides for migration
+ # paths such as ``from_params(generation={...}, structured_generation_backend=...)``.
+ values["structured_generation"] = structured_values | legacy
+ return values
+
+ @property
+ def use_structured_generation(self) -> bool:
+ """Deprecated flat alias for ``structured_generation.enabled``."""
+ warnings.warn(
+ "use_structured_generation is deprecated; use structured_generation.enabled instead.",
+ DeprecationWarning,
+ stacklevel=2,
+ )
+ return self.structured_generation.enabled
+
+ @use_structured_generation.setter
+ def use_structured_generation(self, value: bool) -> None:
+ warnings.warn(
+ "use_structured_generation is deprecated; use structured_generation.enabled instead.",
+ DeprecationWarning,
+ stacklevel=2,
+ )
+ self.structured_generation.enabled = value
+
+ @property
+ def structured_generation_backend(self) -> StructuredGenerationBackend:
+ """Deprecated flat alias for ``structured_generation.backend``."""
+ warnings.warn(
+ "structured_generation_backend is deprecated; use structured_generation.backend instead.",
+ DeprecationWarning,
+ stacklevel=2,
+ )
+ return self.structured_generation.backend
+
+ @structured_generation_backend.setter
+ def structured_generation_backend(self, value: StructuredGenerationBackend) -> None:
+ warnings.warn(
+ "structured_generation_backend is deprecated; use structured_generation.backend instead.",
+ DeprecationWarning,
+ stacklevel=2,
+ )
+ self.structured_generation.backend = value
+
+ @property
+ def structured_generation_schema_method(self) -> StructuredGenerationSchemaMethod:
+ """Deprecated flat alias for ``structured_generation.schema_method``."""
+ warnings.warn(
+ "structured_generation_schema_method is deprecated; use structured_generation.schema_method instead.",
+ DeprecationWarning,
+ stacklevel=2,
+ )
+ return self.structured_generation.schema_method
+
+ @structured_generation_schema_method.setter
+ def structured_generation_schema_method(self, value: StructuredGenerationSchemaMethod) -> None:
+ warnings.warn(
+ "structured_generation_schema_method is deprecated; use structured_generation.schema_method instead.",
+ DeprecationWarning,
+ stacklevel=2,
+ )
+ self.structured_generation.schema_method = value
+
+ @property
+ def structured_generation_use_single_sequence(self) -> bool:
+ """Deprecated flat alias for ``structured_generation.use_single_sequence``."""
+ warnings.warn(
+ "structured_generation_use_single_sequence is deprecated; use structured_generation.use_single_sequence instead.",
+ DeprecationWarning,
+ stacklevel=2,
+ )
+ return self.structured_generation.use_single_sequence
+
+ @structured_generation_use_single_sequence.setter
+ def structured_generation_use_single_sequence(self, value: bool) -> None:
+ warnings.warn(
+ "structured_generation_use_single_sequence is deprecated; use structured_generation.use_single_sequence instead.",
+ DeprecationWarning,
+ stacklevel=2,
+ )
+ self.structured_generation.use_single_sequence = value
diff --git a/src/nemo_safe_synthesizer/config/parameters.py b/src/nemo_safe_synthesizer/config/parameters.py
index 9829de50b..ceee20d14 100644
--- a/src/nemo_safe_synthesizer/config/parameters.py
+++ b/src/nemo_safe_synthesizer/config/parameters.py
@@ -4,6 +4,7 @@
from __future__ import annotations
import warnings
+from collections.abc import Mapping
from typing import Any, Self
from pydantic import BaseModel, Field, model_validator
@@ -65,6 +66,23 @@ def _overlay_set_fields(saved: Parameters, runtime: Parameters) -> Parameters:
return saved.model_validate(merge_dicts(saved.model_dump(), overrides))
+def _section_values(kwargs: dict[str, Any], section: str) -> dict[str, Any]:
+ """Return values for a nested section merged with flat compatibility kwargs."""
+ section_value = kwargs.get(section)
+ match section_value:
+ case BaseModel() as model:
+ section_data = model.model_dump()
+ case Mapping() as mapping:
+ section_data = dict(mapping)
+ case None:
+ section_data = {}
+ case _:
+ section_data = {section: section_value}
+
+ flat_values = {key: value for key, value in kwargs.items() if key != section}
+ return section_data | flat_values
+
+
class SafeSynthesizerParameters(Parameters):
"""Main configuration class for the Safe Synthesizer pipeline.
@@ -212,14 +230,14 @@ def from_params(cls, **kwargs) -> "SafeSynthesizerParameters":
Example:
>>> from nemo_safe_synthesizer.config import SafeSynthesizerParameters
- >>> SafeSynthesizerParameters.from_params(use_structured_generation=True)
+ >>> SafeSynthesizerParameters.from_params(structured_generation={"enabled": True})
"""
- thp = TrainingHyperparams().model_copy(update=kwargs)
- gp = GenerateParameters().model_copy(update=kwargs)
- ep = EvaluationParameters().model_copy(update=kwargs)
- pp = DifferentialPrivacyHyperparams().model_copy(update=kwargs)
- dp = DataParameters().model_copy(update=kwargs)
- tsp = TimeSeriesParameters().model_copy(update=kwargs)
+ thp = TrainingHyperparams.model_validate(_section_values(kwargs, "training"))
+ gp = GenerateParameters.model_validate(_section_values(kwargs, "generation"))
+ ep = EvaluationParameters.model_validate(_section_values(kwargs, "evaluation"))
+ pp = DifferentialPrivacyHyperparams.model_validate(_section_values(kwargs, "privacy"))
+ dp = DataParameters.model_validate(_section_values(kwargs, "data"))
+ tsp = TimeSeriesParameters.model_validate(_section_values(kwargs, "time_series"))
extra: dict[str, Any] = {
"training": thp,
diff --git a/src/nemo_safe_synthesizer/configurator/pydantic_click_options.py b/src/nemo_safe_synthesizer/configurator/pydantic_click_options.py
index d24b7bf9b..073f29bfe 100644
--- a/src/nemo_safe_synthesizer/configurator/pydantic_click_options.py
+++ b/src/nemo_safe_synthesizer/configurator/pydantic_click_options.py
@@ -32,6 +32,14 @@
__all__ = ["pydantic_options", "parse_overrides", "AutoParamType"]
+_LEGACY_CLI_OPTION_PATHS: dict[str, tuple[str, ...]] = {
+ "generation.structured_generation.enabled": ("generation.use_structured_generation",),
+ "generation.structured_generation.backend": ("generation.structured_generation_backend",),
+ "generation.structured_generation.schema_method": ("generation.structured_generation_schema_method",),
+ "generation.structured_generation.use_single_sequence": ("generation.structured_generation_use_single_sequence",),
+}
+"""Hidden compatibility aliases for renamed generated CLI options."""
+
def parse_overrides(values: dict[str, Any] | None = None, field_sep: str = "__") -> dict[str, Any]:
"""Parse Click kwargs into a nested override dict.
@@ -309,6 +317,15 @@ def pydantic_options(model_class: type[BaseModel], field_separator: str = "__"):
A Click decorator that attaches the generated options to a command.
"""
+ def apply_leaf_option(f, name: str, field: FieldInfo, *, hidden: bool = False):
+ names = _option_names(name, field_separator)
+ return click.option(
+ *names,
+ type=_click_type(field.annotation),
+ help=field.description or "",
+ hidden=hidden,
+ )(f)
+
def decorator(f):
for param in sorted(_collect_params(model_class), key=lambda p: p.name):
match param:
@@ -330,12 +347,9 @@ def decorator(f):
help=f"Disable {field_name.replace('_', '-')} entirely.",
)(f)
case LeafParam(field=field):
- names = _option_names(param.name, field_separator)
- f = click.option(
- *names,
- type=_click_type(field.annotation),
- help=field.description or "",
- )(f)
+ f = apply_leaf_option(f, param.name, field)
+ for legacy_name in _LEGACY_CLI_OPTION_PATHS.get(param.name, ()):
+ f = apply_leaf_option(f, legacy_name, field, hidden=True)
return f
return decorator
diff --git a/src/nemo_safe_synthesizer/generation/regex_manager.py b/src/nemo_safe_synthesizer/generation/regex_manager.py
index 990f34437..b67d51665 100644
--- a/src/nemo_safe_synthesizer/generation/regex_manager.py
+++ b/src/nemo_safe_synthesizer/generation/regex_manager.py
@@ -360,7 +360,7 @@ def build_json_based_regex(
# Without grouping, the "sequence" is a single record.
sequence_regex = record_regex
- if config.generation.structured_generation_use_single_sequence and config.data.max_sequences_per_example == 1:
+ if config.generation.structured_generation.use_single_sequence and config.data.max_sequences_per_example == 1:
regex = sequence_regex
else:
regex = rf"({sequence_regex}\n)+"
@@ -423,7 +423,7 @@ def build_json_structural_tag(
else:
sequence_format = record_format
- if config.generation.structured_generation_use_single_sequence and config.data.max_sequences_per_example == 1:
+ if config.generation.structured_generation.use_single_sequence and config.data.max_sequences_per_example == 1:
output_format = sequence_format
elif config.data.group_training_examples_by is not None:
output_format = _plus_format(_sequence_format([sequence_format, _const_string_format("\n")]))
diff --git a/src/nemo_safe_synthesizer/generation/vllm_backend.py b/src/nemo_safe_synthesizer/generation/vllm_backend.py
index aca79ec1a..4b82585f9 100644
--- a/src/nemo_safe_synthesizer/generation/vllm_backend.py
+++ b/src/nemo_safe_synthesizer/generation/vllm_backend.py
@@ -290,7 +290,7 @@ def initialize(self, **kwargs) -> None:
# vllm requires this "config" to set the backend ahead of time.
structured_outputs_config = StructuredOutputsConfig(
- backend=self.config.generation.structured_generation_backend,
+ backend=self.config.generation.structured_generation.backend,
)
model_ref = ModelRef.parse(self.config.training.pretrained_model)
@@ -344,13 +344,13 @@ def _build_structured_output_params(self) -> StructuredOutputsParams | None:
Returns:
StructuredOutputsParams if structured generation is enabled, None otherwise.
"""
- if not self.config.generation.use_structured_generation:
+ if not self.config.generation.structured_generation.enabled:
return None
params: dict[str, Any] = {}
schema_method = resolve_structured_generation_schema_method(
- self.config.generation.structured_generation_schema_method,
- self.config.generation.structured_generation_backend,
+ self.config.generation.structured_generation.schema_method,
+ self.config.generation.structured_generation.backend,
)
if schema_method == "regex":
@@ -366,7 +366,7 @@ def _build_structured_output_params(self) -> StructuredOutputsParams | None:
elif schema_method == "json_schema":
params["json"] = self.schema
elif schema_method == "structural_tag":
- backend = self.config.generation.structured_generation_backend
+ backend = self.config.generation.structured_generation.backend
if message := structural_tag_backend_error_message(backend):
raise ParameterError(message)
logger.info("Structured generation is enabled, using an XGrammar Structural Tag")
diff --git a/src/nemo_safe_synthesizer/sdk/config_builder.py b/src/nemo_safe_synthesizer/sdk/config_builder.py
index fed86bbdb..8f5cd668f 100644
--- a/src/nemo_safe_synthesizer/sdk/config_builder.py
+++ b/src/nemo_safe_synthesizer/sdk/config_builder.py
@@ -6,7 +6,7 @@
from __future__ import annotations
from collections.abc import Mapping
-from typing import Any, Self, TypeAlias, TypeVar, cast
+from typing import Any, Self, TypeAlias, TypeVar
import pandas as pd
from pydantic import BaseModel
@@ -23,6 +23,7 @@
)
from ..observability import get_logger
from ..telemetry import _telemetry_enabled
+from ..utils import merge_dicts
logger = get_logger(__name__)
@@ -119,11 +120,12 @@ def _resolve_config(self, values: ParamDict | NSSParameters | None, cls: type[Pa
overrides = kwargs
match values:
case BaseModel() as model:
- return cast(ParamT, model.model_copy(update=overrides))
+ data = model.model_dump()
+ return cls.model_validate(merge_dicts(data, overrides))
case dict() as d:
- return cls.model_validate(d).model_copy(update=overrides)
+ return cls.model_validate(merge_dicts(d, overrides))
case None:
- return cls(**overrides)
+ return cls.model_validate(overrides)
case _:
raise TypeError(f"Unsupported config type: {type(values)}")
diff --git a/tests/benchmarks/test_generation_structured_methods.py b/tests/benchmarks/test_generation_structured_methods.py
index 302aeb00a..2be384453 100644
--- a/tests/benchmarks/test_generation_structured_methods.py
+++ b/tests/benchmarks/test_generation_structured_methods.py
@@ -197,15 +197,15 @@ def _build_command(method: GenerationMethod, output_file: Path, root: Path) -> l
str(output_file),
"--generation__num_records",
num_records,
- "--generation__use_structured_generation",
+ "--generation__structured_generation__enabled",
str(method.use_structured_generation).lower(),
]
if method.use_structured_generation:
command.extend(
[
- "--generation__structured_generation_backend",
+ "--generation__structured_generation__backend",
method.backend,
- "--generation__structured_generation_schema_method",
+ "--generation__structured_generation__schema_method",
method.schema_method or "regex",
]
)
diff --git a/tests/cli/test_utils.py b/tests/cli/test_utils.py
index de3fc3d35..8fe75efec 100644
--- a/tests/cli/test_utils.py
+++ b/tests/cli/test_utils.py
@@ -270,7 +270,8 @@ def test_overrides_config_registry_and_cli(
learning_rate: 0.05
batch_size: 8
generation:
- use_structured_generation: true
+ structured_generation:
+ enabled: true
""")
# Overrides from registry for reference (see registry_with_dataset fixture)
@@ -298,7 +299,7 @@ def test_overrides_config_registry_and_cli(
# Only given in config file
assert config.training.num_input_records_to_sample == 150
- assert config.generation.use_structured_generation
+ assert config.generation.structured_generation.enabled
# Only given in registry
assert config.replace_pii is None
# Only given in CLI
@@ -364,7 +365,7 @@ def test_partial_config_preserves_explicit_field_metadata(self, tmp_path: Path):
config = merge_overrides(config_file, {})
assert config.generation.num_records == 77
- assert config.generation.use_structured_generation is False
+ assert config.generation.structured_generation.enabled is False
assert config.model_dump(exclude_unset=True) == {"generation": {"num_records": 77}}
diff --git a/tests/config/conftest.py b/tests/config/conftest.py
index 5f15c9bde..1d6d34117 100644
--- a/tests/config/conftest.py
+++ b/tests/config/conftest.py
@@ -12,6 +12,7 @@
DifferentialPrivacyHyperparams,
GenerateParameters,
SafeSynthesizerParameters,
+ StructuredGenerationParameters,
TrainingHyperparams,
)
@@ -39,6 +40,9 @@ def fixture_simple_safe_synthesizer_parameters() -> SafeSynthesizerParameters:
num_input_records_to_sample=100,
batch_size=10,
),
- generation=GenerateParameters(num_records=1000, use_structured_generation=True),
+ generation=GenerateParameters(
+ num_records=1000,
+ structured_generation=StructuredGenerationParameters(enabled=True),
+ ),
privacy=DifferentialPrivacyHyperparams(dp_enabled=False),
)
diff --git a/tests/config/test_generate.py b/tests/config/test_generate.py
index 3749e4d35..8f8fc5bc2 100644
--- a/tests/config/test_generate.py
+++ b/tests/config/test_generate.py
@@ -8,6 +8,7 @@
from nemo_safe_synthesizer.config.generate import (
GenerateParameters,
+ StructuredGenerationParameters,
resolve_structured_generation_schema_method,
)
from nemo_safe_synthesizer.config.parameters import SafeSynthesizerParameters
@@ -35,38 +36,176 @@ def test_explicit_methods_pass_through(self, method: str) -> None:
@pytest.mark.unit
-class TestGenerateParametersStructuralTagValidation:
+class TestStructuredGenerationParametersStructuralTagValidation:
@staticmethod
- def _generation_kwargs(*, schema_method: str = "structural_tag", backend: str = "xgrammar") -> dict[str, Any]:
+ def _structured_generation_kwargs(
+ *, schema_method: str = "structural_tag", backend: str = "xgrammar"
+ ) -> dict[str, Any]:
return {
- "use_structured_generation": True,
- "structured_generation_schema_method": schema_method,
- "structured_generation_backend": backend,
+ "enabled": True,
+ "schema_method": schema_method,
+ "backend": backend,
}
@pytest.mark.parametrize("backend", ["xgrammar", "auto"])
def test_compatible_backends_validate(self, backend: str) -> None:
- GenerateParameters(**self._generation_kwargs(backend=backend))
+ StructuredGenerationParameters(
+ **self._structured_generation_kwargs(backend=backend),
+ )
def test_incompatible_backend_raises_validation_error(self) -> None:
- with pytest.raises(ValidationError, match="requires `structured_generation_backend`"):
- GenerateParameters(**self._generation_kwargs(backend="outlines"))
+ with pytest.raises(ValidationError, match="requires `backend`"):
+ StructuredGenerationParameters(**self._structured_generation_kwargs(backend="outlines"))
@pytest.mark.parametrize("backend", ["outlines", "guidance", "lm-format-enforcer"])
def test_auto_with_incompatible_backend_validates(self, backend: str) -> None:
- GenerateParameters(**self._generation_kwargs(schema_method="auto", backend=backend))
+ StructuredGenerationParameters(
+ **self._structured_generation_kwargs(schema_method="auto", backend=backend),
+ )
def test_default_schema_method_is_auto(self) -> None:
params = GenerateParameters()
- assert params.structured_generation_schema_method == "auto"
+ assert params.structured_generation.schema_method == "auto"
def test_skipped_when_structured_generation_disabled(self) -> None:
- GenerateParameters(
- use_structured_generation=False,
- structured_generation_schema_method="structural_tag",
- structured_generation_backend="outlines",
+ StructuredGenerationParameters(
+ enabled=False,
+ schema_method="structural_tag",
+ backend="outlines",
+ )
+
+ def test_generate_parameters_accepts_nested_structured_generation(self) -> None:
+ params = GenerateParameters(
+ structured_generation={ # ty: ignore[invalid-argument-type]
+ "enabled": True,
+ "schema_method": "json_schema",
+ "backend": "outlines",
+ }
)
+ assert params.structured_generation.enabled is True
+ assert params.structured_generation.schema_method == "json_schema"
+ assert params.structured_generation.backend == "outlines"
- def test_from_params_rejects_incompatible_backend(self) -> None:
+ def test_generate_parameters_migrates_legacy_flat_keys(self) -> None:
+ params = GenerateParameters.model_validate(
+ {
+ "use_structured_generation": True,
+ "structured_generation_schema_method": "json_schema",
+ "structured_generation_backend": "outlines",
+ }
+ )
+ assert params.structured_generation.enabled is True
+ assert params.structured_generation.schema_method == "json_schema"
+ assert params.structured_generation.backend == "outlines"
+
+ def test_from_params_rejects_incompatible_legacy_backend(self) -> None:
with pytest.raises(ValidationError, match="outlines"):
- SafeSynthesizerParameters.from_params(**self._generation_kwargs(backend="outlines"))
+ SafeSynthesizerParameters.from_params(
+ use_structured_generation=True,
+ structured_generation_schema_method="structural_tag",
+ structured_generation_backend="outlines",
+ )
+
+ def test_from_params_accepts_nested_structured_generation(self) -> None:
+ params = SafeSynthesizerParameters.from_params(
+ generation={
+ "structured_generation": {
+ "enabled": True,
+ "schema_method": "json_schema",
+ "backend": "outlines",
+ }
+ }
+ )
+ assert params.generation.structured_generation.enabled is True
+ assert params.generation.structured_generation.schema_method == "json_schema"
+ assert params.generation.structured_generation.backend == "outlines"
+
+
+@pytest.mark.unit
+class TestMixedInputMigration:
+ """Regression coverage for mixed-input: legacy flat keys coexist with nested structured_generation dict.
+
+ The migration contract: when both forms are present, legacy flat keys win as explicit overrides.
+ Un-overridden nested dict fields are preserved.
+ """
+
+ def test_legacy_backend_overrides_nested_dict_backend(self) -> None:
+ """A legacy flat backend key beats the nested dict backend when both are supplied."""
+ params = GenerateParameters.model_validate(
+ {
+ "structured_generation": {"enabled": True, "backend": "outlines"},
+ "structured_generation_backend": "xgrammar",
+ }
+ )
+ assert params.structured_generation.backend == "xgrammar"
+ assert params.structured_generation.enabled is True # nested value preserved
+
+ def test_legacy_enabled_overrides_nested_dict_enabled(self) -> None:
+ """A legacy use_structured_generation key beats the nested dict enabled field."""
+ params = GenerateParameters.model_validate(
+ {
+ "use_structured_generation": True,
+ "structured_generation": {"enabled": False, "schema_method": "json_schema"},
+ }
+ )
+ assert params.structured_generation.enabled is True
+ assert params.structured_generation.schema_method == "json_schema" # preserved
+
+ def test_legacy_schema_method_overrides_nested_dict_schema_method(self) -> None:
+ """A legacy schema_method key beats the nested dict schema_method field."""
+ params = GenerateParameters.model_validate(
+ {
+ "structured_generation_schema_method": "regex",
+ "structured_generation": {"enabled": True, "schema_method": "json_schema"},
+ }
+ )
+ assert params.structured_generation.schema_method == "regex"
+ assert params.structured_generation.enabled is True # preserved
+
+ def test_from_params_legacy_backend_overrides_generation_section_backend(self) -> None:
+ """from_params: a top-level legacy backend key overrides the backend inside generation={}."""
+ params = SafeSynthesizerParameters.from_params(
+ generation={"structured_generation": {"enabled": True, "backend": "outlines"}},
+ structured_generation_backend="xgrammar",
+ )
+ assert params.generation.structured_generation.backend == "xgrammar"
+ assert params.generation.structured_generation.enabled is True # preserved
+
+ def test_from_params_legacy_enabled_overrides_structured_generation_kwarg(self) -> None:
+ """from_params: a top-level legacy use_structured_generation overrides structured_generation kwarg."""
+ params = SafeSynthesizerParameters.from_params(
+ structured_generation={"enabled": False, "schema_method": "json_schema"},
+ use_structured_generation=True,
+ )
+ assert params.generation.structured_generation.enabled is True
+ assert params.generation.structured_generation.schema_method == "json_schema" # preserved
+
+ def test_nested_dict_with_no_legacy_keys_uses_dict_values(self) -> None:
+ """When no legacy flat keys are present, nested dict values are used as-is."""
+ params = GenerateParameters.model_validate(
+ {
+ "structured_generation": {
+ "enabled": True,
+ "backend": "outlines",
+ "schema_method": "json_schema",
+ }
+ }
+ )
+ assert params.structured_generation.enabled is True
+ assert params.structured_generation.backend == "outlines"
+ assert params.structured_generation.schema_method == "json_schema"
+
+ def test_legacy_keys_with_no_nested_dict_are_migrated(self) -> None:
+ """When only legacy flat keys are present, migration produces the correct nested values."""
+ params = GenerateParameters.model_validate(
+ {
+ "use_structured_generation": True,
+ "structured_generation_backend": "xgrammar",
+ "structured_generation_schema_method": "structural_tag",
+ "structured_generation_use_single_sequence": True,
+ }
+ )
+ assert params.structured_generation.enabled is True
+ assert params.structured_generation.backend == "xgrammar"
+ assert params.structured_generation.schema_method == "structural_tag"
+ assert params.structured_generation.use_single_sequence is True
diff --git a/tests/config/test_parameters.py b/tests/config/test_parameters.py
index 1dda3cd2e..a67da2150 100644
--- a/tests/config/test_parameters.py
+++ b/tests/config/test_parameters.py
@@ -146,7 +146,7 @@ def _saved_config() -> SafeSynthesizerParameters:
saved.training.batch_size = 8
saved.emit_telemetry = False
saved.generation.num_records = 3000
- saved.generation.use_structured_generation = True
+ saved.generation.structured_generation.enabled = True
saved.generation.validation.group_by_ignore_invalid_records = True
return saved
@@ -175,7 +175,7 @@ class TestWithRuntimeOverrides:
SafeSynthesizerParameters,
{
"generation.num_records": 3000,
- "generation.use_structured_generation": True,
+ "generation.structured_generation.enabled": True,
"generation.validation.group_by_ignore_invalid_records": True,
"training.batch_size": 8,
"emit_telemetry": False,
@@ -186,7 +186,7 @@ class TestWithRuntimeOverrides:
_runtime_num_records,
{
"generation.num_records": 100,
- "generation.use_structured_generation": True, # unset runtime field preserved
+ "generation.structured_generation.enabled": True, # unset runtime field preserved
"training.batch_size": 8, # non-overridable section inherited
},
id="top-level-generation-override",
@@ -240,10 +240,10 @@ def test_returned_config_is_independent_of_saved(self):
# Inherited section, overridden section, and an unchanged section.
merged.data.holdout = 0.42
merged.training.batch_size = 99
- merged.generation.use_structured_generation = False
+ merged.generation.structured_generation.enabled = False
assert merged.data is not saved.data
assert merged.training is not saved.training
assert saved.data.holdout != 0.42
assert saved.training.batch_size == 8
- assert saved.generation.use_structured_generation is True
+ assert saved.generation.structured_generation.enabled is True
diff --git a/tests/configurator/test_pydantic_click_options.py b/tests/configurator/test_pydantic_click_options.py
index 10b8162e8..78ebe98e0 100644
--- a/tests/configurator/test_pydantic_click_options.py
+++ b/tests/configurator/test_pydantic_click_options.py
@@ -478,3 +478,55 @@ def cmd(**kwargs):
result = CliRunner().invoke(cmd, ["--replace_pii__globals__seed", "42"])
assert result.exit_code == 0, result.output
assert captured["replace_pii"]["globals"]["seed"] == 42
+
+
+def test_structured_generation_nested_option_end_to_end_via_click_runner():
+ """The canonical nested structured-generation option flows through Click parsing."""
+ captured: dict = {}
+
+ @pydantic_options(SafeSynthesizerParameters, field_separator="__")
+ @click.command()
+ def cmd(**kwargs):
+ captured.update(parse_overrides(kwargs))
+
+ result = CliRunner().invoke(cmd, ["--generation__structured_generation__backend", "outlines"])
+ assert result.exit_code == 0, result.output
+ assert captured["generation"]["structured_generation"]["backend"] == "outlines"
+
+
+@pytest.mark.parametrize(
+ ("option", "legacy_key", "value", "expected"),
+ [
+ ("--generation__use_structured_generation", "use_structured_generation", "true", True),
+ ("--generation__structured_generation_backend", "structured_generation_backend", "outlines", "outlines"),
+ (
+ "--generation__structured_generation_schema_method",
+ "structured_generation_schema_method",
+ "json_schema",
+ "json_schema",
+ ),
+ (
+ "--generation__structured_generation_use_single_sequence",
+ "structured_generation_use_single_sequence",
+ "true",
+ True,
+ ),
+ ],
+)
+def test_structured_generation_legacy_options_end_to_end_via_click_runner(
+ option: str,
+ legacy_key: str,
+ value: str,
+ expected: object,
+):
+ """Legacy flat structured-generation CLI aliases remain accepted during migration."""
+ captured: dict = {}
+
+ @pydantic_options(SafeSynthesizerParameters, field_separator="__")
+ @click.command()
+ def cmd(**kwargs):
+ captured.update(parse_overrides(kwargs))
+
+ result = CliRunner().invoke(cmd, [option, value])
+ assert result.exit_code == 0, result.output
+ assert captured["generation"][legacy_key] == expected
diff --git a/tests/conftest.py b/tests/conftest.py
index 94072b75e..ee89adfe4 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -109,9 +109,10 @@ def fixture_yaml_config_str() -> str:
num_records: 1000
patience: 1
repetition_penalty: 1.0
+ structured_generation:
+ enabled: false
temperature: 0.9
top_p: 1.0
- use_structured_generation: false
privacy:
delta: auto
dp_enabled: false
diff --git a/tests/e2e/required_configs/mistral-dp.yaml b/tests/e2e/required_configs/mistral-dp.yaml
index 1660a6f38..10f9ac660 100644
--- a/tests/e2e/required_configs/mistral-dp.yaml
+++ b/tests/e2e/required_configs/mistral-dp.yaml
@@ -5,8 +5,9 @@ data:
random_state: 187260
generation:
num_records: 500
- use_structured_generation: true
- structured_generation_backend: outlines
+ structured_generation:
+ enabled: true
+ backend: outlines
privacy:
dp_enabled: true
training:
diff --git a/tests/e2e/required_configs/smollm3-dp.yaml b/tests/e2e/required_configs/smollm3-dp.yaml
index c2589e723..3ee8b2764 100644
--- a/tests/e2e/required_configs/smollm3-dp.yaml
+++ b/tests/e2e/required_configs/smollm3-dp.yaml
@@ -5,7 +5,8 @@ data:
random_state: 18726
generation:
num_records: 500
- use_structured_generation: true # somehow the auto backend works better than outlines for smollm3 on clinc_oos
+ structured_generation:
+ enabled: true # somehow the auto backend works better than outlines for smollm3 on clinc_oos
privacy:
dp_enabled: true
training:
diff --git a/tests/e2e/required_configs/tinyllama-dp.yaml b/tests/e2e/required_configs/tinyllama-dp.yaml
index ed52fefb2..e21e435d1 100644
--- a/tests/e2e/required_configs/tinyllama-dp.yaml
+++ b/tests/e2e/required_configs/tinyllama-dp.yaml
@@ -5,8 +5,9 @@ data:
random_state: 329896
generation:
num_records: 500
- use_structured_generation: true
- structured_generation_backend: outlines
+ structured_generation:
+ enabled: true
+ backend: outlines
privacy:
dp_enabled: true
training:
diff --git a/tests/e2e/test_dataset_config.py b/tests/e2e/test_dataset_config.py
index 5d2132344..d6cb2d322 100644
--- a/tests/e2e/test_dataset_config.py
+++ b/tests/e2e/test_dataset_config.py
@@ -144,7 +144,7 @@ def test_dow_jones_index_dataset(
config = update_group_by_config(config, "stock", "date")
config = AutoConfigResolver(df, config)()
if config_file == "smollm3-dp.yaml":
- config.generation.structured_generation_backend = "outlines"
+ config.generation.structured_generation.backend = "outlines"
logger.info(f"Running test with config: {config}")
nss = SafeSynthesizer(config=config).with_data_source(df)
diff --git a/tests/e2e/test_safe_synthesizer.py b/tests/e2e/test_safe_synthesizer.py
index 336ef112f..16cf909c1 100644
--- a/tests/e2e/test_safe_synthesizer.py
+++ b/tests/e2e/test_safe_synthesizer.py
@@ -66,7 +66,9 @@ def test_train_and_generate_dp(fixture_financial_transactions_dataset, fixture_s
dp_enabled=True,
epsilon=100.0,
num_records=100,
- use_structured_generation=True,
+ structured_generation={
+ "enabled": True,
+ },
)
logger.info(f"Running DP test with config: {config}")
diff --git a/tests/generation/test_regex_manager.py b/tests/generation/test_regex_manager.py
index 27eb1bd06..06d9c3fd9 100644
--- a/tests/generation/test_regex_manager.py
+++ b/tests/generation/test_regex_manager.py
@@ -26,7 +26,7 @@ def fixture_safe_synthesizer_config():
return SafeSynthesizerParameters.from_params(
group_training_examples_by=None,
max_sequences_per_example=None,
- structured_generation_use_single_sequence=False,
+ structured_generation={"use_single_sequence": False},
)
@@ -107,7 +107,7 @@ def test_build_json_structural_tag_uses_schema_constrained_jsonl(fixture_safe_sy
def test_build_json_structural_tag_with_single_sequence(fixture_safe_synthesizer_config):
"""Single-sequence Structural Tag emits exactly one schema-constrained object."""
fixture_safe_synthesizer_config.data.max_sequences_per_example = 1
- fixture_safe_synthesizer_config.generation.structured_generation_use_single_sequence = True
+ fixture_safe_synthesizer_config.generation.structured_generation.use_single_sequence = True
schema = {"type": "object", "properties": {"name": {"type": "string"}}}
structural_tag = json.loads(
@@ -168,7 +168,7 @@ def test_build_json_structural_tag_with_groupby_single_sequence(fixture_safe_syn
"""Grouped single-sequence Structural Tag emits exactly one BOS/EOS-wrapped sequence."""
fixture_safe_synthesizer_config.data.group_training_examples_by = "id"
fixture_safe_synthesizer_config.data.max_sequences_per_example = 1
- fixture_safe_synthesizer_config.generation.structured_generation_use_single_sequence = True
+ fixture_safe_synthesizer_config.generation.structured_generation.use_single_sequence = True
schema = {"type": "object", "properties": {"name": {"type": "string"}}}
structural_tag = json.loads(
@@ -563,7 +563,7 @@ def test_property_regex(fixture_safe_synthesizer_config):
[
({}, '{"name":"alice"}\n{"name":"bob"}\n'),
(
- {"max_sequences_per_example": 1, "structured_generation_use_single_sequence": True},
+ {"max_sequences_per_example": 1, "structured_generation.use_single_sequence": True},
'{"name":"alice"}',
),
(
@@ -574,7 +574,7 @@ def test_property_regex(fixture_safe_synthesizer_config):
{
"group_training_examples_by": "id",
"max_sequences_per_example": 1,
- "structured_generation_use_single_sequence": True,
+ "structured_generation.use_single_sequence": True,
},
'{"name":"alice"}\n',
),
@@ -592,8 +592,8 @@ def test_structural_tag_accepts_training_jsonl_shapes(
fixture_safe_synthesizer_config.data.group_training_examples_by = value
elif key == "max_sequences_per_example":
fixture_safe_synthesizer_config.data.max_sequences_per_example = value
- elif key == "structured_generation_use_single_sequence":
- fixture_safe_synthesizer_config.generation.structured_generation_use_single_sequence = value
+ elif key == "structured_generation.use_single_sequence":
+ fixture_safe_synthesizer_config.generation.structured_generation.use_single_sequence = value
schema = {
"type": "object",
@@ -670,10 +670,10 @@ def test_round_trip_doc_summaries_schema_regex_matches(fixture_doc_summaries_dat
# Purpose: Validates single-sequence mode without grouping produces a bare record regex (no \n, no repetition).
-# Data: Simple enum schema; structured_generation_use_single_sequence=True, max_sequences_per_example=1, no grouping.
+# Data: Simple enum schema; structured_generation.use_single_sequence=True, max_sequences_per_example=1, no grouping.
# Asserts: Regex is the raw record pattern; matches a single record without trailing newline; rejects newline-terminated and multi-record inputs.
def test_single_sequence_no_grouping(fixture_safe_synthesizer_config):
- fixture_safe_synthesizer_config.generation.structured_generation_use_single_sequence = True
+ fixture_safe_synthesizer_config.generation.structured_generation.use_single_sequence = True
fixture_safe_synthesizer_config.data.max_sequences_per_example = 1
schema = {
@@ -695,10 +695,10 @@ def test_single_sequence_no_grouping(fixture_safe_synthesizer_config):
# Purpose: Validates single-sequence mode with grouping produces a single BOS/EOS group (no trailing \n, no group repetition).
-# Data: Simple enum schema; structured_generation_use_single_sequence=True, max_sequences_per_example=1, group_by="id".
+# Data: Simple enum schema; structured_generation.use_single_sequence=True, max_sequences_per_example=1, group_by="id".
# Asserts: Regex matches exactly one group; rejects trailing newline after EOS and multiple groups.
def test_single_sequence_with_groupby(fixture_safe_synthesizer_config):
- fixture_safe_synthesizer_config.generation.structured_generation_use_single_sequence = True
+ fixture_safe_synthesizer_config.generation.structured_generation.use_single_sequence = True
fixture_safe_synthesizer_config.data.max_sequences_per_example = 1
fixture_safe_synthesizer_config.data.group_training_examples_by = "id"
@@ -721,12 +721,12 @@ def test_single_sequence_with_groupby(fixture_safe_synthesizer_config):
assert re.fullmatch(regex, '{"foo":"a"}\n\n{"foo":"b"}\n\n') is None
-# Purpose: Validates that structured_generation_use_single_sequence is ignored when max_sequences_per_example != 1.
+# Purpose: Validates that structured_generation.use_single_sequence is ignored when max_sequences_per_example != 1.
# Data: Simple enum schema; flag=True but max_sequences_per_example is None or 2.
# Asserts: Regex uses the repeated (...\n)+ form despite the flag being set.
@pytest.mark.parametrize("max_seq", [None, 2])
def test_single_sequence_flag_ignored_when_max_sequences_not_one(max_seq, fixture_safe_synthesizer_config):
- fixture_safe_synthesizer_config.generation.structured_generation_use_single_sequence = True
+ fixture_safe_synthesizer_config.generation.structured_generation.use_single_sequence = True
fixture_safe_synthesizer_config.data.max_sequences_per_example = max_seq
schema = {
diff --git a/tests/generation/test_timeseries_backend.py b/tests/generation/test_timeseries_backend.py
index c46a019f3..f9a65415b 100644
--- a/tests/generation/test_timeseries_backend.py
+++ b/tests/generation/test_timeseries_backend.py
@@ -94,7 +94,6 @@ def timeseries_base_params():
),
generation=GenerateParameters(
num_records=100,
- use_structured_generation=False,
),
time_series=TimeSeriesParameters(
is_timeseries=True,
@@ -128,7 +127,6 @@ def timeseries_elapsed_params():
),
generation=GenerateParameters(
num_records=100,
- use_structured_generation=False,
),
time_series=TimeSeriesParameters(
is_timeseries=True,
diff --git a/tests/generation/test_vllm_backend.py b/tests/generation/test_vllm_backend.py
index e00c74f9a..a76fa6d23 100644
--- a/tests/generation/test_vllm_backend.py
+++ b/tests/generation/test_vllm_backend.py
@@ -93,7 +93,6 @@ def base_params():
),
generation=GenerateParameters(
num_records=100,
- use_structured_generation=False,
),
)
@@ -101,35 +100,35 @@ def base_params():
@pytest.fixture
def params_with_structured_generation_auto(base_params):
"""Create params with structured generation enabled using auto schema method."""
- base_params.generation.use_structured_generation = True
- base_params.generation.structured_generation_schema_method = "auto"
+ base_params.generation.structured_generation.enabled = True
+ base_params.generation.structured_generation.schema_method = "auto"
return base_params
@pytest.fixture
def params_with_structured_generation_regex(base_params):
"""Create params with structured generation enabled using regex."""
- base_params.generation.use_structured_generation = True
- base_params.generation.structured_generation_schema_method = "regex"
- base_params.generation.structured_generation_backend = "xgrammar"
+ base_params.generation.structured_generation.enabled = True
+ base_params.generation.structured_generation.schema_method = "regex"
+ base_params.generation.structured_generation.backend = "xgrammar"
return base_params
@pytest.fixture
def params_with_structured_generation_json(base_params):
"""Create params with structured generation enabled using json_schema."""
- base_params.generation.use_structured_generation = True
- base_params.generation.structured_generation_schema_method = "json_schema"
- base_params.generation.structured_generation_backend = "xgrammar"
+ base_params.generation.structured_generation.enabled = True
+ base_params.generation.structured_generation.schema_method = "json_schema"
+ base_params.generation.structured_generation.backend = "xgrammar"
return base_params
@pytest.fixture
def params_with_structured_generation_structural_tag(base_params):
"""Create params with structured generation enabled using structural_tag."""
- base_params.generation.use_structured_generation = True
- base_params.generation.structured_generation_schema_method = "structural_tag"
- base_params.generation.structured_generation_backend = "xgrammar"
+ base_params.generation.structured_generation.enabled = True
+ base_params.generation.structured_generation.schema_method = "structural_tag"
+ base_params.generation.structured_generation.backend = "xgrammar"
return base_params
@@ -266,7 +265,7 @@ def test_auto_resolves_to_structural_tag_on_xgrammar_backends(
backend,
):
"""Auto schema method uses structural_tag on xgrammar-capable backends."""
- params_with_structured_generation_auto.generation.structured_generation_backend = backend
+ params_with_structured_generation_auto.generation.structured_generation.backend = backend
backend_instance = create_backend(
params_with_structured_generation_auto,
mock_model_metadata,
@@ -298,7 +297,7 @@ def test_auto_resolves_to_regex_on_other_backends(
backend,
):
"""Auto schema method falls back to regex on non-xgrammar backends."""
- params_with_structured_generation_auto.generation.structured_generation_backend = backend
+ params_with_structured_generation_auto.generation.structured_generation.backend = backend
backend_instance = create_backend(
params_with_structured_generation_auto,
mock_model_metadata,
@@ -330,7 +329,7 @@ def test_structural_tag_rejects_non_xgrammar_backends(
backend,
):
"""Structural Tag requires vLLM's xgrammar backend."""
- params_with_structured_generation_structural_tag.generation.structured_generation_backend = backend
+ params_with_structured_generation_structural_tag.generation.structured_generation.backend = backend
backend_instance = create_backend(
params_with_structured_generation_structural_tag,
mock_model_metadata,
@@ -338,7 +337,7 @@ def test_structural_tag_rejects_non_xgrammar_backends(
mock_workdir,
)
- with pytest.raises(ParameterError, match="requires `structured_generation_backend`"):
+ with pytest.raises(ParameterError, match="requires `backend`"):
backend_instance._build_structured_output_params()
def test_config_with_grouping_passed_to_build_regex(
diff --git a/tests/generation/test_vllm_shutdown.py b/tests/generation/test_vllm_shutdown.py
index 47da08868..9a9d40c5a 100644
--- a/tests/generation/test_vllm_shutdown.py
+++ b/tests/generation/test_vllm_shutdown.py
@@ -35,7 +35,7 @@ def backend(_mock_vllm_cleanup, fixture_session_cache_dir):
mock_config.time_series.is_timeseries = False
mock_config.data.group_training_examples_by = None
# Pin to a valid literal so StructuredOutputsConfig Pydantic validation passes in initialize().
- mock_config.generation.structured_generation_backend = "xgrammar"
+ mock_config.generation.structured_generation.backend = "xgrammar"
mock_config.generation.attention_backend = None
mock_workdir = MagicMock()
diff --git a/tests/sdk/test_process_data.py b/tests/sdk/test_process_data.py
index 3c6bde491..f9d1c891f 100644
--- a/tests/sdk/test_process_data.py
+++ b/tests/sdk/test_process_data.py
@@ -456,14 +456,14 @@ def test_load_applies_runtime_generation_and_evaluation_config(
saved_config = SafeSynthesizerParameters()
saved_config.training.batch_size = 8
saved_config.generation.num_records = 3000
- saved_config.generation.use_structured_generation = True
- saved_config.generation.structured_generation_schema_method = "auto"
+ saved_config.generation.structured_generation.enabled = True
+ saved_config.generation.structured_generation.schema_method = "auto"
workdir.config.write_text(saved_config.model_dump_json())
runtime_config = SafeSynthesizerParameters()
runtime_config.training.batch_size = 32
runtime_config.generation.num_records = 100
- runtime_config.generation.structured_generation_schema_method = "auto"
+ runtime_config.generation.structured_generation.schema_method = "auto"
runtime_config.evaluation.enabled = False
mock_metadata_cls.from_metadata_json.return_value = MagicMock()
@@ -474,8 +474,8 @@ def test_load_applies_runtime_generation_and_evaluation_config(
assert builder._nss_config.training.batch_size == 8
assert builder._nss_config.generation.num_records == 100
# Saved value not re-specified at runtime is preserved (field-level merge).
- assert builder._nss_config.generation.use_structured_generation is True
- assert builder._nss_config.generation.structured_generation_schema_method == "auto"
+ assert builder._nss_config.generation.structured_generation.enabled is True
+ assert builder._nss_config.generation.structured_generation.schema_method == "auto"
assert builder._nss_config.evaluation.enabled is False
@patch("nemo_safe_synthesizer.sdk.library_builder.ModelMetadata")
@@ -494,8 +494,8 @@ def test_load_preserves_saved_generation_when_runtime_config_has_no_overrides(
)
saved_config = SafeSynthesizerParameters()
saved_config.generation.num_records = 3000
- saved_config.generation.use_structured_generation = True
- saved_config.generation.structured_generation_schema_method = "auto"
+ saved_config.generation.structured_generation.enabled = True
+ saved_config.generation.structured_generation.schema_method = "auto"
saved_config.evaluation.enabled = True
workdir.config.write_text(saved_config.model_dump_json())
@@ -506,8 +506,8 @@ def test_load_preserves_saved_generation_when_runtime_config_has_no_overrides(
assert builder._nss_config is not None
assert builder._nss_config.generation.num_records == 3000
- assert builder._nss_config.generation.use_structured_generation is True
- assert builder._nss_config.generation.structured_generation_schema_method == "auto"
+ assert builder._nss_config.generation.structured_generation.enabled is True
+ assert builder._nss_config.generation.structured_generation.schema_method == "auto"
assert builder._nss_config.evaluation.enabled is True
@patch("nemo_safe_synthesizer.sdk.library_builder.ModelMetadata")
@@ -561,8 +561,8 @@ def test_load_full_runtime_config_replaces_supported_runtime_sections(
saved_config = SafeSynthesizerParameters()
saved_config.training.batch_size = 8
saved_config.generation.num_records = 3000
- saved_config.generation.use_structured_generation = True
- saved_config.generation.structured_generation_schema_method = "auto"
+ saved_config.generation.structured_generation.enabled = True
+ saved_config.generation.structured_generation.schema_method = "auto"
saved_config.evaluation.enabled = True
workdir.config.write_text(saved_config.model_dump_json())
@@ -579,8 +579,8 @@ def test_load_full_runtime_config_replaces_supported_runtime_sections(
assert builder._nss_config is not None
assert builder._nss_config.training.batch_size == 8
assert builder._nss_config.generation.num_records == 100
- assert builder._nss_config.generation.use_structured_generation is False
- assert builder._nss_config.generation.structured_generation_schema_method == "auto"
+ assert builder._nss_config.generation.structured_generation.enabled is False
+ assert builder._nss_config.generation.structured_generation.schema_method == "auto"
assert builder._nss_config.evaluation.enabled is False
@patch("nemo_safe_synthesizer.sdk.library_builder.ModelMetadata")
diff --git a/tests/smoke/test_nss_structured_gen_gpu.py b/tests/smoke/test_nss_structured_gen_gpu.py
index fae8ff996..8a007c4c3 100644
--- a/tests/smoke/test_nss_structured_gen_gpu.py
+++ b/tests/smoke/test_nss_structured_gen_gpu.py
@@ -35,9 +35,11 @@ def test_nss_structured_generation(fixture_local_tinyllama_dir, fixture_gpu_smok
lora_r=8,
holdout=0,
max_holdout=0,
- use_structured_generation=True,
- structured_generation_backend="outlines",
- structured_generation_schema_method="json_schema",
+ structured_generation={
+ "enabled": True,
+ "backend": "outlines",
+ "schema_method": "json_schema",
+ },
)
nss = SafeSynthesizer(config=config, save_path=tmp_path)
nss.with_data_source(fixture_gpu_smoke_df).process_data().train()