From f50076736fe6961808bae48d843564c5d9482988 Mon Sep 17 00:00:00 2001 From: wassname Date: Sat, 11 May 2024 11:09:43 +0800 Subject: [PATCH] rename to be closer to json schema --- README.md | 52 +-- dev.ipynb | 214 +++++++++++ example.ipynb | 805 ++++++++++++++++++---------------------- prob_jsonformer/main.py | 18 +- 4 files changed, 598 insertions(+), 491 deletions(-) create mode 100644 dev.ipynb diff --git a/README.md b/README.md index ed5867e..96a5e11 100644 --- a/README.md +++ b/README.md @@ -19,60 +19,46 @@ json_schema = { "type": "object", "properties": { # we can return the probability of each choice, even if they are multiple tokens - "age_probs": {"type": "choice_probs", "enum": [str(s) for s in range(10, 30)]}, + "age_probs": {"type": "p_enum", "values": [str(s) for s in range(10, 20)]}, # we can return the probabilistic weighted mean of a range - "age_wmean": {"type": "range_mean", "minimum": 10, "maximum": 30}, + "age_wmean": {"type": "p_integer", "minimum": 10, "maximum": 20}, # the prob of true and false - "is_student_probs": {"type": "choice_probs", "enum": ["true", "false"]}, + "is_student_probs": {"type": "p_enum", "values": ["true", "false"]}, "is_student": {"type": "boolean"}, # we've merged patches for enum, integer, null, union - currently mising from jsonformer "name": {"type": "string", "maxLength": 4}, "age": {"type": "integer"}, "unit_time": {"type": "number"}, - "courses": { - "type": "array", - "items": {"type": "string"} - }, + "courses": {"type": "array", "items": {"type": "string"}}, "trim": {"type": ["string", "null"]}, "color": { "type": "enum", "values": ["red", "green", "blue", "brown", "white", "black"], }, - } + }, } - prompt = "Generate a young person's information based on the following schema:" jsonformer = Jsonformer(model, tokenizer, json_schema, prompt, temperature=0) generated_data = jsonformer() generated_data = { "age_probs": [ - {"prob": 0.94091796875, "choice": "10"}, - {"prob": 0.033233642578125, "choice": "20"}, - {"prob": 0.0122222900390625, "choice": "12"}, - {"prob": 0.00412750244140625, "choice": "21"}, - {"prob": 0.0028362274169921875, "choice": "16"}, - {"prob": 0.0018453598022460938, "choice": "15"}, - {"prob": 0.00113677978515625, "choice": "11"}, - {"prob": 0.0011110305786132812, "choice": "18"}, - {"prob": 0.0005083084106445312, "choice": "25"}, - {"prob": 0.0004558563232421875, "choice": "23"}, - {"prob": 0.0002498626708984375, "choice": "14"}, - {"prob": 0.00023281574249267578, "choice": "13"}, - {"prob": 0.0002238750457763672, "choice": "22"}, - {"prob": 0.00018131732940673828, "choice": "26"}, - {"prob": 0.0001690387725830078, "choice": "24"}, - {"prob": 0.00012552738189697266, "choice": "19"}, - {"prob": 7.796287536621094e-05, "choice": "27"}, - {"prob": 7.265806198120117e-05, "choice": "28"}, - {"prob": 4.106760025024414e-05, "choice": "17"}, - {"prob": 2.5033950805664062e-06, "choice": "29"}, + {"prob": 0.62353515625, "choice": "10"}, + {"prob": 0.349609375, "choice": "12"}, + {"prob": 0.01123809814453125, "choice": "11"}, + {"prob": 0.00760650634765625, "choice": "16"}, + {"prob": 0.0025482177734375, "choice": "13"}, + {"prob": 0.0025081634521484375, "choice": "15"}, + {"prob": 0.0018062591552734375, "choice": "14"}, + {"prob": 0.00104522705078125, "choice": "18"}, + {"prob": 0.00011551380157470703, "choice": "17"}, + {"prob": 5.042552947998047e-05, "choice": "19"}, ], - "age_wmean": 17.816404402256012, + "age_wmean": 15.544570922851562, "is_student_probs": [ - {"prob": 0.974609375, "choice": "true"}, - {"prob": 0.025177001953125, "choice": "false"}, + {"prob": 0.962890625, "choice": "true"}, + {"prob": 0.037322998046875, "choice": "false"}, ], "is_student": False, "name": "John", @@ -80,7 +66,7 @@ generated_data = { "unit_time": 0.5, "courses": ["C++"], "trim": None, - "color": "white", + "color": "green", } ``` diff --git a/dev.ipynb b/dev.ipynb new file mode 100644 index 0000000..c5b9b3a --- /dev/null +++ b/dev.ipynb @@ -0,0 +1,214 @@ +{ + "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": "stderr", + "output_type": "stream", + "text": [ + "/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" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "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", + "/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" + ] + } + ], + "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": 18, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'age_probs': [{'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", + " 'age_wmean': 17.816404402256012,\n", + " 'is_student_probs': [{'prob': 0.974609375, 'choice': 'true'},\n", + " {'prob': 0.025177001953125, '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': 'white'}" + ] + }, + "execution_count": 18, + "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\", \"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": null, + "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", + "}" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "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.9.16" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/example.ipynb b/example.ipynb index 60d776f..bca0cf8 100644 --- a/example.ipynb +++ b/example.ipynb @@ -1,452 +1,359 @@ { - "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": "stderr", - "output_type": "stream", - "text": [ - "/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" - ] + "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": "stderr", + "output_type": "stream", + "text": [ + "/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" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "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", + "/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" + ] + } + ], + "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": 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", + " \"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": 9, + "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": ".venv", + "language": "python", + "name": "python3" + }, + "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.9.16" + }, + "orig_nbformat": 4 }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "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", - "/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" - ] - } - ], - "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": 4, - "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\"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.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", - "}\n" - ] - } - ], - "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", - " \"choice_probs\": {\"type\": \"choice_probs\", \"enum\": [\"ski\", \"snowboard\", \"walk\", \"pretend\"]},\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\": \"choice_probs\", \"enum\": [\"Mazda\", \"Kea\"]},\n", - " \"new\": {\"type\": \"choice_probs\", \"enum\": [\"true\", \"false\"]},\n", - " \"rating\": {\"type\": \"choice_probs\", \"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\": {\"type\": \"choice_probs\", \"enum\": [\"red\", \"green\", \"blue\", \"black\", \"white\"]},\n", - " \"as_new\": {\"type\": \"choice_probs\", \"enum\": [\"true\", \"false\"]},\n", - " \"rating\": {\"type\": \"choice_probs\", \"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 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)" - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/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", - "/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", - "/media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/.venv/lib/python3.9/site-packages/transformers/generation/configuration_utils.py:492: UserWarning: `do_sample` is set to `False`. However, `temperature` is set to `0` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `temperature`.\n", - " warnings.warn(\n" - ] - }, - { - "data": { - "text/plain": [ - "{'name': 'John',\n", - " 'age': 20,\n", - " 'age_probs': [{'prob': 0.856144905090332, 'choice': '12'},\n", - " {'prob': 0.045701637864112854, 'choice': '10'},\n", - " {'prob': 0.030096691101789474, 'choice': '20'},\n", - " {'prob': 0.01899518258869648, 'choice': '11'},\n", - " {'prob': 0.013291668146848679, 'choice': '16'},\n", - " {'prob': 0.013288195244967937, 'choice': '14'},\n", - " {'prob': 0.011642636731266975, 'choice': '18'},\n", - " {'prob': 0.005356263369321823, 'choice': '15'},\n", - " {'prob': 0.0035301733296364546, 'choice': '13'},\n", - " {'prob': 0.0010820770403370261, 'choice': '21'},\n", - " {'prob': 0.0003798121470026672, 'choice': '19'},\n", - " {'prob': 0.0002950581256300211, 'choice': '17'},\n", - " {'prob': 7.64212163630873e-05, 'choice': '22'},\n", - " {'prob': 4.703202284872532e-05, 'choice': '23'},\n", - " {'prob': 2.3594444428454153e-05, 'choice': '25'},\n", - " {'prob': 1.987080577237066e-05, 'choice': '24'},\n", - " {'prob': 1.821534169721417e-05, 'choice': '26'},\n", - " {'prob': 9.411132850800641e-06, 'choice': '28'},\n", - " {'prob': 7.120665941329207e-07, 'choice': '27'},\n", - " {'prob': 4.4407053678696684e-07, 'choice': '29'}],\n", - " 'age_wmean': 18.284568134928122,\n", - " 'unit_time': 0.5,\n", - " 'is_student': True,\n", - " 'is_student_probs': [{'prob': 0.885669469833374, 'choice': 'true'},\n", - " {'prob': 0.1143304631114006, 'choice': 'false'}],\n", - " 'courses': ['C1'],\n", - " 'trim': None,\n", - " 'color': 'white'}" - ] - }, - "execution_count": 1, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "json_schema = {\n", - " \"type\": \"object\",\n", - " \"properties\": {\n", - " \"name\": {\"type\": \"string\", \"maxLength\": 4},\n", - " \"age_probs\": {\"type\": \"choice_probs\", \"enum\": [str(s) for s in range(10, 30)]},\n", - " \"age_wmean\": {\"type\": \"range_mean\", \"minimum\": 10, \"maximum\": 30},\n", - " \"is_student_probs\": {\"type\": \"choice_probs\", \"enum\": [\"true\", \"false\"]},\n", - " \"is_student\": {\"type\": \"boolean\"},\n", - " \"age\": {\"type\": \"integer\"},\n", - " \"unit_time\": {\"type\": \"number\"},\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": { - "kernelspec": { - "display_name": ".venv", - "language": "python", - "name": "python3" - }, - "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.9.16" - }, - "orig_nbformat": 4 - }, - "nbformat": 4, - "nbformat_minor": 2 + "nbformat": 4, + "nbformat_minor": 2 } diff --git a/prob_jsonformer/main.py b/prob_jsonformer/main.py index 538febb..47a8142 100644 --- a/prob_jsonformer/main.py +++ b/prob_jsonformer/main.py @@ -181,12 +181,12 @@ class Jsonformer: return response.split('"')[0].strip() - def generate_choice_probs(self, choices) -> str: + def generate_p_enum(self, choices) -> str: """ This is not in the json schema, but can be usefull for effeciently getting the prob distibution over choices """ prompt = self.get_prompt() + '"' - self.debug("[generate_string_prob]", prompt, is_prompt=True) + self.debug("[generate_p_enum]", prompt, is_prompt=True) input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to( self.model.device )[0] @@ -194,14 +194,14 @@ class Jsonformer: choices_tokens = [torch.tensor(c) for c in choices_tokens] r = list(choice_tree(self.model, self.tokenizer, input_ids, choices_tokens)) - return r # json.dumps(r) + return r - def generate_range_mean(self, range_min: float, range_max: float) -> float: + def generate_p_integer(self, range_min: float, range_max: float) -> float: """ This is not in the json schema, but can be usefull for effeciently generating the weighted mean from a range of integers """ choices = [str(n) for n in range(int(range_min), int(range_max) + 1)] - result = self.generate_choice_probs(choices) + result = self.generate_p_enum(choices) # now do a weighted average total = 0.0 @@ -332,18 +332,18 @@ class Jsonformer: return self.generate_string( schema["maxLength"] if "maxLength" in schema else None ) - elif schema_type == "choice_probs": + elif schema_type == "p_enum": if key: obj[key] = self.generation_marker else: obj.append(self.generation_marker) - return self.generate_choice_probs(schema["enum"]) - elif schema_type == "range_mean": + return self.generate_p_enum(schema["values"]) + elif schema_type == "p_integer": if key: obj[key] = self.generation_marker else: obj.append(self.generation_marker) - return self.generate_range_mean(schema["minimum"], schema["maximum"]) + return self.generate_p_integer(schema["minimum"], schema["maximum"]) elif schema_type == "enum": if key: obj[key] = self.generation_marker