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_name = "failspy/kappa-3-phi-abliterated"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    use_cache=True,
    torch_dtype=torch.float16,
    #  device="cuda:0",
    #  device_map="auto",
    attn_implementation='eager',
    trust_remote_code=True,
).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(
`flash-attention` package not found, consider installing for better performance: No module named 'flash_attn'.
Current `flash-attenton` does not support `window_size`. Either upgrade or use `attn_implementation='eager'`.
Loading checkpoint shards: 100%|██████████| 4/4 [00:02<00:00,  1.94it/s]
/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

Scratch

In [3]:

# from jaxtyping import Float, Int
# import torch
# from torch.nn import functional as F
# from torch import Tensor
# from typing import List, Callable, Tuple, Dict, Optional
# import pandas as pd

In [4]:

# # initital state
# prompt = "The necromancer in his tower, what's his top problem? "
# choices = ["1", "the skeleton", "the boney boys"]
# choices_tokens = tokenizer(choices).input_ids
# choices_tokens = [torch.tensor(c) for c in choices_tokens]
# # current_tokens = torch.tensor([])

# # next
# input_ids = tokenizer([prompt], return_tensors="pt").to(model.device).input_ids[0]
# choices_tokens

# # for each next choice, continue down the tree, recording the log probs

In [5]:

# def get_valid_next_choices(choices_tokens, current_tokens):
#     next_choices = []
#     for choice_tokens in choices_tokens:
#         # if we have some more slots left
#         if len(current_tokens)<len(choice_tokens):
#             # see if current_tokens matches
#             if (choice_tokens[: len(current_tokens)] == current_tokens).all():
#                 c = choice_tokens[len(current_tokens)].item()
#                 next_choices.append(c)

#     next_choices = list(set(next_choices))
#     return torch.LongTensor(next_choices)


# def gen_choice_probs(
#     model: AutoModelForCausalLM,
#     tokenizer: AutoTokenizer,
#     input_ids: Int[Tensor, "seq"],
#     choices_tokens: List[Int[Tensor, "seq"]],
#     choice: Optional[Int[Tensor, ""]] = None,
#     prob: float = 1,
#     current_tokens: Int[Tensor, "seq"] = torch.LongTensor([]),
#     z=[],
# ):
#     if choice is not None:
#         c = choice[None].to(current_tokens.device)
#         current_tokens = torch.cat([current_tokens, c], dim=-1)
#         print(current_tokens, 'current_tokens')
#         c = choice[None].to(input_ids.device)
#         input_ids = torch.cat([input_ids, c], dim=-1)

#     next_choices = get_valid_next_choices(choices_tokens, current_tokens)
#     if len(next_choices) == 0:
#         s = tokenizer.decode(current_tokens)
#         r = dict(tokens=current_tokens.cpu(), prob=prob, choice=s)
#         yield r
#     else:
#         o = model(input_ids[None])
#         logits_constrained = o.logits[0, -1][next_choices]
#         probs = F.softmax(logits_constrained, dim=-1)
#         for i in range(len(next_choices)):
#             next_choice = next_choices[i]
#             next_prob = prob * probs[i].item()
#             yield from gen_choice_probs(
#                 model=model,
#                 tokenizer=tokenizer,
#                 choices_tokens=choices_tokens,
#                 input_ids=input_ids,
#                 choice=next_choice,
#                 prob=next_prob,
#                 current_tokens=current_tokens,
#                 z=z + [i],
#             )


# r = list(gen_choice_probs(model, tokenizer, input_ids, choices_tokens))

In [6]:

# r

In [7]:

# pd.DataFrame(r).sort_values("prob", ascending=False).drop(columns=["tokens"])

Continue

In [8]:

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"},
                "choices": {"type": "choices", "enum": ["1", "the", "they walked the old dog", "1 the"]},
                "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...

You are not running the flash-attention implementation, expect numerical differences.

In [ ]:

car = {
    "type": "object",
    "properties": {
        "make": {"type": "string"},
        "model": {"type": "string"},
        "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)

In [ ]:

complex_car = {
    "type": "object",
    "properties": {
        "car": {
            "type": "object",
            "properties": {
                "make": {"type": "string"},
                "model": {"type": "string"},
                "year": {"type": "number"},
                "colors": {"type": "array", "items": {"type": "string"}},
                "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)