prob_jsonformer/example.ipynb

In [1]:

# autoreload your package
%load_ext autoreload
%autoreload 2

In [2]:

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

print("Loading model and tokenizer...")
model_name = "databricks/dolly-v2-3b"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    use_cache=True,
    torch_dtype=torch.float16,
    attn_implementation='eager',
).to("cuda:0")
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, use_cache=True)
print("Loaded model and tokenizer")

/media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/.venv/lib/python3.9/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...

/media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/.venv/lib/python3.9/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.
  warnings.warn(
/media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/.venv/lib/python3.9/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.
  warnings.warn(
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.

Loaded model and tokenizer

Continue

In [4]:

from prob_jsonformer.format import highlight_values
from prob_jsonformer.main import Jsonformer

ecomm = {
    "type": "object",
    "properties": {
        "store": {
            "type": "object",
            "properties": {
                "name": {"type": "string"},
                "location": {"type": "string"},
                "choice_probs": {"type": "choice_probs", "enum": ["ski", "snowboard", "walk", "pretend"]},
                "inventory": {
                    "type": "array",
                    "items": {
                        "type": "object",
                        "properties": {
                            "productId": {"type": "string"},
                            "name": {"type": "string"},
                            "description": {"type": "string"},
                            "category": {"type": "string"},
                            "price": {"type": "number"},
                            "inStock": {"type": "boolean"},
                            "rating": {"type": "number"},
                            "images": {"type": "array", "items": {"type": "string"}},
                        },
                    },
                },
            },
        }
    },
}


builder = Jsonformer(
    model=model,
    tokenizer=tokenizer,
    json_schema=ecomm,
    prompt="write a description about mike's ski shop which sells premium skis and snowboards",
    max_string_token_length=20,
)

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

highlight_values(output)

Generating...
{
  store: {
    name: "Mike's Ski Shop",
    location: "Somewhere",
    choice_probs: [
      {
        prob: 0.01739501953125,
        choice: "pretend"
      },
      {
        prob: 0.002094268798828125,
        choice: "snowboard"
      },
      {
        prob: 0.0007467269897460938,
        choice: "walk"
      },
      {
        prob: 0.97998046875,
        choice: "ski"
      }
    ],
    inventory: [
      {
        productId: "1",
        name: "Snowboard X-15",
        description: "Snowboard for all levels",
        category: "Snowboards",
        price: 20.0375,
        inStock: True,
        rating: 5.0,
        images: [
          "https://s3.amazonaws.com/mikesskisport/images/Snow"
        ]
      }
    ]
  }
}

In [5]:

car = {
    "type": "object",
    "properties": {
        "make": {"type": "string"},
        "model": {"type": "choice_probs", "enum": ["Mazda", "Kea"]},
        "new": {"type": "choice_probs", "enum": ["true", "false"]},
        "rating": {"type": "choice_probs", "enum": ["1", "2", "3", "4"]},
        "year": {"type": "number"},
        "colors_available": {
            "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: "Mazda",
  model: [
    {
      prob: 0.8154296875,
      choice: "Kea"
    },
    {
      prob: 0.184814453125,
      choice: "Mazda"
    }
  ],
  new: [
    {
      prob: 0.90185546875,
      choice: "true"
    },
    {
      prob: 0.09808349609375,
      choice: "false"
    }
  ],
  rating: [
    {
      prob: 0.221435546875,
      choice: "1"
    },
    {
      prob: 0.394775390625,
      choice: "2"
    },
    {
      prob: 0.382568359375,
      choice: "3"
    },
    {
      prob: 0.0013370513916015625,
      choice: "4"
    }
  ],
  year: 2016.0,
  colors_available: [
    "red"
  ]
}

In [6]:

complex_car = {
    "type": "object",
    "properties": {
        "car": {
            "type": "object",
            "properties": {
                "make": {"type": "string"},
                "model": {"type": "string"},
                "year": {"type": "number"},
                "colors": {"type": "choice_probs", "enum": ["red", "green", "blue", "black", "white"]},
                "as_new": {"type": "choice_probs", "enum": ["true", "false"]},
                "rating": {"type": "choice_probs", "enum": ["1", "2", "3", "4"]},
                "features": {
                    "type": "object",
                    "properties": {
                        "audio": {
                            "type": "object",
                            "properties": {
                                "brand": {"type": "string"},
                                "speakers": {"type": "number"},
                                "hasBluetooth": {"type": "boolean"},
                            },
                        },
                        "safety": {
                            "type": "object",
                            "properties": {
                                "airbags": {"type": "number"},
                                "parkingSensors": {"type": "boolean"},
                                "laneAssist": {"type": "boolean"},
                            },
                        },
                        "performance": {
                            "type": "object",
                            "properties": {
                                "engine": {"type": "string"},
                                "horsepower": {"type": "number"},
                                "topSpeed": {"type": "number"},
                            },
                        },
                    },
                },
            },
        },
        "owner": {
            "type": "object",
            "properties": {
                "firstName": {"type": "string"},
                "lastName": {"type": "string"},
                "age": {"type": "number"},
            },
        },
    },
}
builder = Jsonformer(
    model=model,
    tokenizer=tokenizer,
    json_schema=complex_car,
    prompt="generate an example Rolls Royce Phantom",
)

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

highlight_values(output)

Generating...
{
  car: {
    make: "Rolls Royce",
    model: "Phantom",
    year: 2014.0,
    colors: [
      {
        prob: 0.001560211181640625,
        choice: "white"
      },
      {
        prob: 0.833984375,
        choice: "red"
      },
      {
        prob: 0.0865478515625,
        choice: "black"
      },
      {
        prob: 0.048553466796875,
        choice: "blue"
      },
      {
        prob: 0.0294342041015625,
        choice: "green"
      }
    ],
    as_new: [
      {
        prob: 0.96533203125,
        choice: "true"
      },
      {
        prob: 0.03460693359375,
        choice: "false"
      }
    ],
    rating: [
      {
        prob: 0.05462646484375,
        choice: "1"
      },
      {
        prob: 0.233642578125,
        choice: "2"
      },
      {
        prob: 0.2093505859375,
        choice: "3"
      },
      {
        prob: 0.50244140625,
        choice: "4"
      }
    ],
    features: {
      audio: {
        brand: "Mercedes-Benz",
        speakers: 2.09999,
        hasBluetooth: True
      },
      safety: {
        airbags: 2.09999,
        parkingSensors: True,
        laneAssist: True
      },
      performance: {
        engine: "Mercedes-Benz 6.2 L",
        horsepower: 423.09999,
        topSpeed: 220.09999
      }
    }
  },
  owner: {
    firstName: "John",
    lastName: "Doe",
    age: 40.09999
  }
}