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()