wassname/prob_jsonformer
1
0
Fork 0
mirror of https://github.com/wassname/prob_jsonformer.git synced 2026-08-28 12:54:00 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
main
prob_jsonformer/example.ipynb
T
wassname b639079d04 stopping
2025-03-23 17:59:48 +08:00

366 lines
23 KiB
Plaintext
Raw Permalink Blame History

{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"# autoreload your package\n",
"%load_ext autoreload\n",
"%autoreload 2"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Loading model and tokenizer...\n",
"Loaded model and tokenizer\n"
]
}
],
"source": [
"from transformers import AutoModelForCausalLM, AutoTokenizer\n",
"import torch\n",
"\n",
"print(\"Loading model and tokenizer...\")\n",
"model_name = \"databricks/dolly-v2-3b\"\n",
"model = AutoModelForCausalLM.from_pretrained(\n",
" model_name,\n",
" use_cache=True,\n",
" torch_dtype=torch.float16,\n",
" attn_implementation=\"eager\",\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": [
"# Continue"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"The attention mask is not set and cannot be inferred from input because pad token is same as eos token. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Generating...\n"
]
},
{
"ename": "KeyError",
"evalue": "'values'",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mKeyError\u001b[0m Traceback (most recent call last)",
"Cell \u001b[0;32mIn[3], line 47\u001b[0m\n\u001b[1;32m 38\u001b[0m builder \u001b[38;5;241m=\u001b[39m Jsonformer(\n\u001b[1;32m 39\u001b[0m model\u001b[38;5;241m=\u001b[39mmodel,\n\u001b[1;32m 40\u001b[0m tokenizer\u001b[38;5;241m=\u001b[39mtokenizer,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 43\u001b[0m max_string_token_length\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m20\u001b[39m,\n\u001b[1;32m 44\u001b[0m )\n\u001b[1;32m 46\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mGenerating...\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m---> 47\u001b[0m output \u001b[38;5;241m=\u001b[39m \u001b[43mbuilder\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 49\u001b[0m highlight_values(output)\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/prob_jsonformer/main.py:439\u001b[0m, in \u001b[0;36mJsonformer.__call__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 437\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Dict[\u001b[38;5;28mstr\u001b[39m, Any]:\n\u001b[1;32m 438\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvalue \u001b[38;5;241m=\u001b[39m {}\n\u001b[0;32m--> 439\u001b[0m generated_data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgenerate_object\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 440\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mjson_schema\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mproperties\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvalue\u001b[49m\n\u001b[1;32m 441\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 442\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m generated_data\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/prob_jsonformer/main.py:274\u001b[0m, in \u001b[0;36mJsonformer.generate_object\u001b[0;34m(self, properties, obj)\u001b[0m\n\u001b[1;32m 272\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m key, schema \u001b[38;5;129;01min\u001b[39;00m properties\u001b[38;5;241m.\u001b[39mitems():\n\u001b[1;32m 273\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdebug(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m[generate_object] generating value for\u001b[39m\u001b[38;5;124m\"\u001b[39m, key)\n\u001b[0;32m--> 274\u001b[0m obj[key] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgenerate_value\u001b[49m\u001b[43m(\u001b[49m\u001b[43mschema\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mobj\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 275\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m obj\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/prob_jsonformer/main.py:376\u001b[0m, in \u001b[0;36mJsonformer.generate_value\u001b[0;34m(self, schema, obj, key)\u001b[0m\n\u001b[1;32m 374\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 375\u001b[0m obj\u001b[38;5;241m.\u001b[39mappend(new_obj)\n\u001b[0;32m--> 376\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgenerate_object\u001b[49m\u001b[43m(\u001b[49m\u001b[43mschema\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mproperties\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnew_obj\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 377\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m schema_type \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mnull\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[1;32m 378\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/prob_jsonformer/main.py:274\u001b[0m, in \u001b[0;36mJsonformer.generate_object\u001b[0;34m(self, properties, obj)\u001b[0m\n\u001b[1;32m 272\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m key, schema \u001b[38;5;129;01min\u001b[39;00m properties\u001b[38;5;241m.\u001b[39mitems():\n\u001b[1;32m 273\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdebug(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m[generate_object] generating value for\u001b[39m\u001b[38;5;124m\"\u001b[39m, key)\n\u001b[0;32m--> 274\u001b[0m obj[key] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgenerate_value\u001b[49m\u001b[43m(\u001b[49m\u001b[43mschema\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mobj\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 275\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m obj\n",
"File \u001b[0;32m/media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/prob_jsonformer/main.py:351\u001b[0m, in \u001b[0;36mJsonformer.generate_value\u001b[0;34m(self, schema, obj, key)\u001b[0m\n\u001b[1;32m 349\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 350\u001b[0m obj\u001b[38;5;241m.\u001b[39mappend(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgeneration_marker)\n\u001b[0;32m--> 351\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgenerate_p_enum(\u001b[43mschema\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mvalues\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m, \u001b[38;5;28mround\u001b[39m\u001b[38;5;241m=\u001b[39mschema\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mround\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;241m3\u001b[39m))\n\u001b[1;32m 352\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m schema_type \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mp_integer\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[1;32m 353\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m key:\n",
"\u001b[0;31mKeyError\u001b[0m: 'values'"
]
}
],
"source": [
"from prob_jsonformer.format import highlight_values\n",
"from prob_jsonformer.main import Jsonformer\n",
"\n",
"ecomm = {\n",
" \"type\": \"object\",\n",
" \"properties\": {\n",
" \"store\": {\n",
" \"type\": \"object\",\n",
" \"properties\": {\n",
" \"name\": {\"type\": \"string\"},\n",
" \"location\": {\"type\": \"string\"},\n",
" \"p_enum\": {\n",
" \"type\": \"p_enum\",\n",
" \"enum\": [\"ski\", \"snowboard\", \"walk\", \"pretend\"],\n",
" },\n",
" \"inventory\": {\n",
" \"type\": \"array\",\n",
" \"items\": {\n",
" \"type\": \"object\",\n",
" \"properties\": {\n",
" \"productId\": {\"type\": \"string\"},\n",
" \"name\": {\"type\": \"string\"},\n",
" \"description\": {\"type\": \"string\"},\n",
" \"category\": {\"type\": \"string\"},\n",
" \"price\": {\"type\": \"number\"},\n",
" \"inStock\": {\"type\": \"boolean\"},\n",
" \"rating\": {\"type\": \"number\"},\n",
" \"images\": {\"type\": \"array\", \"items\": {\"type\": \"string\"}},\n",
" },\n",
" },\n",
" },\n",
" },\n",
" }\n",
" },\n",
"}\n",
"\n",
"\n",
"builder = Jsonformer(\n",
" model=model,\n",
" tokenizer=tokenizer,\n",
" json_schema=ecomm,\n",
" prompt=\"write a description about mike's ski shop which sells premium skis and snowboards\",\n",
" max_string_token_length=20,\n",
")\n",
"\n",
"print(\"Generating...\")\n",
"output = builder()\n",
"\n",
"highlight_values(output)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"car = {\n",
" \"type\": \"object\",\n",
" \"properties\": {\n",
" \"make\": {\"type\": \"string\"},\n",
" \"model\": {\"type\": \"p_enum\", \"enum\": [\"Mazda\", \"Kea\"]},\n",
" \"new\": {\"type\": \"p_enum\", \"enum\": [\"true\", \"false\"]},\n",
" \"rating\": {\"type\": \"p_enum\", \"enum\": [\"1\", \"2\", \"3\", \"4\"]},\n",
" \"year\": {\"type\": \"number\"},\n",
" \"colors_available\": {\n",
" \"type\": \"array\",\n",
" \"items\": {\"type\": \"string\"},\n",
" },\n",
" },\n",
"}\n",
"\n",
"builder = Jsonformer(\n",
" model=model,\n",
" tokenizer=tokenizer,\n",
" json_schema=car,\n",
" prompt=\"generate an example car\",\n",
")\n",
"\n",
"print(\"Generating...\")\n",
"output = builder()\n",
"\n",
"highlight_values(output)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"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\": {\n",
" \"type\": \"p_enum\",\n",
" \"enum\": [\"red\", \"green\", \"blue\", \"black\", \"white\"],\n",
" },\n",
" \"as_new\": {\"type\": \"p_enum\", \"enum\": [\"true\", \"false\"]},\n",
" \"rating\": {\"type\": \"p_enum\", \"enum\": [\"1\", \"2\", \"3\", \"4\"]},\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",
" json_schema=complex_car,\n",
" prompt=\"generate an example Rolls Royce Phantom\",\n",
")\n",
"\n",
"print(\"Generating...\")\n",
"output = builder()\n",
"\n",
"highlight_values(output)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Readme example"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"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)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'age_probs': [{'prob': 0.62353515625, 'choice': '10'},\n",
" {'prob': 0.349609375, 'choice': '12'},\n",
" {'prob': 0.01123809814453125, 'choice': '11'},\n",
" {'prob': 0.00760650634765625, 'choice': '16'},\n",
" {'prob': 0.0025482177734375, 'choice': '13'},\n",
" {'prob': 0.0025081634521484375, 'choice': '15'},\n",
" {'prob': 0.0018062591552734375, 'choice': '14'},\n",
" {'prob': 0.00104522705078125, 'choice': '18'},\n",
" {'prob': 0.00011551380157470703, 'choice': '17'},\n",
" {'prob': 5.042552947998047e-05, 'choice': '19'}],\n",
" 'age_wmean': 15.544570922851562,\n",
" 'is_student_probs': [{'prob': 0.962890625, 'choice': 'true'},\n",
" {'prob': 0.037322998046875, 'choice': 'false'}],\n",
" 'is_student': False,\n",
" 'name': 'John',\n",
" 'age': 17,\n",
" 'unit_time': 0.5,\n",
" 'courses': ['C++'],\n",
" 'trim': None,\n",
" 'color': 'green'}"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from prob_jsonformer import Jsonformer\n",
"\n",
"json_schema = {\n",
" \"type\": \"object\",\n",
" \"properties\": {\n",
" # we can return the probability of each choice, even if they are multiple tokens\n",
" \"age_probs\": {\"type\": \"p_enum\", \"values\": [str(s) for s in range(10, 20)]},\n",
" # we can return the probabilistic weighted mean of a range\n",
" \"age_wmean\": {\"type\": \"p_integer\", \"minimum\": 10, \"maximum\": 20},\n",
" # the prob of true and false\n",
" \"is_student_probs\": {\"type\": \"p_enum\", \"values\": [\"true\", \"false\"]},\n",
" \"is_student\": {\"type\": \"boolean\"},\n",
" # we've merged patches for enum, integer, null, union - currently mising from jsonformer\n",
" \"name\": {\"type\": \"string\", \"maxLength\": 4},\n",
" \"age\": {\"type\": \"integer\"},\n",
" \"unit_time\": {\"type\": \"number\"},\n",
" \"courses\": {\"type\": \"array\", \"items\": {\"type\": \"string\"}},\n",
" \"trim\": {\"type\": [\"string\", \"null\"]},\n",
" \"color\": {\n",
" \"type\": \"enum\",\n",
" \"values\": [\"red\", \"green\", \"blue\", \"brown\", \"white\", \"black\"],\n",
" },\n",
" },\n",
"}\n",
"prompt = \"Generate a young person's information based on the following schema:\"\n",
"jsonformer = Jsonformer(model, tokenizer, json_schema, prompt)\n",
"generated_data = jsonformer()\n",
"\n",
"generated_data"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "pytorch_hf_env",
"language": "python",
"name": "pytorch_hf_env"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.16"
},
"orig_nbformat": 4
},
"nbformat": 4,
"nbformat_minor": 2
}
Reference in New Issue View Git Blame Copy Permalink