mirror of
https://github.com/wassname/prob_jsonformer.git
synced 2026-09-12 12:51:29 +08:00
add example and fix
This commit is contained in:
wassname
@@ -2,6 +2,48 @@
|
||||
|
||||
This fork has been modified to include the token probabilities. The original README is below.
|
||||
|
||||
|
||||
## Example
|
||||
|
||||
```python
|
||||
from jsonformer import Jsonformer
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained("databricks/dolly-v2-12b")
|
||||
tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b")
|
||||
|
||||
json_schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
"age": {"type": "choice_probs", "enum": ["8", "9", "10", "11"]},
|
||||
"age2": {"type": "number"},
|
||||
"is_student": {"type": "boolean"},
|
||||
"is_student2": {"type": "choice_probs", "enum": ["true", "false"]},
|
||||
"courses": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
prompt = "Generate a person's information based on the following schema:"
|
||||
jsonformer = Jsonformer(model, tokenizer, json_schema, prompt)
|
||||
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++']}
|
||||
```
|
||||
|
||||
# Jsonformer: A Bulletproof Way to Generate Structured JSON from Language Models.
|
||||
|
||||
### Problem: Getting models to output structured JSON is hard
|
||||
|
||||
+201
-11
@@ -38,7 +38,15 @@
|
||||
"/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",
|
||||
"/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"
|
||||
" 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"
|
||||
]
|
||||
}
|
||||
],
|
||||
@@ -48,15 +56,11 @@
|
||||
"\n",
|
||||
"print(\"Loading model and tokenizer...\")\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",
|
||||
" torch_dtype=torch.float16,\n",
|
||||
" # device=\"cuda:0\",\n",
|
||||
" # device_map=\"auto\",\n",
|
||||
" attn_implementation='eager',\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\")"
|
||||
@@ -71,9 +75,55 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"execution_count": 4,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"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\"Somewhere\"\u001b[0m,\n",
|
||||
" choice_probs: [\n",
|
||||
" {\n",
|
||||
" prob: \u001b[32m0.01739501953125\u001b[0m,\n",
|
||||
" choice: \u001b[32m\"pretend\"\u001b[0m\n",
|
||||
" },\n",
|
||||
" {\n",
|
||||
" prob: \u001b[32m0.002094268798828125\u001b[0m,\n",
|
||||
" choice: \u001b[32m\"snowboard\"\u001b[0m\n",
|
||||
" },\n",
|
||||
" {\n",
|
||||
" prob: \u001b[32m0.0007467269897460938\u001b[0m,\n",
|
||||
" choice: \u001b[32m\"walk\"\u001b[0m\n",
|
||||
" },\n",
|
||||
" {\n",
|
||||
" prob: \u001b[32m0.97998046875\u001b[0m,\n",
|
||||
" choice: \u001b[32m\"ski\"\u001b[0m\n",
|
||||
" }\n",
|
||||
" ],\n",
|
||||
" inventory: [\n",
|
||||
" {\n",
|
||||
" productId: \u001b[32m\"1\"\u001b[0m,\n",
|
||||
" 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",
|
||||
" 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",
|
||||
" }\n",
|
||||
"}\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from prob_jsonformer.format import highlight_values\n",
|
||||
"from prob_jsonformer.main import Jsonformer\n",
|
||||
@@ -125,9 +175,62 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"execution_count": 5,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"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"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"car = {\n",
|
||||
" \"type\": \"object\",\n",
|
||||
@@ -159,9 +262,96 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"execution_count": 6,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"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"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"complex_car = {\n",
|
||||
" \"type\": \"object\",\n",
|
||||
|
||||
@@ -80,7 +80,7 @@ class OutputNumbersTokens(LogitsWarper):
|
||||
self.allowed_mask[token_id] = True
|
||||
|
||||
def __call__(self, _, scores):
|
||||
mask = self.allowed_mask.expand_as(scores[0])
|
||||
mask = self.allowed_mask.expand_as(scores)
|
||||
scores[~mask] = -float("inf")
|
||||
|
||||
return scores
|
||||
|
||||
Reference in New Issue
Block a user