diff --git a/README.md b/README.md index f7b342e..c07806d 100644 --- a/README.md +++ b/README.md @@ -1,8 +1,8 @@ # 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](https://github.com/1rgs/jsonformer) is includesbelow. +This fork has been modified to include the token probabilities. This is not complaint with json schema, but it can be useful for efficient extracting of a range of possible values. -I've also merged some of hte recent pr's for enum, integer, null, union. You can see them all below in this example: +I've also merged some of the recent PR's for enum, integer, null, union. They are not yet included in the upstream Jsonformer. You can see them all below in this example: ## Example @@ -18,12 +18,17 @@ tokenizer = AutoTokenizer.from_pretrained(model_name) json_schema = { "type": "object", "properties": { - "name": {"type": "string"}, - "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"}, + # 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)]}, + # we can return the probabilistic weighted mean of a range + "age_wmean": {"type": "range_mean", "minimum": 10, "maximum": 30}, + # the prob of true and false "is_student_probs": {"type": "choice_probs", "enum": ["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"} @@ -36,31 +41,45 @@ json_schema = { } } + 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 = {'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'], +generated_data = {'name': 'John', + 'age_probs': [{'prob': 0.62353515625, 'choice': '10'}, + {'prob': 0.276611328125, 'choice': '12'}, + {'prob': 0.05364990234375, 'choice': '20'}, + {'prob': 0.0257415771484375, 'choice': '11'}, + {'prob': 0.0047607421875, 'choice': '15'}, + {'prob': 0.004688262939453125, 'choice': '16'}, + {'prob': 0.002910614013671875, 'choice': '18'}, + {'prob': 0.0024127960205078125, 'choice': '13'}, + {'prob': 0.0015821456909179688, 'choice': '14'}, + {'prob': 0.0013532638549804688, 'choice': '23'}, + {'prob': 0.0012521743774414062, 'choice': '21'}, + {'prob': 0.00042247772216796875, 'choice': '17'}, + {'prob': 0.0003342628479003906, 'choice': '22'}, + {'prob': 0.0002484321594238281, 'choice': '19'}, + {'prob': 0.0001995563507080078, 'choice': '25'}, + {'prob': 4.851818084716797e-05, 'choice': '24'}, + {'prob': 3.30805778503418e-05, 'choice': '26'}, + {'prob': 2.6404857635498047e-05, 'choice': '28'}, + {'prob': 1.728534698486328e-05, 'choice': '27'}, + {'prob': 2.9802322387695312e-06, 'choice': '29'}], + 'age_wmean': 17.32853078842163, + 'is_student_probs': [{'prob': 0.8173828125, 'choice': 'true'}, + {'prob': 0.182373046875, 'choice': 'false'}], + 'is_student': False, + 'age': 17, + 'unit_time': 0.5, + 'courses': ['CS101'], 'trim': None, 'color': 'white'} ``` + The original [README](https://github.com/1rgs/jsonformer) is includes below. + # ORIGINAL: Jsonformer: A Bulletproof Way to Generate Structured JSON from Language Models. ### Problem: Getting models to output structured JSON is hard diff --git a/example.ipynb b/example.ipynb index 7807ed8..60d776f 100644 --- a/example.ipynb +++ b/example.ipynb @@ -312,65 +312,87 @@ }, { "cell_type": "code", - "execution_count": 22, + "execution_count": null, "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" - } - ], + "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)\n", - "\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\"},\n", - " \"age\": {\"type\": \"integer\"},\n", + " \"name\": {\"type\": \"string\", \"maxLength\": 4},\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", + " \"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", diff --git a/prob_jsonformer/logits_processors.py b/prob_jsonformer/logits_processors.py index 4e66d26..0e763b8 100644 --- a/prob_jsonformer/logits_processors.py +++ b/prob_jsonformer/logits_processors.py @@ -10,7 +10,6 @@ class StringStoppingCriteria(StoppingCriteria): self.tokenizer = tokenizer self.prompt_length = prompt_length self.max_length = max_length - print(max_length, ", max_length") def __call__( self, @@ -28,7 +27,6 @@ class StringStoppingCriteria(StoppingCriteria): if self.max_length is not None: str_l = len(self.tokenizer.decode(input_ids[0], skip_special_tokens=True)) if str_l > self.max_length: - print("maxlen", str_l) return True return result