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")

Loading model and tokenizer...

config.json:   0%|          | 0.00/819 [00:00<?, ?B/s]

pytorch_model.bin:   0%|          | 0.00/5.68G [00:00<?, ?B/s]

model.safetensors:   0%|          | 0.00/5.68G [00:00<?, ?B/s]

tokenizer_config.json:   0%|          | 0.00/450 [00:00<?, ?B/s]

tokenizer.json:   0%|          | 0.00/2.11M [00:00<?, ?B/s]

special_tokens_map.json:   0%|          | 0.00/228 [00:00<?, ?B/s]

Loaded model and tokenizer

In [3]:

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

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[3], line 29
     27 prompt = "Generate a young person's information based on the following schema:"
     28 jsonformer = Jsonformer(model, tokenizer, json_schema, prompt)
---> 29 generated_data = jsonformer()
     31 generated_data

File /media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/prob_jsonformer/main.py:439, in Jsonformer.__call__(self)
    437 def __call__(self) -> Dict[str, Any]:
    438     self.value = {}
--> 439     generated_data = self.generate_object(
    440         self.json_schema["properties"], self.value
    441     )
    442     return generated_data

File /media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/prob_jsonformer/main.py:274, in Jsonformer.generate_object(self, properties, obj)
    272 for key, schema in properties.items():
    273     self.debug("[generate_object] generating value for", key)
--> 274     obj[key] = self.generate_value(schema, obj, key)
    275 return obj

File /media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/prob_jsonformer/main.py:351, in Jsonformer.generate_value(self, schema, obj, key)
    349     else:
    350         obj.append(self.generation_marker)
--> 351     return self.generate_p_enum(schema["values"], round=schema.get("round", 3))
    352 elif schema_type == "p_integer":
    353     if key:

KeyError: 'values'

In [4]:

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",
}

In [ ]: