prob_jsonformer/example.ipynb

In [1]:

from transformers import AutoModelForCausalLM, AutoTokenizer

print("Loading model and tokenizer...")
model_name = "databricks/dolly-v2-12b"
model = AutoModelForCausalLM.from_pretrained(model_name, use_cache=True)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, use_cache=True)
print("Loaded model and tokenizer")

/home/ubuntu/jsonllm/.venv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

Loading model and tokenizer...
Loaded model and tokenizer

In [3]:

from jsonformer.format import highlight_values
from jsonformer.main import Jsonformer

weather_schema = {
    "type": "object",
    "properties": {
        "temperature": {"type": "number"},
        "humidity": {
            "type": "number",
        },
        "wind_speed": {
            "type": "object",
            "properties": {
                "value": {"type": "number"},
                "unit": {"type": "string"},
            },
        },
    },
}

builder = Jsonformer(
    model=model,
    tokenizer=tokenizer,
    json_schema=weather_schema,
    prompt="generate the weather",
)

print("Generating...")
output = builder()

highlight_values(output)

Generating...
{
  temperature: 2.2225,
  humidity: 1.0,
  wind_speed: {
    value: 0.0,
    unit: "value"
  }
}

In [8]:

car = {
    "type": "object",
    "properties": {
        "make": {"type": "string"},
        "model": {"type": "string"},
        "year": {"type": "number"},
        "colors": {
            "type": "array",
            "items": {"type": "string"},
        }
    },
}

builder = Jsonformer(
    model=model,
    tokenizer=tokenizer,
    json_schema=car,
    prompt="generate an example car",
)

print("Generating...")
output = builder()

highlight_values(output)

Generating...
{
  make: "Ford",
  model: "Mustang",
  year: 10.0,
  colors: [
    "red",
    "white",
    "blue",
    "black",
    "yellow",
    "orange",
    "green",
    "pink",
    "purple",
    "violet"
  ]
}

In [9]:

test = {
  "temperature": 24,
  "humidity": 48.0,
  "wind_speed": {
    "value": 12,
    "unit": "mph"
  }
}

highlight_values(test)

{
  temperature: 24,
  humidity: 48.0,
  wind_speed: {
    value: 12,
    unit: "mph"
  }
}

In [ ]:

In [ ]: