prob_jsonformer/dev.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

In [18]:

from prob_jsonformer import Jsonformer

json_schema = {
    "type": "object",
    "properties": {
        # we can return the probability of each choice, even if they are multiple tokens
        "age_probs": {"type": "p_enum", "enum": [str(s) for s in range(10, 30)]},
        # we can return the probabilistic weighted mean of a range
        "age_wmean": {"type": "p_integer", "minimum": 10, "maximum": 30},
        # the prob of true and false
        "is_student_probs": {"type": "p_enum", "enum": ["true", "false"]},
        "is_student": {"type": "boolean"},
        # we've merged patches for enum, integer, null, union - currently mising from jsonformer
        "name": {"type": "string", "maxLength": 4},
        "age": {"type": "integer"},
        "unit_time": {"type": "number"},
        "courses": {"type": "array", "items": {"type": "string"}},
        "trim": {"type": ["string", "null"]},
        "color": {
            "type": "enum",
            "values": ["red", "green", "blue", "brown", "white", "black"],
        },
    },
}


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

generated_data

Out [18]:

{'age_probs': [{'prob': 0.94091796875, 'choice': '10'},
  {'prob': 0.033233642578125, 'choice': '20'},
  {'prob': 0.0122222900390625, 'choice': '12'},
  {'prob': 0.00412750244140625, 'choice': '21'},
  {'prob': 0.0028362274169921875, 'choice': '16'},
  {'prob': 0.0018453598022460938, 'choice': '15'},
  {'prob': 0.00113677978515625, 'choice': '11'},
  {'prob': 0.0011110305786132812, 'choice': '18'},
  {'prob': 0.0005083084106445312, 'choice': '25'},
  {'prob': 0.0004558563232421875, 'choice': '23'},
  {'prob': 0.0002498626708984375, 'choice': '14'},
  {'prob': 0.00023281574249267578, 'choice': '13'},
  {'prob': 0.0002238750457763672, 'choice': '22'},
  {'prob': 0.00018131732940673828, 'choice': '26'},
  {'prob': 0.0001690387725830078, 'choice': '24'},
  {'prob': 0.00012552738189697266, 'choice': '19'},
  {'prob': 7.796287536621094e-05, 'choice': '27'},
  {'prob': 7.265806198120117e-05, 'choice': '28'},
  {'prob': 4.106760025024414e-05, 'choice': '17'},
  {'prob': 2.5033950805664062e-06, 'choice': '29'}],
 'age_wmean': 17.816404402256012,
 'is_student_probs': [{'prob': 0.974609375, 'choice': 'true'},
  {'prob': 0.025177001953125, 'choice': 'false'}],
 'is_student': False,
 'name': 'John',
 'age': 17,
 'unit_time': 0.5,
 'courses': ['C++'],
 'trim': None,
 'color': 'white'}

In [ ]:

generated_data = {
    "age_probs": [
        {"prob": 0.94091796875, "choice": "10"},
        {"prob": 0.033233642578125, "choice": "20"},
        {"prob": 0.0122222900390625, "choice": "12"},
        {"prob": 0.00412750244140625, "choice": "21"},
        {"prob": 0.0028362274169921875, "choice": "16"},
        {"prob": 0.0018453598022460938, "choice": "15"},
        {"prob": 0.00113677978515625, "choice": "11"},
        {"prob": 0.0011110305786132812, "choice": "18"},
        {"prob": 0.0005083084106445312, "choice": "25"},
        {"prob": 0.0004558563232421875, "choice": "23"},
        {"prob": 0.0002498626708984375, "choice": "14"},
        {"prob": 0.00023281574249267578, "choice": "13"},
        {"prob": 0.0002238750457763672, "choice": "22"},
        {"prob": 0.00018131732940673828, "choice": "26"},
        {"prob": 0.0001690387725830078, "choice": "24"},
        {"prob": 0.00012552738189697266, "choice": "19"},
        {"prob": 7.796287536621094e-05, "choice": "27"},
        {"prob": 7.265806198120117e-05, "choice": "28"},
        {"prob": 4.106760025024414e-05, "choice": "17"},
        {"prob": 2.5033950805664062e-06, "choice": "29"},
    ],
    "age_wmean": 17.816404402256012,
    "is_student_probs": [
        {"prob": 0.974609375, "choice": "true"},
        {"prob": 0.025177001953125, "choice": "false"},
    ],
    "is_student": False,
    "name": "John",
    "age": 17,
    "unit_time": 0.5,
    "courses": ["C++"],
    "trim": None,
    "color": "white",
}