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

262 lines
12 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"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "c7620cdf45c54ab5abaf479b141f4479",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"config.json: 0%| | 0.00/819 [00:00<?, ?B/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "b9dba56663274820997ef5a07e7acea4",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"pytorch_model.bin: 0%| | 0.00/5.68G [00:00<?, ?B/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "8571f2926f9541469ed34140333cdf81",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"model.safetensors: 0%| | 0.00/5.68G [00:00<?, ?B/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "93c30a28fbdd4257ab783cc659b24687",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"tokenizer_config.json: 0%| | 0.00/450 [00:00<?, ?B/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "10ad2621da06464bbdc5b4aa49687d50",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"tokenizer.json: 0%| | 0.00/2.11M [00:00<?, ?B/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "3abe2819228e4da4ad859c54abefc479",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"special_tokens_map.json: 0%| | 0.00/228 [00:00<?, ?B/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"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": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"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 29\u001b[0m\n\u001b[1;32m 27\u001b[0m prompt \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mGenerate a young person\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124ms information based on the following schema:\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 28\u001b[0m jsonformer \u001b[38;5;241m=\u001b[39m Jsonformer(model, tokenizer, json_schema, prompt)\n\u001b[0;32m---> 29\u001b[0m generated_data \u001b[38;5;241m=\u001b[39m \u001b[43mjsonformer\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 31\u001b[0m generated_data\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: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 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\", \"enum\": [str(s) for s in range(10, 30)]},\n",
" # we can return the probabilistic weighted mean of a range\n",
" \"age_wmean\": {\"type\": \"p_integer\", \"minimum\": 10, \"maximum\": 30},\n",
" # the prob of true and false\n",
" \"is_student_probs\": {\"type\": \"p_enum\", \"enum\": [\"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",
"\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": 4,
"metadata": {},
"outputs": [],
"source": [
"generated_data = {\n",
" \"age_probs\": [\n",
" {\"prob\": 0.94091796875, \"choice\": \"10\"},\n",
" {\"prob\": 0.033233642578125, \"choice\": \"20\"},\n",
" {\"prob\": 0.0122222900390625, \"choice\": \"12\"},\n",
" {\"prob\": 0.00412750244140625, \"choice\": \"21\"},\n",
" {\"prob\": 0.0028362274169921875, \"choice\": \"16\"},\n",
" {\"prob\": 0.0018453598022460938, \"choice\": \"15\"},\n",
" {\"prob\": 0.00113677978515625, \"choice\": \"11\"},\n",
" {\"prob\": 0.0011110305786132812, \"choice\": \"18\"},\n",
" {\"prob\": 0.0005083084106445312, \"choice\": \"25\"},\n",
" {\"prob\": 0.0004558563232421875, \"choice\": \"23\"},\n",
" {\"prob\": 0.0002498626708984375, \"choice\": \"14\"},\n",
" {\"prob\": 0.00023281574249267578, \"choice\": \"13\"},\n",
" {\"prob\": 0.0002238750457763672, \"choice\": \"22\"},\n",
" {\"prob\": 0.00018131732940673828, \"choice\": \"26\"},\n",
" {\"prob\": 0.0001690387725830078, \"choice\": \"24\"},\n",
" {\"prob\": 0.00012552738189697266, \"choice\": \"19\"},\n",
" {\"prob\": 7.796287536621094e-05, \"choice\": \"27\"},\n",
" {\"prob\": 7.265806198120117e-05, \"choice\": \"28\"},\n",
" {\"prob\": 4.106760025024414e-05, \"choice\": \"17\"},\n",
" {\"prob\": 2.5033950805664062e-06, \"choice\": \"29\"},\n",
" ],\n",
" \"age_wmean\": 17.816404402256012,\n",
" \"is_student_probs\": [\n",
" {\"prob\": 0.974609375, \"choice\": \"true\"},\n",
" {\"prob\": 0.025177001953125, \"choice\": \"false\"},\n",
" ],\n",
" \"is_student\": False,\n",
" \"name\": \"John\",\n",
" \"age\": 17,\n",
" \"unit_time\": 0.5,\n",
" \"courses\": [\"C++\"],\n",
" \"trim\": None,\n",
" \"color\": \"white\",\n",
"}"
]
},
{
"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"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
Reference in New Issue View Git Blame Copy Permalink