fix: upgrade xgrammar and use new structured gen method

docs/user-guide/configuration.md

@@ -160,7 +160,7 @@ for the full API reference.
 | `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` | `"regex"` | Schema method (`"regex"` or `"json_schema"`) | Leave at `"regex"` |
+| `generation.structured_generation_schema_method` | `"structural_tag"` | Schema method (`"structural_tag"`, `"json_schema"`, or `"regex"`) | Leave at `"structural_tag"` for XGrammar schema-constrained JSONL constraints |
 | `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"` |

docs/user-guide/running.md

@@ -883,9 +883,10 @@ records. Use it when the pipeline struggles to produce valid records.
     ```yaml
     generation:
       use_structured_generation: true
-      structured_generation_schema_method: "regex"
+      structured_generation_schema_method: "structural_tag"
     ```
 
+- `"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.
 

pyproject.toml

@@ -137,6 +137,7 @@ cpu = [
   "triton>=2.0.0; sys_platform=='linux'",
   "trl>=0.23.0",
   "vllm==0.20.0; sys_platform=='linux'",
+  "xgrammar>=0.2.0; sys_platform=='linux'",
 ]
 
 cu129 = [
@@ -161,6 +162,7 @@ cu129 = [
   "triton>=2.0.0; sys_platform == 'linux'",
   "trl>=0.23.0",
   "vllm==0.20.0+cu129; sys_platform == 'linux'",
+  "xgrammar>=0.2.0; sys_platform == 'linux'",
 ]
 
 # at some point, do per-subpackage dependencies

src/nemo_safe_synthesizer/cli/run.py

@@ -568,7 +568,7 @@ def run_generate(
 
         try:
             nss = (
-                nss.load_from_save_path()
+                nss.load_from_save_path(runtime_config=config)
                 .process_data()
                 .generate()
                 .evaluate()

src/nemo_safe_synthesizer/config/generate.py

@@ -147,16 +147,17 @@ class GenerateParameters(Parameters, BaseModel):
     ] = "auto"
 
     structured_generation_schema_method: Annotated[
-        Literal["regex", "json_schema"],
+        Literal["regex", "json_schema", "structural_tag"],
         Field(
             title="structured_generation_schema_method",
             description=(
                 "The method used to generate the schema from your dataset and pass it to the generation backend. "
                 "'regex' uses a custom regex construction method that tends to be more comprehensive "
-                "than 'json_schema' at the cost of speed."
+                "than 'json_schema' at the cost of speed. 'structural_tag' uses XGrammar Structural Tag "
+                "to compose schema-constrained JSONL output."
             ),
         ),
-    ] = "regex"
+    ] = "structural_tag"
 
     structured_generation_use_single_sequence: Annotated[
         bool,

src/nemo_safe_synthesizer/generation/regex_manager.py

@@ -366,3 +366,74 @@ def build_json_based_regex(
         regex = rf"({sequence_regex}\n)+"
 
     return regex
+
+
+def _const_string_format(value: str) -> dict[str, str]:
+    """Return a Structural Tag constant-string format."""
+    return {"type": "const_string", "value": value}
+
+
+def _sequence_format(elements: list[dict[str, Any]]) -> dict[str, Any]:
+    """Return a Structural Tag sequence format."""
+    return {"type": "sequence", "elements": elements}
+
+
+def _plus_format(content: dict[str, Any]) -> dict[str, Any]:
+    """Return a Structural Tag one-or-more repetition format."""
+    return {"type": "plus", "content": content}
+
+
+def build_json_structural_tag(
+    schema: dict[str, Any],
+    config: SafeSynthesizerParameters,
+    bos_token: str,
+    eos_token: str,
+) -> str:
+    """Build an XGrammar Structural Tag for schema-constrained JSONL records.
+
+    The raw vLLM ``json`` constraint describes a single JSON value. Structural
+    Tag lets NSS describe the larger generation shape directly: one or more
+    schema-constrained JSON records separated by newlines, optionally wrapped
+    in BOS/EOS group delimiters.
+
+    Args:
+        schema: JSON schema dictionary describing one record.
+        config: Pipeline configuration (used for grouping and
+            structured-generation settings).
+        bos_token: Beginning-of-sequence token (used when grouping).
+        eos_token: End-of-sequence token (used when grouping).
+
+    Returns:
+        JSON string suitable for ``StructuredOutputsParams(structural_tag=...)``.
+    """
+    record_format: dict[str, Any] = {
+        "type": "json_schema",
+        "json_schema": schema,
+    }
+    record_line_format = _sequence_format([record_format, _const_string_format("\n")])
+
+    if config.data.group_training_examples_by is not None:
+        sequence_format = _sequence_format(
+            [
+                _const_string_format(bos_token),
+                _plus_format(record_line_format),
+                _const_string_format(eos_token),
+            ]
+        )
+    else:
+        sequence_format = record_format
+
+    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")]))
+    else:
+        output_format = _plus_format(record_line_format)
+
+    return json.dumps(
+        {
+            "type": "structural_tag",
+            "format": output_format,
+        },
+        ensure_ascii=True,
+    )

src/nemo_safe_synthesizer/generation/vllm_backend.py

@@ -26,11 +26,11 @@
 from ..cli.artifact_structure import Workdir
 from ..config import SafeSynthesizerParameters
 from ..defaults import DEFAULT_SAMPLING_PARAMETERS, FIXED_RUNTIME_GENERATE_ARGS
-from ..errors import InternalError
+from ..errors import InternalError, ParameterError
 from ..generation.backend import GeneratorBackend
 from ..generation.batch import Batch
 from ..generation.processors import Processor, TabularDataProcessor, create_processor
-from ..generation.regex_manager import build_json_based_regex
+from ..generation.regex_manager import build_json_based_regex, build_json_structural_tag
 from ..generation.results import GenerateJobResults, GenerationBatches, GenerationStatus
 from ..llm.metadata import ModelMetadata
 from ..llm.utils import ModelRef, cleanup_memory, get_max_vram
@@ -337,6 +337,22 @@ def _build_structured_output_params(self) -> StructuredOutputsParams | None:
             params["regex"] = regex
         elif self.config.generation.structured_generation_schema_method == "json_schema":
             params["json"] = self.schema
+        elif self.config.generation.structured_generation_schema_method == "structural_tag":
+            backend = self.config.generation.structured_generation_backend
+            if backend not in {"auto", "xgrammar"}:
+                raise ParameterError(
+                    "Invalid structured generation configuration: "
+                    "`structured_generation_schema_method='structural_tag'` requires "
+                    f"`structured_generation_backend` to be 'xgrammar' or 'auto', got {backend!r}."
+                )
+            logger.info("Structured generation is enabled, using an XGrammar Structural Tag")
+            pc = self.model_metadata.prompt_config
+            params["structural_tag"] = build_json_structural_tag(
+                self.schema,
+                self.config,
+                bos_token=pc.bos_token,
+                eos_token=pc.eos_token,
+            )
 
         return StructuredOutputsParams(**params)
 

src/nemo_safe_synthesizer/sdk/library_builder.py

@@ -237,12 +237,15 @@ def _ensure_observability(self) -> None:
         initialize_observability()
 
     @traced("SafeSynthesizer.load_from_save_path", category=LogCategory.RUNTIME)
-    def load_from_save_path(self) -> SafeSynthesizer:
+    def load_from_save_path(self, runtime_config: SafeSynthesizerParameters | None = None) -> SafeSynthesizer:
         """Load the Safe Synthesizer configuration from the save path.
 
         Loads the configuration from the source run directory's config file.
         When resuming from a trained model for generation, the source paths
         point to the parent workdir that contains the trained adapter.
+        Optional ``runtime_config`` values for generation and evaluation are
+        applied after loading the saved training-run config so resume-time CLI
+        overrides work without mutating the persisted train config.
 
         Always prefers cached train/test splits from the training run to ensure
         evaluation metrics are consistent and privacy guarantees are maintained.
@@ -256,7 +259,19 @@ def load_from_save_path(self) -> SafeSynthesizer:
         # Use source paths which point to parent workdir when resuming for generation
         config_file = self._workdir.source_config
 
-        self._nss_config = SafeSynthesizerParameters.from_json(config_file)
+        saved_config = SafeSynthesizerParameters.from_json(config_file)
+        if runtime_config is not None:
+            saved_config = saved_config.model_copy(
+                update={
+                    "generation": runtime_config.generation,
+                    "evaluation": runtime_config.evaluation,
+                    "emit_telemetry": runtime_config.emit_telemetry,
+                },
+            )
+        self._nss_config = saved_config
+        self._generation_config = self._nss_config.generation
+        self._evaluation_config = self._nss_config.evaluation
+        self._emit_telemetry_config = self._nss_config.emit_telemetry
 
         # Load model metadata from saved file (contains initial_prefill for timeseries)
         # rather than creating new metadata from config