Structured Output (JSON)

Predibase endpoints allow you to enforce that responses contain only of valid JSON and adhere to a provided schema.

The schema can be provided either using JSON schema (REST, Python) or Pydantic (Python).

Example

Python (Pydantic)

from pydantic import BaseModel, constr
from predibase import Predibase


# Initialize Predibase client
pb = Predibase(api_token="<PREDIBASE API TOKEN>")


# Define a schema for the response
class Character(BaseModel):
    name: constr(max_length=10)
    age: int
    strength: int


# Get a handle to the base LLM deployment
lorax_client = pb.deployments.client("mistral-7b-instruct")

# Generate a response that adheres to the schema
response = lorax_client.generate(
    "[INST] Generate a new character for my awesome game. Strength 1-10. [/INST]",
    response_format={
        "type": "json_object",
        "schema": Character.model_json_schema(),
    },
    max_new_tokens=128,
)

# Load the response as JSON and init an object of the desired schema
response_json = json.loads(response.generated_text)
my_character = Character(**response_json)

REST (JSON Schema)

# Export environment variables
export PREDIBASE_API_TOKEN="<YOUR TOKEN HERE>" # Settings > My Profile > Generate API Token
export PREDIBASE_TENANT_ID="<YOUR TENANT ID>" # Settings > My Profile > Overview > Tenant ID
export PREDIBASE_DEPLOYMENT="mistral-7b-instruct"

# query the LLM deployment
curl https://serving.app.predibase.com/$PREDIBASE_TENANT_ID/deployments/v2/llms/$PREDIBASE_DEPLOYMENT/generate \
    -H "Content-Type: application/json" \
    -X POST \
    -d '{
        "inputs": "[INST] Generate a new character for my awesome game. Strength 1-10. [/INST]",
        "parameters": {
            "response_format": {
                "type": "json_object",
                "schema": {"properties": {
                    "name": {"maxLength": 10, "title": "Name", "type": "string"}, 
                    "age": {"title": "Age", "type": "integer"}, 
                    "strength": {"title": "Strength", "type": "integer"}}, 
                    "required": ["name", "age", "strength"], 
                    "title": "Character", 
                    "type": "object"
                }
            },
            "max_new_tokens": 128
        }
    }' \
    -H "Authorization: Bearer ${PREDIBASE_API_TOKEN}"