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` | `"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"` |

docs/user-guide/running.md

@@ -883,9 +883,11 @@ 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: "auto"
     ```
 
+- `"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.
 

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,57 @@
     ValueValidator,
     range_validator,
 )
+from ..errors import ParameterError
+
+StructuredGenerationSchemaMethod = Literal["auto", "regex", "json_schema", "structural_tag"]
+ResolvedStructuredGenerationSchemaMethod = Literal["regex", "json_schema", "structural_tag"]
+StructuredGenerationBackend = Literal["auto", "xgrammar", "guidance", "outlines", "lm-format-enforcer"]
+
+STRUCTURAL_TAG_COMPATIBLE_BACKENDS = frozenset({"auto", "xgrammar"})
+
+__all__ = [
+    "GenerateParameters",
+    "ResolvedStructuredGenerationSchemaMethod",
+    "STRUCTURAL_TAG_COMPATIBLE_BACKENDS",
+    "StructuredGenerationBackend",
+    "StructuredGenerationSchemaMethod",
+    "ValidationParameters",
+    "resolve_structured_generation_schema_method",
+    "structural_tag_backend_error_message",
+]
+
+
+def resolve_structured_generation_schema_method(
+    schema_method: StructuredGenerationSchemaMethod,
+    backend: StructuredGenerationBackend | str,
+) -> ResolvedStructuredGenerationSchemaMethod:
+    """Resolve ``auto`` schema method from the configured structured-output backend.
+
+    ``auto`` picks ``structural_tag`` on xgrammar-capable backends and ``regex``
+    elsewhere, preserving legacy behavior for outlines/guidance configs that omit
+    an explicit schema method.
+    """
+    if schema_method != "auto":
+        return schema_method
+    if backend in STRUCTURAL_TAG_COMPATIBLE_BACKENDS:
+        return "structural_tag"
+    return "regex"
+
 
-__all__ = ["GenerateParameters", "ValidationParameters"]
+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 +197,18 @@ class GenerateParameters(Parameters, BaseModel):
     ] = "auto"
 
     structured_generation_schema_method: Annotated[
-        Literal["regex", "json_schema"],
+        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."
+                "than 'json_schema' at the cost of speed. 'structural_tag' uses XGrammar Structural Tag "
+                "to compose schema-constrained JSONL output."
             ),
         ),
-    ] = "regex"
+    ] = "auto"
 
     structured_generation_use_single_sequence: Annotated[
         bool,
@@ -191,3 +243,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,16 @@
 from .. import utils
 from ..cli.artifact_structure import Workdir
 from ..config import SafeSynthesizerParameters
+from ..config.generate import (
+    resolve_structured_generation_schema_method,
+    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
@@ -324,8 +328,12 @@ def _build_structured_output_params(self) -> StructuredOutputsParams | None:
             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,
+        )
 
-        if self.config.generation.structured_generation_schema_method == "regex":
+        if schema_method == "regex":
             logger.info("Structured generation is enabled, using a regex to enforce the schema")
             pc = self.model_metadata.prompt_config
             regex = build_json_based_regex(
@@ -335,8 +343,20 @@ def _build_structured_output_params(self) -> StructuredOutputsParams | None:
                 eos_token=pc.eos_token,
             )
             params["regex"] = regex
-        elif self.config.generation.structured_generation_schema_method == "json_schema":
+        elif schema_method == "json_schema":
             params["json"] = self.schema
+        elif 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)