wassname/prob_jsonformer
1
0
Fork 0
mirror of https://github.com/wassname/prob_jsonformer.git synced 2026-09-09 11:29:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

prototype a enum with prob

Browse Source
This commit is contained in:
wassname
2024-05-10 16:49:12 +08:00
parent bfad031876
commit 637420044c
5 changed files with 311 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+4
View File
@@ -1,3 +1,7 @@
# prob_jsonformer: A Bulletproof Way to Generate Probabilistic Structured JSON from Language Models.
This fork has been modified to include the token probabilities. The original README is below.
# Jsonformer: A Bulletproof Way to Generate Structured JSON from Language Models.
### Problem: Getting models to output structured JSON is hard
example.ipynb
+286 -104
View File
@@ -9,7 +9,7 @@
"name": "stderr",
"output_type": "stream",
"text": [
"/home/ubuntu/jsonformer/.venv/lib/python3.8/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\n",
"/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\n",
" from .autonotebook import tqdm as notebook_tqdm\n"
]
},
@@ -17,7 +17,27 @@
"name": "stdout",
"output_type": "stream",
"text": [
"Loading model and tokenizer...\n",
"Loading model and tokenizer...\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"/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`.\n",
" warnings.warn(\n",
"`flash-attention` package not found, consider installing for better performance: No module named 'flash_attn'.\n",
"Current `flash-attenton` does not support `window_size`. Either upgrade or use `attn_implementation='eager'`.\n",
"Loading checkpoint shards: 100%|██████████| 4/4 [00:04<00:00, 1.10s/it]\n",
"/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`.\n",
" warnings.warn(\n",
"Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Loaded model and tokenizer\n"
]
}
@@ -26,60 +46,228 @@
"from transformers import AutoModelForCausalLM, AutoTokenizer\n",
"\n",
"print(\"Loading model and tokenizer...\")\n",
"model_name = \"databricks/dolly-v2-3b\"\n",
"model = AutoModelForCausalLM.from_pretrained(model_name, use_cache=True, device_map=\"auto\")\n",
"# model_name = \"databricks/dolly-v2-3b\"\n",
"model_name = \"failspy/kappa-3-phi-abliterated\"\n",
"model = AutoModelForCausalLM.from_pretrained(\n",
" model_name,\n",
" use_cache=True,\n",
" # device=\"cuda:0\",\n",
" # device_map=\"auto\",\n",
" trust_remote_code=True,\n",
").to(\"cuda:0\")\n",
"tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, use_cache=True)\n",
"print(\"Loaded model and tokenizer\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Scratch"
]
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": 65,
"metadata": {},
"outputs": [],
"source": [
"from jaxtyping import Float, Int\n",
"import torch\n",
"from torch.nn import functional as F\n",
"from torch import Tensor\n",
"from typing import List, Callable, Tuple, Dict, Optional\n",
"import pandas as pd"
]
},
{
"cell_type": "code",
"execution_count": 43,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[tensor([ 1, 29871, 29896]),\n",
" tensor([ 1, 278, 18109, 11285]),\n",
" tensor([ 1, 278, 289, 4992, 12544])]"
]
},
"execution_count": 43,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# initital state\n",
"prompt = \"The necromancer in his tower, what's his top problem? \"\n",
"choices = [\"1\", \"the skeleton\", \"the boney boys\"]\n",
"choices_tokens = tokenizer(choices).input_ids\n",
"choices_tokens = [torch.tensor(c) for c in choices_tokens]\n",
"# current_tokens = torch.tensor([])\n",
"\n",
"# next\n",
"input_ids = tokenizer([prompt], return_tensors=\"pt\").to(model.device).input_ids[0]\n",
"choices_tokens\n",
"\n",
"# for each next choice, continue down the tree, recording the log probs"
]
},
{
"cell_type": "code",
"execution_count": 70,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Generating...\n",
"{\n",
" store: {\n",
" name: \u001b[32m\"Mike's Ski Shop\"\u001b[0m,\n",
" location: \u001b[32m\"Sugarloaf\"\u001b[0m,\n",
" inventory: [\n",
" {\n",
" productId: \u001b[32m\"1234567890\"\u001b[0m,\n",
" name: \u001b[32m\"Ski Shop\"\u001b[0m,\n",
" description: \u001b[32m\"Ski Shop sells premium skis and snowboards\"\u001b[0m,\n",
" category: \u001b[32m\"Sports\"\u001b[0m,\n",
" price: \u001b[32m20.09999\u001b[0m,\n",
" inStock: \u001b[32mTrue\u001b[0m,\n",
" rating: \u001b[32m5.09999\u001b[0m,\n",
" images: [\n",
" \u001b[32m\"https://s3.amazonaws.com/images.skisnow.com/skis\"\u001b[0m\n",
" ]\n",
" },\n",
" {\n",
" productId: \u001b[32m\"12345678910\"\u001b[0m,\n",
" name: \u001b[32m\"Snowboard\"\u001b[0m,\n",
" description: \u001b[32m\"Snowboard for sale\"\u001b[0m,\n",
" category: \u001b[32m\"Sports\"\u001b[0m,\n",
" price: \u001b[32m20.09999\u001b[0m,\n",
" inStock: \u001b[32mTrue\u001b[0m,\n",
" rating: \u001b[32m5.09999\u001b[0m,\n",
" images: [\n",
" \u001b[32m\"https://s3.amazonaws.com/images.skisnow.com/snow\"\u001b[0m\n",
" ]\n",
" }\n",
" ]\n",
" }\n",
"}\n"
"tensor([1]) current_tokens\n",
"tensor([ 1, 278]) current_tokens\n",
"tensor([ 1, 278, 289]) current_tokens\n",
"tensor([ 1, 278, 289, 4992]) current_tokens\n",
"tensor([ 1, 278, 289, 4992, 12544]) current_tokens\n",
"tensor([ 1, 278, 18109]) current_tokens\n",
"tensor([ 1, 278, 18109, 11285]) current_tokens\n",
"tensor([ 1, 29871]) current_tokens\n",
"tensor([ 1, 29871, 29896]) current_tokens\n"
]
}
],
"source": [
"from jsonformer.format import highlight_values\n",
"from jsonformer.main import Jsonformer\n",
"def get_valid_next_choices(choices_tokens, current_tokens):\n",
" next_choices = []\n",
" for choice_tokens in choices_tokens:\n",
" # if we have some more slots left\n",
" if len(current_tokens)<len(choice_tokens):\n",
" # see if current_tokens matches\n",
" if (choice_tokens[: len(current_tokens)] == current_tokens).all():\n",
" c = choice_tokens[len(current_tokens)].item()\n",
" next_choices.append(c)\n",
"\n",
" next_choices = list(set(next_choices))\n",
" return torch.LongTensor(next_choices)\n",
"\n",
"\n",
"def next(\n",
" input_ids: Int[Tensor, \"seq\"],\n",
" choice: Optional[Int[Tensor, \"\"]] = None,\n",
" prob: float = 1,\n",
" current_tokens: Int[Tensor, \"seq\"] = torch.LongTensor([]),\n",
" z=[],\n",
"):\n",
" if choice is not None:\n",
" c = choice[None].to(current_tokens.device)\n",
" current_tokens = torch.cat([current_tokens, c], dim=-1)\n",
" print(current_tokens, 'current_tokens')\n",
" c = choice[None].to(input_ids.device)\n",
" input_ids = torch.cat([input_ids, c], dim=-1)\n",
"\n",
" next_choices = get_valid_next_choices(choices_tokens, current_tokens)\n",
" if len(next_choices) == 0:\n",
" s = tokenizer.decode(current_tokens)\n",
" r = dict(tokens=current_tokens.cpu(), prob=prob, choice=s)\n",
" yield r\n",
" else:\n",
" o = model(input_ids[None])\n",
" logits_constrained = o.logits[0, -1][next_choices]\n",
" probs = F.softmax(logits_constrained, dim=-1)\n",
" for i in range(len(next_choices)):\n",
" next_choice = next_choices[i]\n",
" next_prob = prob * probs[i].item()\n",
" yield from next(\n",
" input_ids=input_ids,\n",
" choice=next_choice,\n",
" prob=next_prob,\n",
" current_tokens=current_tokens,\n",
" z=z + [i],\n",
" )\n",
"\n",
"\n",
"r = list(next(input_ids=input_ids))"
]
},
{
"cell_type": "code",
"execution_count": 74,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>prob</th>\n",
" <th>choice</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>2</th>\n",
" <td>0.995732</td>\n",
" <td>&lt;s&gt; 1</td>\n",
" </tr>\n",
" <tr>\n",
" <th>0</th>\n",
" <td>0.004187</td>\n",
" <td>&lt;s&gt; the boney boys</td>\n",
" </tr>\n",
" <tr>\n",
" <th>1</th>\n",
" <td>0.000081</td>\n",
" <td>&lt;s&gt; the skeleton</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"text/plain": [
" prob choice\n",
"2 0.995732 <s> 1\n",
"0 0.004187 <s> the boney boys\n",
"1 0.000081 <s> the skeleton"
]
},
"execution_count": 74,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"pd.DataFrame(r).sort_values(\"prob\", ascending=False).drop(columns=[\"tokens\"])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Continue"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from prob_jsonformer.format import highlight_values\n",
"from prob_jsonformer.main import Jsonformer\n",
"\n",
"ecomm = {\n",
" \"type\": \"object\",\n",
@@ -127,25 +315,9 @@
},
{
"cell_type": "code",
"execution_count": 3,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Generating...\n",
"{\n",
" make: \u001b[32m\"audi\"\u001b[0m,\n",
" model: \u001b[32m\"a4\"\u001b[0m,\n",
" year: \u001b[32m2016.0\u001b[0m,\n",
" colors_available: [\n",
" \u001b[32m\"blue\"\u001b[0m\n",
" ]\n",
"}\n"
]
}
],
"outputs": [],
"source": [
"car = {\n",
" \"type\": \"object\",\n",
@@ -170,56 +342,66 @@
"print(\"Generating...\")\n",
"output = builder()\n",
"\n",
"highlight_values(output)\n"
"highlight_values(output)"
]
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Generating...\n",
"{\n",
" car: {\n",
" make: \u001b[32m\"Rolls Royce\"\u001b[0m,\n",
" model: \u001b[32m\"Phantom\"\u001b[0m,\n",
" year: \u001b[32m2016.0\u001b[0m,\n",
" colors: [\n",
" \u001b[32m\"Gold\"\u001b[0m\n",
" ],\n",
" features: {\n",
" audio: {\n",
" brand: \u001b[32m\"Mercedes\"\u001b[0m,\n",
" speakers: \u001b[32m2.0\u001b[0m,\n",
" hasBluetooth: \u001b[32mTrue\u001b[0m\n",
" },\n",
" safety: {\n",
" airbags: \u001b[32m2.0\u001b[0m,\n",
" parkingSensors: \u001b[32mTrue\u001b[0m,\n",
" laneAssist: \u001b[32mTrue\u001b[0m\n",
" },\n",
" performance: {\n",
" engine: \u001b[32m\"Mercedes-Benz OM615\"\u001b[0m,\n",
" horsepower: \u001b[32m350.0\u001b[0m,\n",
" topSpeed: \u001b[32m220.0\u001b[0m\n",
" }\n",
" }\n",
" },\n",
" owner: {\n",
" firstName: \u001b[32m\"John\"\u001b[0m,\n",
" lastName: \u001b[32m\"Doe\"\u001b[0m,\n",
" age: \u001b[32m40.0\u001b[0m\n",
" }\n",
"}\n"
]
}
],
"outputs": [],
"source": [
"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\"}}}}}\n",
"complex_car = {\n",
" \"type\": \"object\",\n",
" \"properties\": {\n",
" \"car\": {\n",
" \"type\": \"object\",\n",
" \"properties\": {\n",
" \"make\": {\"type\": \"string\"},\n",
" \"model\": {\"type\": \"string\"},\n",
" \"year\": {\"type\": \"number\"},\n",
" \"colors\": {\"type\": \"array\", \"items\": {\"type\": \"string\"}},\n",
" \"features\": {\n",
" \"type\": \"object\",\n",
" \"properties\": {\n",
" \"audio\": {\n",
" \"type\": \"object\",\n",
" \"properties\": {\n",
" \"brand\": {\"type\": \"string\"},\n",
" \"speakers\": {\"type\": \"number\"},\n",
" \"hasBluetooth\": {\"type\": \"boolean\"},\n",
" },\n",
" },\n",
" \"safety\": {\n",
" \"type\": \"object\",\n",
" \"properties\": {\n",
" \"airbags\": {\"type\": \"number\"},\n",
" \"parkingSensors\": {\"type\": \"boolean\"},\n",
" \"laneAssist\": {\"type\": \"boolean\"},\n",
" },\n",
" },\n",
" \"performance\": {\n",
" \"type\": \"object\",\n",
" \"properties\": {\n",
" \"engine\": {\"type\": \"string\"},\n",
" \"horsepower\": {\"type\": \"number\"},\n",
" \"topSpeed\": {\"type\": \"number\"},\n",
" },\n",
" },\n",
" },\n",
" },\n",
" },\n",
" },\n",
" \"owner\": {\n",
" \"type\": \"object\",\n",
" \"properties\": {\n",
" \"firstName\": {\"type\": \"string\"},\n",
" \"lastName\": {\"type\": \"string\"},\n",
" \"age\": {\"type\": \"number\"},\n",
" },\n",
" },\n",
" },\n",
"}\n",
"builder = Jsonformer(\n",
" model=model,\n",
" tokenizer=tokenizer,\n",
@@ -230,7 +412,7 @@
"print(\"Generating...\")\n",
"output = builder()\n",
"\n",
"highlight_values(output)\n"
"highlight_values(output)"
]
}
],
@@ -250,7 +432,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.10"
"version": "3.9.16"
},
"orig_nbformat": 4
},
jsonformer/__init__.py
+2 -2
View File
@@ -1,2 +1,2 @@
from jsonformer.main import Jsonformer
from jsonformer.format import highlight_values
from prob_jsonformer.main import prob_jsonformer
from prob_jsonformer.format import highlight_values
jsonformer/main.py
+14 -9
View File
@@ -1,6 +1,6 @@
from typing import List, Union, Dict, Any
from jsonformer.logits_processors import (
from prob_jsonformer.logits_processors import (
NumberStoppingCriteria,
OutputNumbersTokens,
StringStoppingCriteria,
@@ -12,7 +12,7 @@ import json
GENERATION_MARKER = "|GENERATION|"
class Jsonformer:
class prob_jsonformer:
value: Dict[str, Any] = {}
def __init__(
@@ -80,9 +80,11 @@ class Jsonformer:
if iterations > 3:
raise ValueError("Failed to generate a valid number")
return self.generate_number(temperature=self.temperature * 1.3, iterations=iterations+1)
return self.generate_number(
temperature=self.temperature * 1.3, iterations=iterations + 1
)
def generate_boolean(self) -> bool:
def generate_boolean(self, prob=False) -> bool:
prompt = self.get_prompt()
self.debug("[generate_boolean]", prompt, is_prompt=True)
@@ -90,13 +92,17 @@ class Jsonformer:
output = self.model.forward(input_tensor.to(self.model.device))
logits = output.logits[0, -1]
# todo: this assumes that "true" and "false" are both tokenized to a single token
# TODO: this assumes that "true" and "false" are both tokenized to a single token
# this is probably not true for all tokenizers
# this can be fixed by looking at only the first token of both "true" and "false"
# FIXME: consdier " True", " true", "\ntrue" etc
true_token_id = self.tokenizer.convert_tokens_to_ids("true")
false_token_id = self.tokenizer.convert_tokens_to_ids("false")
result = logits[true_token_id] > logits[false_token_id]
if prob:
result = dict(true=logits[true_token_id], false=logits[false_token_id])
else:
result = logits[true_token_id] > logits[false_token_id]
self.debug("[generate_boolean]", result)
@@ -199,7 +205,6 @@ class Jsonformer:
output = self.model.forward(input_tensor.to(self.model.device))
logits = output.logits[0, -1]
top_indices = logits.topk(30).indices
sorted_token_ids = top_indices[logits[top_indices].argsort(descending=True)]
@@ -208,10 +213,10 @@ class Jsonformer:
for token_id in sorted_token_ids:
decoded_token = self.tokenizer.decode(token_id)
if ',' in decoded_token:
if "," in decoded_token:
found_comma = True
break
if ']' in decoded_token:
if "]" in decoded_token:
found_close_bracket = True
break
pyproject.toml
+5 -3
View File
@@ -1,13 +1,15 @@
[tool.poetry]
name = "jsonformer"
name = "prob_jsonformer"
version = "0.12.0"
description = ""
authors = ["1rgs <rgsduke@gmail.com>"]
readme = "README.md"
[tool.poetry.dependencies]
python = "^3.8"
python = "^3.9,<4.0"
termcolor = "^2.3.0"
jaxtyping = "^0.2.28"
pandas = "^2.2.2"
[tool.poetry.group.dev.dependencies]
ipykernel = "^6.22.0"
@@ -23,4 +25,4 @@ build-backend = "poetry.core.masonry.api"
[virtualenvs]
create = true
in-project = true
in-project = true