transformers-grammar-constraint 0.2.2

Grammar-constrained LLM token generation via Lark + HuggingFace Transformers

transformers-grammar-constraint

Enforce formal grammar constraints on LLM token generation with HuggingFace Transformers. Uses Lark to define grammars and masks invalid tokens at each generation step, so the model can only produce output that conforms to your grammar.

jsonformer exists if you need something more light weight that only supports JsonSchema.

from grammar_constrain import GrammarConstrainedLogitsProcessor, JsonGrammar

grammar = JsonGrammar(schema={"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]})
processor = GrammarConstrainedLogitsProcessor(grammar=grammar, tokenizer=tokenizer)

output = model.generate(input_ids, logits_processor=[processor], max_new_tokens=50)
# Guaranteed to produce valid JSON matching the schema

Installation

uv add transformers-grammar-constraint

Requirements: Python >= 3.12, PyTorch >= 2.10, Transformers 4.x or 5.x (>= 4.17)

How It Works

GrammarConstrainedLogitsProcessor is a LogitsProcessor that plugs into the model.generate() call. Before each token is sampled, it computes which token IDs are valid given the current parse state and sets all others to -inf. The model never sees invalid candidates.

---

Usage

With model.generate()

from transformers import AutoTokenizer, AutoModelForCausalLM
from grammar_constrain import GrammarConstrainedLogitsProcessor, JsonGrammar

tokenizer = AutoTokenizer.from_pretrained("your-model")
model = AutoModelForCausalLM.from_pretrained("your-model")

grammar = JsonGrammar()
processor = GrammarConstrainedLogitsProcessor(grammar=grammar, tokenizer=tokenizer)

input_ids = tokenizer("Output JSON: ", return_tensors="pt").input_ids
output = model.generate(
    input_ids,
    logits_processor=[processor],
    max_new_tokens=100,
    do_sample=False,
)

text = tokenizer.decode(output[0][input_ids.shape[1]:], skip_special_tokens=True)
# text is guaranteed to be valid JSON

With pipeline()

from transformers import pipeline
from grammar_constrain import GrammarConstrainedLogitsProcessor, JsonGrammar

pipe = pipeline("text-generation", model="your-model")

grammar = JsonGrammar()
processor = GrammarConstrainedLogitsProcessor(grammar=grammar, tokenizer=pipe.tokenizer)

result = pipe(
    "Output JSON: ",
    logits_processor=[processor],
    max_new_tokens=100,
    return_full_text=False,
    do_sample=False,
)

text = result[0]["generated_text"].strip()
# text is guaranteed to be valid JSON

Reusing a processor across calls

Create one processor and call it repeatedly — it resets its state automatically between independent generate() / pipeline calls.

Thinking models

For models that use <think>...</think> reasoning blocks, pass the think tokens so constraints are suspended inside them:

processor = GrammarConstrainedLogitsProcessor(
    grammar=grammar,
    tokenizer=tokenizer,
    think_start_token="<think>",
    think_end_token="</think>",
)

Set either to None to disable think-block handling entirely.

---

Available Grammars

JsonGrammar

Constrains generation to valid JSON. Optionally accepts a JSON Schema to further restrict structure.

from grammar_constrain import JsonGrammar

# Any valid JSON
grammar = JsonGrammar()

# Object matching a schema
grammar = JsonGrammar(schema={
    "type": "object",
    "properties": {
        "name": {"type": "string"},
        "age":  {"type": "integer"},
    },
    "required": ["age", "name"],            # enforced in alphabetical order
    "additionalProperties": False,
})

# Array of strings
grammar = JsonGrammar(schema={
    "type": "array",
    "items": {"type": "string"},
})

# Enum
grammar = JsonGrammar(schema={
    "type": "string",
    "enum": ["red", "green", "blue"],
})

Supported JSON Schema keywords: type, properties, required, items, enum, additionalProperties.

Note: required keys are enforced in alphabetical order. Schema keywords such as oneOf, anyOf, allOf, $ref, pattern, minLength, maxLength, minimum, and maximum are not yet supported and will be ignored with a warning.

---

SanGrammar and PgnGrammar

Constrain generation to valid chess notation. Both are implemented as hand-crafted character-level DFAs (not Lark grammars), because SAN disambiguation (e.g. Nbd2, R1e2) is ambiguous for LALR(1). PgnGrammar embeds the SAN DFA for individual move bodies and adds structural states for move numbers, spacing, and game results.

from grammar_constrain import SanGrammar, PgnGrammar

# Single move in Standard Algebraic Notation
grammar = SanGrammar()
# Valid output examples: "e4", "Nf3", "Rxe5+", "O-O", "e8=Q#", "Nbd2"

# Full game in Portable Game Notation
grammar = PgnGrammar()
# Valid output example: "1. e4 e5 2. Nf3 Nc6 1-0"

---

Extending to New Grammars

Subclass Grammar and implement the grammar_string property using Lark EBNF syntax.

from grammar_constrain import Grammar

class TemperatureGrammar(Grammar):
    @property
    def grammar_string(self) -> str:
        return r"""
        start: number unit
        number: /-?\d+(\.\d+)?/
        unit: "°C" | "°F" | "K"

        %import common.WS
        %ignore WS
        """

Use it like any built-in grammar:

processor = GrammarConstrainedLogitsProcessor(
    grammar=TemperatureGrammar(),
    tokenizer=tokenizer,
)

Choosing a parser backend

The default parser is LALR (fast). If your grammar is ambiguous or has shift-reduce conflicts, switch to Earley:

class MyGrammar(Grammar):
    lark_parser = "earley"   # override class variable

    @property
    def grammar_string(self) -> str:
        return r"..."

LALR is significantly faster; prefer it when possible.

Inline grammars (no subclass needed)

Pass a raw Lark grammar string directly to the processor:

processor = GrammarConstrainedLogitsProcessor(
    grammar=r"""
    start: ("yes" | "no") /\n/
    """,
    tokenizer=tokenizer,
)

Validating output

Every Grammar exposes is_valid_complete(text) to check whether a string is a complete, valid parse:

grammar = JsonGrammar()
assert grammar.is_valid_complete('{"key": 1}')   # True
assert not grammar.is_valid_complete('{"key":')  # False — incomplete

---

Running Tests

# Unit tests (no model download required)
uv run pytest tests/ -v

# Integration tests with a real model (requires network access)
uv run pytest tests/test_integration.py -v -m slow

---

License

MIT