LM Format Enforcer

Warnings

• breaking The return type of `TokenEnforcer.get_allowed_tokens()` changed in v0.11.1 to be torch tensor bitmask based for vLLM V1 integration. This is a breaking change if you directly call or rely on the return type of this internal function.
  Fix: Avoid direct calls to `TokenEnforcer.get_allowed_tokens()`. If you need similar functionality, check the latest integration patterns or raise a GitHub issue. For vLLM integration, ensure `use_bitmask=True` is handled if customizing.
• gotcha The primary user-facing API for integrating `lm-format-enforcer` with LLMs (e.g., Transformers, vLLM) involves specific `build_..._prefix_allowed_tokens_fn` functions, rather than directly instantiating a generic `LMFormatEnforcer` class. This is a common point of confusion for new users.
  Fix: Refer to the library's examples and documentation for the correct integration pattern, typically importing `build_transformers_prefix_allowed_tokens_fn` or `build_vllm_prefix_allowed_tokens_fn` from their respective `integrations` submodules.
• gotcha The library has specific Python version requirements (`>=3.8, <4.0`). Using unsupported Python versions may lead to unexpected errors or installation issues.
  Fix: Ensure your environment uses a compatible Python version (e.g., Python 3.8, 3.9, 3.10, 3.11).
• gotcha Compatibility with `transformers` and `pydantic` libraries can be strict. Older versions of `transformers` (pre-4.38.0) and `pydantic` (pre-2.0.0) might cause issues or not be fully supported.
  Fix: Always use the officially recommended or latest compatible versions of `transformers` (>=4.38.0) and `pydantic` (>=2.0.0,<3.0.0) as specified in the `pyproject.toml` or `setup.py`.
• gotcha When integrating with vLLM, ensure your vLLM version is compatible with the `lm-format-enforcer` version. Specific vLLM versions (e.g., vLLM V1) may require particular `lm-format-enforcer` versions and features like the `use_bitmask` flag.
  Fix: Consult the `lm-format-enforcer` GitHub README and release notes for the recommended vLLM version and any specific flags or configuration needed for stable integration.

Install

• pip install lm-format-enforcer Standard installation
• pip install lm-format-enforcer[vllm] Installation with vLLM support

Imports

• JsonSchemaParser

  from lm_format_enforcer.json_schema_parser import JsonSchemaParser

• RegexParser

  from lm_format_enforcer.regex_parser import RegexParser

• build_transformers_prefix_allowed_tokens_fn

  from lm_format_enforcer.integrations.transformers import build_transformers_prefix_allowed_tokens_fn

  The core 'LMFormatEnforcer' class is usually an internal component; users interact with integration-specific builder functions like this one, combined with a parser.
• build_vllm_prefix_allowed_tokens_fn

  from lm_format_enforcer.integrations.vllm import build_vllm_prefix_allowed_tokens_fn

  Similar to transformers integration, vLLM users should import the builder function rather than a direct enforcer class.

Quickstart

This quickstart demonstrates how to enforce a JSON Schema output using a Hugging Face Transformers model. It initializes a tokenizer and model, defines a JSON schema, creates a `JsonSchemaParser`, and then uses `build_transformers_prefix_allowed_tokens_fn` to generate text that strictly adheres to the defined format.

from transformers import AutoTokenizer, AutoModelForCausalLM
from lm_format_enforcer.json_schema_parser import JsonSchemaParser
from lm_format_enforcer.integrations.transformers import build_transformers_prefix_allowed_tokens_fn
import torch

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")

# Define the JSON schema
json_schema = {
    "type": "object",
    "properties": {
        "name": {"type": "string"},
        "age": {"type": "integer", "minimum": 0},
        "isStudent": {"type": "boolean"}
    },
    "required": ["name", "age", "isStudent"]
}

# Create the parser
json_parser = JsonSchemaParser(json_schema)

# Build the prefix_allowed_tokens_fn for transformers integration
prefix_allowed_tokens_fn = build_transformers_prefix_allowed_tokens_fn(tokenizer, json_parser)

prompt = "Please generate a JSON object describing a person with name, age, and student status:\n"

# Encode the prompt
input_ids = tokenizer.encode(prompt, return_tensors="pt")

# Generate text with format enforcement
# GPT2 might not perfectly follow instructions but the *format* will be enforced.
output = model.generate(
    input_ids,
    max_new_tokens=100,
    prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
    pad_token_id=tokenizer.eos_token_id,
    do_sample=False, # For deterministic generation where possible
    num_beams=1
)

# Decode and print the result
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
print(generated_text)
# Example output (format enforced):
# {"name": "Alice", "age": 25, "isStudent": true}