wassname/prob_jsonformer
1
0
Fork 0
mirror of https://github.com/wassname/prob_jsonformer.git synced 2026-08-20 12:40:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix, update readme and example

Browse Source
This commit is contained in:
wassname
2024-05-10 20:52:30 +08:00
parent 5a68bc5345
commit ef99113aa3
4 changed files with 211 additions and 185 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+39 -22
View File
@@ -1,50 +1,67 @@
# 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.
This fork has been modified to include the token probabilities. The original [README](https://github.com/1rgs/jsonformer) is includesbelow.
I've also merged some of hte recent pr's for enum, integer, null, union. You can see them all below in this example:
## Example
```python
from jsonformer import Jsonformer
from prob_jsonformer import Jsonformer
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("databricks/dolly-v2-12b")
tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b")
model_name = "databricks/dolly-v2-3b"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
json_schema = {
"type": "object",
"properties": {
"name": {"type": "string"},
"age": {"type": "choice_probs", "enum": ["8", "9", "10", "11"]},
"age2": {"type": "number"},
"age": {"type": "integer"},
"age_probs": {"type": "choice_probs", "enum": [str(s) for s in range(10, 20)]},
"unit_time": {"type": "number"},
"is_student": {"type": "boolean"},
"is_student2": {"type": "choice_probs", "enum": ["true", "false"]},
"is_student_probs": {"type": "choice_probs", "enum": ["true", "false"]},
"courses": {
"type": "array",
"items": {"type": "string"}
}
},
"trim": {"type": ["string", "null"]},
"color": {
"type": "enum",
"values": ["red", "green", "blue", "brown", "white", "black"],
},
}
}
prompt = "Generate a person's information based on the following schema:"
jsonformer = Jsonformer(model, tokenizer, json_schema, prompt)
prompt = "Generate a young person's information based on the following schema:"
jsonformer = Jsonformer(model, tokenizer, json_schema, prompt, temperature=0)
generated_data = jsonformer()
print(generated_data)
# {'name': 'John Doe',
# 'age': [{'prob': 0.1497802734375, 'choice': '8'},
# {'prob': 0.159423828125, 'choice': '9'},
# {'prob': 0.0982666015625, 'choice': '11'},
# {'prob': 0.59228515625, 'choice': '10'}],
# 'age2': 10.0201,
# 'is_student': True,
# 'is_student2': [{'prob': 0.94580078125, 'choice': 'true'},
# {'prob': 0.05419921875, 'choice': 'false'}],
# 'courses': ['C++']}
generated_data = {'name': 'John Doe',
'age': 20,
'age_probs': [{'prob': 0.794921875, 'choice': '12'},
{'prob': 0.068359375, 'choice': '10'},
{'prob': 0.04345703125, 'choice': '16'},
{'prob': 0.03228759765625, 'choice': '14'},
{'prob': 0.0175628662109375, 'choice': '11'},
{'prob': 0.0157318115234375, 'choice': '15'},
{'prob': 0.006664276123046875, 'choice': '18'},
{'prob': 0.0046539306640625, 'choice': '13'},
{'prob': 0.00041294097900390625, 'choice': '17'},
{'prob': 0.00028824806213378906, 'choice': '19'},,
'unit_time': 0.01,
'is_student': True,
'is_student_probs': [{'prob': 0.8310546875, 'choice': 'true'},
{'prob': 0.1688232421875, 'choice': 'false'}],
'courses': ['C1'],
'trim': None,
'color': 'white'}
```
# Jsonformer: A Bulletproof Way to Generate Structured JSON from Language Models.
# ORIGINAL: Jsonformer: A Bulletproof Way to Generate Structured JSON from Language Models.
### Problem: Getting models to output structured JSON is hard
example.ipynb
+130 -145
View File
@@ -111,12 +111,36 @@
" name: \u001b[32m\"Snowboard X-15\"\u001b[0m,\n",
" description: \u001b[32m\"Snowboard for all levels\"\u001b[0m,\n",
" category: \u001b[32m\"Snowboards\"\u001b[0m,\n",
" price: \u001b[32m20.0375\u001b[0m,\n",
" price: \u001b[32m20.0\u001b[0m,\n",
" inStock: \u001b[32mTrue\u001b[0m,\n",
" rating: \u001b[32m5.0\u001b[0m,\n",
" images: [\n",
" \u001b[32m\"https://s3.amazonaws.com/mikesskisport/images/Snow\"\u001b[0m\n",
" ]\n",
" },\n",
" {\n",
" productId: \u001b[32m\"2\"\u001b[0m,\n",
" name: \u001b[32m\"Mike's Ski Shop Exclusive\"\u001b[0m,\n",
" description: \u001b[32m\"Mike's Ski Shop Exclusive\"\u001b[0m,\n",
" category: \u001b[32m\"Ski Shops\"\u001b[0m,\n",
" price: \u001b[32m20.0\u001b[0m,\n",
" inStock: \u001b[32mTrue\u001b[0m,\n",
" rating: \u001b[32m5.0\u001b[0m,\n",
" images: [\n",
" \u001b[32m\"https://s3.amazonaws.com/mikesskisport/images/Mike\"\u001b[0m\n",
" ]\n",
" },\n",
" {\n",
" productId: \u001b[32m\"3\"\u001b[0m,\n",
" name: \u001b[32m\"Mike's Ski Shop Exclusive\"\u001b[0m,\n",
" description: \u001b[32m\"Mike's Ski Shop Exclusive\"\u001b[0m,\n",
" category: \u001b[32m\"Ski Shops\"\u001b[0m,\n",
" price: \u001b[32m20.0\u001b[0m,\n",
" inStock: \u001b[32mTrue\u001b[0m,\n",
" rating: \u001b[32m5.0\u001b[0m,\n",
" images: [\n",
" \u001b[32m\"https://s3.amazonaws.com/mikesskisport/images/Mike\"\u001b[0m\n",
" ]\n",
" }\n",
" ]\n",
" }\n",
@@ -175,62 +199,9 @@
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Generating...\n",
"{\n",
" make: \u001b[32m\"Mazda\"\u001b[0m,\n",
" model: [\n",
" {\n",
" prob: \u001b[32m0.8154296875\u001b[0m,\n",
" choice: \u001b[32m\"Kea\"\u001b[0m\n",
" },\n",
" {\n",
" prob: \u001b[32m0.184814453125\u001b[0m,\n",
" choice: \u001b[32m\"Mazda\"\u001b[0m\n",
" }\n",
" ],\n",
" new: [\n",
" {\n",
" prob: \u001b[32m0.90185546875\u001b[0m,\n",
" choice: \u001b[32m\"true\"\u001b[0m\n",
" },\n",
" {\n",
" prob: \u001b[32m0.09808349609375\u001b[0m,\n",
" choice: \u001b[32m\"false\"\u001b[0m\n",
" }\n",
" ],\n",
" rating: [\n",
" {\n",
" prob: \u001b[32m0.221435546875\u001b[0m,\n",
" choice: \u001b[32m\"1\"\u001b[0m\n",
" },\n",
" {\n",
" prob: \u001b[32m0.394775390625\u001b[0m,\n",
" choice: \u001b[32m\"2\"\u001b[0m\n",
" },\n",
" {\n",
" prob: \u001b[32m0.382568359375\u001b[0m,\n",
" choice: \u001b[32m\"3\"\u001b[0m\n",
" },\n",
" {\n",
" prob: \u001b[32m0.0013370513916015625\u001b[0m,\n",
" choice: \u001b[32m\"4\"\u001b[0m\n",
" }\n",
" ],\n",
" year: \u001b[32m2016.0\u001b[0m,\n",
" colors_available: [\n",
" \u001b[32m\"red\"\u001b[0m\n",
" ]\n",
"}\n"
]
}
],
"outputs": [],
"source": [
"car = {\n",
" \"type\": \"object\",\n",
@@ -262,96 +233,9 @@
},
{
"cell_type": "code",
"execution_count": 6,
"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[32m2014.0\u001b[0m,\n",
" colors: [\n",
" {\n",
" prob: \u001b[32m0.001560211181640625\u001b[0m,\n",
" choice: \u001b[32m\"white\"\u001b[0m\n",
" },\n",
" {\n",
" prob: \u001b[32m0.833984375\u001b[0m,\n",
" choice: \u001b[32m\"red\"\u001b[0m\n",
" },\n",
" {\n",
" prob: \u001b[32m0.0865478515625\u001b[0m,\n",
" choice: \u001b[32m\"black\"\u001b[0m\n",
" },\n",
" {\n",
" prob: \u001b[32m0.048553466796875\u001b[0m,\n",
" choice: \u001b[32m\"blue\"\u001b[0m\n",
" },\n",
" {\n",
" prob: \u001b[32m0.0294342041015625\u001b[0m,\n",
" choice: \u001b[32m\"green\"\u001b[0m\n",
" }\n",
" ],\n",
" as_new: [\n",
" {\n",
" prob: \u001b[32m0.96533203125\u001b[0m,\n",
" choice: \u001b[32m\"true\"\u001b[0m\n",
" },\n",
" {\n",
" prob: \u001b[32m0.03460693359375\u001b[0m,\n",
" choice: \u001b[32m\"false\"\u001b[0m\n",
" }\n",
" ],\n",
" rating: [\n",
" {\n",
" prob: \u001b[32m0.05462646484375\u001b[0m,\n",
" choice: \u001b[32m\"1\"\u001b[0m\n",
" },\n",
" {\n",
" prob: \u001b[32m0.233642578125\u001b[0m,\n",
" choice: \u001b[32m\"2\"\u001b[0m\n",
" },\n",
" {\n",
" prob: \u001b[32m0.2093505859375\u001b[0m,\n",
" choice: \u001b[32m\"3\"\u001b[0m\n",
" },\n",
" {\n",
" prob: \u001b[32m0.50244140625\u001b[0m,\n",
" choice: \u001b[32m\"4\"\u001b[0m\n",
" }\n",
" ],\n",
" features: {\n",
" audio: {\n",
" brand: \u001b[32m\"Mercedes-Benz\"\u001b[0m,\n",
" speakers: \u001b[32m2.09999\u001b[0m,\n",
" hasBluetooth: \u001b[32mTrue\u001b[0m\n",
" },\n",
" safety: {\n",
" airbags: \u001b[32m2.09999\u001b[0m,\n",
" parkingSensors: \u001b[32mTrue\u001b[0m,\n",
" laneAssist: \u001b[32mTrue\u001b[0m\n",
" },\n",
" performance: {\n",
" engine: \u001b[32m\"Mercedes-Benz 6.2 L\"\u001b[0m,\n",
" horsepower: \u001b[32m423.09999\u001b[0m,\n",
" topSpeed: \u001b[32m220.09999\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.09999\u001b[0m\n",
" }\n",
"}\n"
]
}
],
"outputs": [],
"source": [
"complex_car = {\n",
" \"type\": \"object\",\n",
@@ -418,6 +302,107 @@
"\n",
"highlight_values(output)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Readme example"
]
},
{
"cell_type": "code",
"execution_count": 22,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'name': 'John Doe',\n",
" 'age': 20,\n",
" 'age_probs': [{'prob': 0.794921875, 'choice': '12'},\n",
" {'prob': 0.068359375, 'choice': '10'},\n",
" {'prob': 0.04345703125, 'choice': '16'},\n",
" {'prob': 0.03228759765625, 'choice': '14'},\n",
" {'prob': 0.0175628662109375, 'choice': '11'},\n",
" {'prob': 0.0157318115234375, 'choice': '15'},\n",
" {'prob': 0.013671875, 'choice': '20'},\n",
" {'prob': 0.006664276123046875, 'choice': '18'},\n",
" {'prob': 0.0046539306640625, 'choice': '13'},\n",
" {'prob': 0.0018215179443359375, 'choice': '21'},\n",
" {'prob': 0.00041294097900390625, 'choice': '17'},\n",
" {'prob': 0.00028824806213378906, 'choice': '19'},\n",
" {'prob': 0.00014495849609375, 'choice': '22'},\n",
" {'prob': 6.955862045288086e-05, 'choice': '23'},\n",
" {'prob': 2.968311309814453e-05, 'choice': '25'},\n",
" {'prob': 2.8789043426513672e-05, 'choice': '26'},\n",
" {'prob': 1.901388168334961e-05, 'choice': '24'},\n",
" {'prob': 1.1742115020751953e-05, 'choice': '28'},\n",
" {'prob': 1.1920928955078125e-06, 'choice': '27'},\n",
" {'prob': 7.748603820800781e-07, 'choice': '29'}],\n",
" 'unit_time': 0.01,\n",
" 'is_student': True,\n",
" 'is_student_probs': [{'prob': 0.8310546875, 'choice': 'true'},\n",
" {'prob': 0.1688232421875, 'choice': 'false'}],\n",
" 'courses': ['C1'],\n",
" 'trim': None,\n",
" 'color': 'white'}"
]
},
"execution_count": 22,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from prob_jsonformer import Jsonformer\n",
"from transformers import AutoModelForCausalLM, AutoTokenizer\n",
"\n",
"model_name = \"databricks/dolly-v2-3b\"\n",
"model = AutoModelForCausalLM.from_pretrained(model_name)\n",
"tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
"\n",
"json_schema = {\n",
" \"type\": \"object\",\n",
" \"properties\": {\n",
" \"name\": {\"type\": \"string\"},\n",
" \"age\": {\"type\": \"integer\"},\n",
" \"age_probs\": {\"type\": \"choice_probs\", \"enum\": [str(s) for s in range(10, 30)]},\n",
" \"unit_time\": {\"type\": \"number\"},\n",
" \"is_student\": {\"type\": \"boolean\"},\n",
" \"is_student_probs\": {\"type\": \"choice_probs\", \"enum\": [\"true\", \"false\"]},\n",
" \"courses\": {\n",
" \"type\": \"array\",\n",
" \"items\": {\"type\": \"string\"}\n",
" },\n",
" \"trim\": {\"type\": [\"string\", \"null\"]},\n",
" \"color\": {\n",
" \"type\": \"enum\",\n",
" \"values\": [\"red\", \"green\", \"blue\", \"brown\", \"white\", \"black\"],\n",
" },\n",
" }\n",
"}\n",
"\n",
"prompt = \"Generate a young person's information based on the following schema:\"\n",
"jsonformer = Jsonformer(model, tokenizer, json_schema, prompt, temperature=0)\n",
"generated_data = jsonformer()\n",
"\n",
"generated_data"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
prob_jsonformer/choice_tree.py
+16 -3
View File
@@ -21,7 +21,7 @@ def get_valid_next_choices(choices_tokens, current_tokens):
return torch.LongTensor(next_choices)
def choice_tree(
def _choice_tree(
model: AutoModelForCausalLM,
tokenizer: AutoTokenizer,
input_ids: Int[Tensor, "seq"],
@@ -29,7 +29,6 @@ def choice_tree(
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)
@@ -57,5 +56,19 @@ def choice_tree(
choice=next_choice,
prob=next_prob,
current_tokens=current_tokens,
z=z + [i],
)
def choice_tree(
*args,
**kwargs,
):
choice_json = list(
_choice_tree(
*args,
**kwargs,
)
)
# order by probability
choice_json = sorted(choice_json, key=lambda x: -x["prob"])
return choice_json
prob_jsonformer/main.py
+26 -15
View File
@@ -138,10 +138,7 @@ class Jsonformer:
true_token_id = self.tokenizer.encode("true", return_tensors="pt")[0, 0]
false_token_id = self.tokenizer.encode("false", return_tensors="pt")[0, 0]
if prob:
result = dict(true=logits[true_token_id], false=logits[false_token_id])
else:
result = logits[true_token_id] > logits[false_token_id]
result = logits[true_token_id] > logits[false_token_id]
self.debug("[generate_boolean]", result)
@@ -201,24 +198,38 @@ class Jsonformer:
self.debug("[generate_enum]", prompt, is_prompt=True)
# These are necessary because we don't know if we're at the end or middle of an object/array
terminal_tokens = torch.concat([
self.tokenizer.encode(s, add_special_tokens=False, return_tensors="pt")[:, 0]
for s in ('", "', '"}', '"]')
])
terminal_tokens = torch.concat(
[
self.tokenizer.encode(s, add_special_tokens=False, return_tensors="pt")[
:, 0
]
for s in ('", "', '"}', '"]')
]
)
highest_probability = 0.0
best_option = None
for option in enum_values:
n_option_tokens = self.tokenizer.encode(f'"{option}', add_special_tokens=False, return_tensors="pt").shape[1]
prompt_tokens = self.tokenizer.encode(prompt + f'"{option}', return_tensors="pt")
n_option_tokens = self.tokenizer.encode(
f'"{option}', add_special_tokens=False, return_tensors="pt"
).shape[1]
prompt_tokens = self.tokenizer.encode(
prompt + f'"{option}', return_tensors="pt"
)
option_tokens = prompt_tokens[0, -n_option_tokens:]
with torch.no_grad():
logits = self.model.forward(prompt_tokens.to(self.model.device)).logits[0, -n_option_tokens-1:]
logits = self.model.forward(prompt_tokens.to(self.model.device)).logits[
0, -n_option_tokens - 1 :
]
probabilities = torch.softmax(logits, dim=1)
option_token_probabilities = probabilities[:-1][torch.arange(probabilities.shape[0]-1), option_tokens]
option_token_probabilities = probabilities[:-1][
torch.arange(probabilities.shape[0] - 1), option_tokens
]
termination_probability = torch.max(probabilities[-1, terminal_tokens])
option_probability = torch.prod(option_token_probabilities) * termination_probability
option_probability = (
torch.prod(option_token_probabilities) * termination_probability
)
if option_probability > highest_probability:
best_option = option
@@ -252,7 +263,7 @@ class Jsonformer:
max_type = None
max_logit = -float("inf")
for possible_type in possible_types:
try:
try:
prefix_tokens = self.type_prefix_tokens[possible_type]
except KeyError:
raise ValueError(f"Unsupported schema type: {possible_type}")
@@ -265,7 +276,7 @@ class Jsonformer:
raise Exception("Unable to find best type to generate for union type")
self.debug("[choose_type_to_generate]", max_type)
return max_type
def generate_value(
self,
schema: Dict[str, Any],