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/config/generate.py

@@ -3,11 +3,12 @@
 
 from __future__ import annotations
 
-from typing import Annotated, Literal
+from typing import Annotated, Literal, Self
 
 from pydantic import (
     BaseModel,
     Field,
+    model_validator,
 )
 
 from ..configurator.parameters import (
@@ -17,8 +18,32 @@
     ValueValidator,
     range_validator,
 )
+from ..errors import ParameterError
 
-__all__ = ["GenerateParameters", "ValidationParameters"]
+STRUCTURAL_TAG_COMPATIBLE_BACKENDS = frozenset({"auto", "xgrammar"})
+
+__all__ = [
+    "GenerateParameters",
+    "STRUCTURAL_TAG_COMPATIBLE_BACKENDS",
+    "ValidationParameters",
+    "structural_tag_backend_error_message",
+]
+
+
+def structural_tag_backend_error_message(backend: str) -> str | None:
+    """Return an error message when *backend* cannot serve ``structural_tag``.
+
+    vLLM only supports XGrammar Structural Tag constraints when the guided
+    decoding backend is ``xgrammar`` or ``auto`` (which selects xgrammar for
+    this schema method).
+    """
+    if backend in STRUCTURAL_TAG_COMPATIBLE_BACKENDS:
+        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}."
+    )
 
 
 class ValidationParameters(Parameters, BaseModel):
@@ -147,16 +172,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,
@@ -191,3 +217,13 @@ 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

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

@@ -25,12 +25,13 @@
 from .. import utils
 from ..cli.artifact_structure import Workdir
 from ..config import SafeSynthesizerParameters
+from ..config.generate import structural_tag_backend_error_message
 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 +338,18 @@ 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 message := structural_tag_backend_error_message(backend):
+                raise ParameterError(message)
+            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/preflight/__init__.py

@@ -24,6 +24,7 @@
     OversamplingCheck,
     PseudoColumnCheck,
     SmallDatasetCheck,
+    StructuralTagBackendCheck,
     TimestampColumnCheck,
     TokenBudgetCheck,
     VRAMHeadroomCheck,
@@ -77,6 +78,7 @@
     "PreflightStatus",
     "PreflightStage",
     "PseudoColumnCheck",
+    "StructuralTagBackendCheck",
     "TimestampColumnCheck",
     "TokenBudgetCheck",
     "VRAMHeadroomCheck",

src/nemo_safe_synthesizer/preflight/checks/__init__.py

@@ -30,6 +30,7 @@
     InferenceKeyCheck,
     VRAMHeadroomCheck,
 )
+from .generation import StructuralTagBackendCheck
 from .metadata import TokenBudgetCheck
 
 __all__ = [
@@ -43,6 +44,7 @@
     "OrderbyColumnCheck",
     "OversamplingCheck",
     "PseudoColumnCheck",
+    "StructuralTagBackendCheck",
     "TimestampColumnCheck",
     "TokenBudgetCheck",
     "VRAMHeadroomCheck",
@@ -59,6 +61,7 @@
     CUDAAvailabilityCheck(),
     InferenceKeyCheck(),
     HFModelAvailabilityCheck(),
+    StructuralTagBackendCheck(),
     # DATAFRAME
     DatasetSizeCheck(),
     GroupbyColumnCheck(),

src/nemo_safe_synthesizer/preflight/checks/generation.py

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Generation-config checks that should fail before training starts."""
+
+from __future__ import annotations
+
+from ...config.generate import structural_tag_backend_error_message
+from ..base import ConfigCheck, IssueCollector
+from ..types import ConfigView, PreflightContext
+
+__all__ = ["StructuralTagBackendCheck"]
+
+
+class StructuralTagBackendCheck(ConfigCheck):
+    """Reject ``structural_tag`` when the structured-output backend cannot support it."""
+
+    name = "config.structured_tag_backend"
+    label = "Structural Tag backend"
+    category = "configuration"
+
+    def enabled(self, ctx: PreflightContext) -> bool:
+        if not super().enabled(ctx):
+            return False
+        generation = ctx.config.generation
+        return (
+            generation.use_structured_generation and generation.structured_generation_schema_method == "structural_tag"
+        )
+
+    def check(self, ctx: ConfigView, collector: IssueCollector) -> None:
+        message = structural_tag_backend_error_message(ctx.config.generation.structured_generation_backend)
+        if message is not None:
+            collector.error("structured_tag_backend_incompatible", message)