Jsonformer: Bulletproof Structured JSON Generation from Language Models

Introduction

Generating structured JSON from language models is a significant challenge, often resulting in syntactically incorrect or schema-non-compliant outputs. Traditional methods relying on prompt engineering, fine-tuning, or post-processing are frequently brittle. Jsonformer offers an innovative solution by acting as a wrapper around Hugging Face models. It intelligently fills in fixed JSON tokens during generation, delegating only the content tokens to the language model. This method ensures efficiency and bulletproof reliability for structured data generation, supporting various JSON schema types like number, boolean, string, array, and object.

Installation

To get started with Jsonformer, install it using pip:

pip install jsonformer

Examples

Here's a basic example demonstrating how to use Jsonformer to generate structured data based on a defined schema:

from jsonformer import Jsonformer
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("databricks/dolly-v2-12b")
tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b")

json_schema = {
    "type": "object",
    "properties": {
        "name": {"type": "string"},
        "age": {"type": "number"},
        "is_student": {"type": "boolean"},
        "courses": {
            "type": "array",
            "items": {"type": "string"}
        }
    }
}

prompt = "Generate a person's information based on the following schema:"
jsonformer = Jsonformer(model, tokenizer, json_schema, prompt)
generated_data = jsonformer()

print(generated_data)

Jsonformer also handles complex, nested schemas effectively, even with smaller models.

Why Use Jsonformer?

Jsonformer stands out for several key features:

• Bulletproof JSON Generation: It guarantees that the generated JSON is always syntactically correct and adheres strictly to the specified schema, eliminating common errors.
• Efficiency: By intelligently generating only the variable content tokens and filling in the fixed structural tokens, Jsonformer is significantly more efficient than traditional methods that generate and then parse full JSON strings.
• Flexible and Extendable: Built upon the Hugging Face transformers library, Jsonformer is compatible with any language model that supports the Hugging Face interface, offering broad applicability.

Links

• GitHub Repository: 1rgs/jsonformer
• Colab Example: Jsonformer_example.ipynb