grammars 0.0.4
Enforce language model output using a grammar
License: MIT
Requires Python: >=3.8,<4.0
Keywords: ai,large language models,llm,parser,prompt

# Grammars

**Grammars** is an experimental, early stage Python library for constraining the output of [Transformers](https://github.com/huggingface/transformers) language models using [formal grammars](https://en.wikipedia.org/wiki/Formal_grammar) specified via [parser combinators](https://en.wikipedia.org/wiki/Parser_combinator).


```python
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "gpt2-xl"
tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side='left')
model = AutoModelForCausalLM.from_pretrained(model_id).to("cuda")
tokenizer.pad_token = tokenizer.eos_token

from grammars import G, decode

grammar = G.seq(
    G.literal(" "), 
    G.alt(G.literal("world"), G.literal("computer"))
)

prompt = "Hello"
results = decode(prompt, grammar, model, tokenizer)
assert results[0] == " world"
```

### Installation

`pip install grammars`

## Overview

This library has two packages:

- `parse` — **Streaming parser combinators**: allows you to define a grammar and can parse one character at a time, letting you know if the parse can continue and/or if it is complete
- `decode` — **Constrained decoding**: uses the grammar/streaming parsers to constrain the output of a language model


## Constrained decoding

The constrained decoding takes places by converting a grammar into two different modes — "rank" and "generate":

"Rank" takes a set of possible strings and determines how likely they are according to model, using the LM to decode in parallel.

"Generate" uses a [LogitsProcessor](https://huggingface.co/docs/transformers/internal/generation_utils#transformers.LogitsProcessor) to perform constrained decoding using a regular expression pattern, using partial match support from the [`regex` library](https://github.com/mrabarnett/mrab-regex). 


Here's an example of how a grammar is transformed into these two modes:
```python
from grammars import G

grammar = G.seq(
        G.alt(G.literal("dark"), G.literal("light")),
        G.literal(" "),
        G.re("[a-z]+"),
        G.literal(" "),
        G.alt(G.literal("blue"), G.literal("red"))
    )

```

The grammar begins with two possible strings—"dark " or "light "—and then is followed by a regular expression — followed by " blue" or " red". Respectively, this is converted into a Rank node, followed by a Generate node, followed by a Rank node.

![Modal node diagram](./assets/node.gv.svg)

### Ranking

Each set of subtokens (e.g. "Red", "Red velvet",...) for each possible string (e.g. "Red velvet cake" and "Lemon lime soda") is decoded in parallel. The results' logits are then used to rank the possible strings.

```python
from grammars import G

grammar = G.alt(
    G.literal("Red velvet cake"),
    G.literal("Lemon lime soda")
)
```

*Example using GPT2-XL tokenizer*

|                  | input 0 | input 1 | input 2   | generated 0 |
|------------------|---------|---------|-----------|-------------|
| Candidate 0, N-1 | \<pad\> | \<pad\> | "Red"     | " velvet"   |
| Candidate 0, N   | \<pad\> | "Red"   | " velvet" | " cake"     |
| Candidate 1, N-2 | \<pad\> | \<pad\> | "L"       | "emon"      |
| Candidate 1, N-1 | \<pad\> | "L"     | "emon"    | "lime"      |
| Candidate 1, N   | "L"     | "emon"  | "lime"    | "soda"      |


### Generation

Generate mode uses a [LogitsProcessor](https://huggingface.co/docs/transformers/internal/generation_utils#transformers.LogitsProcessor) parameterize the GenerateNode's parser, to perform constrained decoding. For each possible next token, the `step` method on the current parser is called to determine if the token is a valid (potentially partial) input.

Currently only greedy decoding works; a potentially next step is getting other forms of decoding such as beam search to work.


## Streaming parsers

**Grammars** contains a streaming parser combinators library that can be used independently of the LM-related part of this library.

Using the library, a grammar is formed by combining parsers together by combining base parsers (RegexParser, LiteralParser, FloatParser) with combinators (SeqParser and AltParser). Each parser, has a `step` method, which takes in a string to parse that may be incomplete - returning a new Parser is the parse is valid so far and None if invalid. Each parser has `completed` and `can_continue` boolean flags to convey whether it is in a completed valid state and/or whether there may be a valid continuation.

For example:

```python
from grammars.parse import G

parser = G.literal("hello")

# valid but incomplete
parser = parser.step("hel")
assert not parser.completed
assert parser.can_continue

# valid and complete
parser = parser.step("lo")
assert parser.completed
assert not parser.can_continue

# invalid
parser = parser.step("lo")
assert not parser
```

Although having a RegexParser affords a good deal of expressivity, a user could also define their own custom parsers, which will work fine with the rest of the library, using `generate` mode.

Example of a custom parser:
```python
from dataclasses import dataclass
from grammars.parse import Parser

@dataclass(frozen=True)
class AlphabetParser(Parser):
    completed: bool
    can_continue: bool

    def step(self, substring: str) -> "AlphabetParser":
        if substring.isalpha():
            return AlphabetParser(
                completed=True,
                can_continue=True
            )
        else:
            return None
```


## Motivation

Advances in large language model capabilities have led to a flurry of interest in integrating them as components in larger software projects. However a challenge to integration is being able to depend on or control the output of LLMs.

I've been interested in exploring how constrained decoding can address this need, and have been excited to see projects such as Ben Newhouse's [clownfish](https://github.com/newhouseb/clownfish) library. While Clownfish focuses on parsing JSON that matches a Pydantic schema, this project allows for an arbitrary user-specified grammar. In addition, I've been in interested in exploring more efficient inference/decoding techniques, such as [LLMA](https://arxiv.org/abs/2304.04487) decoding. 


### References / related projects:

- ["Structural Alignment: Modifying Transformers (like GPT) to Follow a JSON Schema"](https://github.com/newhouseb/clownfish), Ben Newhouse
- [Guidance](https://github.com/microsoft/guidance) library from Microsoft
- [Inference with Reference: Lossless Acceleration of Large Language Models](https://arxiv.org/abs/2304.04487), Nan Yang et al.
  - [LLMA: Large Language Model Accelerator](https://github.com/microsoft/LMOps/tree/main/llma) 
- [ReLLM](https://github.com/r2d4/rellm): Exact structure out of any language model completion
- [LMQL](https://lmql.ai/)