Metadata-Version: 2.3
Name: transformers-grammar-constraint
Version: 0.2.2
Summary: Grammar-constrained LLM token generation via Lark + HuggingFace Transformers
Requires-Dist: lark>=1.3.1
Requires-Dist: torch>=2.10.0
Requires-Dist: transformers>=4.17.0,<6.0.0
Requires-Python: >=3.12
Description-Content-Type: text/markdown

# transformers-grammar-constraint

Enforce formal grammar constraints on LLM token generation with HuggingFace Transformers. Uses [Lark](https://github.com/lark-parser/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](https://github.com/1rgs/jsonformer) exists if you need something more light weight that only supports JsonSchema.

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

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

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

```python
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:

```python
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.

```python
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.

```python
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](https://lark-parser.readthedocs.io/en/latest/grammar.html).

```python
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:

```python
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:

```python
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:

```python
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:

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

---

## Running Tests

```bash
# 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
