From 637420044c5ccf0c78e099a861acad176a65c370 Mon Sep 17 00:00:00 2001 From: wassname Date: Fri, 10 May 2024 16:49:12 +0800 Subject: [PATCH] prototype a enum with prob --- README.md | 4 + example.ipynb | 390 ++++++++++++++++++++++++++++++----------- jsonformer/__init__.py | 4 +- jsonformer/main.py | 23 ++- pyproject.toml | 8 +- 5 files changed, 311 insertions(+), 118 deletions(-) diff --git a/README.md b/README.md index a734002..36ad8f8 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,7 @@ +# 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 is below. + # 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 49520e8..2d498a1 100644 --- a/example.ipynb +++ b/example.ipynb @@ -9,7 +9,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "/home/ubuntu/jsonformer/.venv/lib/python3.8/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", + "/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" ] }, @@ -17,7 +17,27 @@ "name": "stdout", "output_type": "stream", "text": [ - "Loading model and tokenizer...\n", + "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", + "`flash-attention` package not found, consider installing for better performance: No module named 'flash_attn'.\n", + "Current `flash-attenton` does not support `window_size`. Either upgrade or use `attn_implementation='eager'`.\n", + "Loading checkpoint shards: 100%|██████████| 4/4 [00:04<00:00, 1.10s/it]\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" ] } @@ -26,60 +46,228 @@ "from transformers import AutoModelForCausalLM, AutoTokenizer\n", "\n", "print(\"Loading model and tokenizer...\")\n", - "model_name = \"databricks/dolly-v2-3b\"\n", - "model = AutoModelForCausalLM.from_pretrained(model_name, use_cache=True, device_map=\"auto\")\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", + " # device=\"cuda:0\",\n", + " # device_map=\"auto\",\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\")" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Scratch" + ] + }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 65, + "metadata": {}, + "outputs": [], + "source": [ + "from jaxtyping import Float, Int\n", + "import torch\n", + "from torch.nn import functional as F\n", + "from torch import Tensor\n", + "from typing import List, Callable, Tuple, Dict, Optional\n", + "import pandas as pd" + ] + }, + { + "cell_type": "code", + "execution_count": 43, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[tensor([ 1, 29871, 29896]),\n", + " tensor([ 1, 278, 18109, 11285]),\n", + " tensor([ 1, 278, 289, 4992, 12544])]" + ] + }, + "execution_count": 43, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# initital state\n", + "prompt = \"The necromancer in his tower, what's his top problem? \"\n", + "choices = [\"1\", \"the skeleton\", \"the boney boys\"]\n", + "choices_tokens = tokenizer(choices).input_ids\n", + "choices_tokens = [torch.tensor(c) for c in choices_tokens]\n", + "# current_tokens = torch.tensor([])\n", + "\n", + "# next\n", + "input_ids = tokenizer([prompt], return_tensors=\"pt\").to(model.device).input_ids[0]\n", + "choices_tokens\n", + "\n", + "# for each next choice, continue down the tree, recording the log probs" + ] + }, + { + "cell_type": "code", + "execution_count": 70, "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\"Sugarloaf\"\u001b[0m,\n", - " inventory: [\n", - " {\n", - " productId: \u001b[32m\"1234567890\"\u001b[0m,\n", - " name: \u001b[32m\"Ski Shop\"\u001b[0m,\n", - " description: \u001b[32m\"Ski Shop sells premium skis and snowboards\"\u001b[0m,\n", - " category: \u001b[32m\"Sports\"\u001b[0m,\n", - " price: \u001b[32m20.09999\u001b[0m,\n", - " inStock: \u001b[32mTrue\u001b[0m,\n", - " rating: \u001b[32m5.09999\u001b[0m,\n", - " images: [\n", - " \u001b[32m\"https://s3.amazonaws.com/images.skisnow.com/skis\"\u001b[0m\n", - " ]\n", - " },\n", - " {\n", - " productId: \u001b[32m\"12345678910\"\u001b[0m,\n", - " name: \u001b[32m\"Snowboard\"\u001b[0m,\n", - " description: \u001b[32m\"Snowboard for sale\"\u001b[0m,\n", - " category: \u001b[32m\"Sports\"\u001b[0m,\n", - " price: \u001b[32m20.09999\u001b[0m,\n", - " inStock: \u001b[32mTrue\u001b[0m,\n", - " rating: \u001b[32m5.09999\u001b[0m,\n", - " images: [\n", - " \u001b[32m\"https://s3.amazonaws.com/images.skisnow.com/snow\"\u001b[0m\n", - " ]\n", - " }\n", - " ]\n", - " }\n", - "}\n" + "tensor([1]) current_tokens\n", + "tensor([ 1, 278]) current_tokens\n", + "tensor([ 1, 278, 289]) current_tokens\n", + "tensor([ 1, 278, 289, 4992]) current_tokens\n", + "tensor([ 1, 278, 289, 4992, 12544]) current_tokens\n", + "tensor([ 1, 278, 18109]) current_tokens\n", + "tensor([ 1, 278, 18109, 11285]) current_tokens\n", + "tensor([ 1, 29871]) current_tokens\n", + "tensor([ 1, 29871, 29896]) current_tokens\n" ] } ], "source": [ - "from jsonformer.format import highlight_values\n", - "from jsonformer.main import Jsonformer\n", + "def get_valid_next_choices(choices_tokens, current_tokens):\n", + " next_choices = []\n", + " for choice_tokens in choices_tokens:\n", + " # if we have some more slots left\n", + " if len(current_tokens)\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
probchoice
20.995732<s> 1
00.004187<s> the boney boys
10.000081<s> the skeleton
\n", + "" + ], + "text/plain": [ + " prob choice\n", + "2 0.995732 1\n", + "0 0.004187 the boney boys\n", + "1 0.000081 the skeleton" + ] + }, + "execution_count": 74, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "pd.DataFrame(r).sort_values(\"prob\", ascending=False).drop(columns=[\"tokens\"])" + ] + }, + { + "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", @@ -127,25 +315,9 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Generating...\n", - "{\n", - " make: \u001b[32m\"audi\"\u001b[0m,\n", - " model: \u001b[32m\"a4\"\u001b[0m,\n", - " year: \u001b[32m2016.0\u001b[0m,\n", - " colors_available: [\n", - " \u001b[32m\"blue\"\u001b[0m\n", - " ]\n", - "}\n" - ] - } - ], + "outputs": [], "source": [ "car = {\n", " \"type\": \"object\",\n", @@ -170,56 +342,66 @@ "print(\"Generating...\")\n", "output = builder()\n", "\n", - "highlight_values(output)\n" + "highlight_values(output)" ] }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "metadata": {}, - "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[32m2016.0\u001b[0m,\n", - " colors: [\n", - " \u001b[32m\"Gold\"\u001b[0m\n", - " ],\n", - " features: {\n", - " audio: {\n", - " brand: \u001b[32m\"Mercedes\"\u001b[0m,\n", - " speakers: \u001b[32m2.0\u001b[0m,\n", - " hasBluetooth: \u001b[32mTrue\u001b[0m\n", - " },\n", - " safety: {\n", - " airbags: \u001b[32m2.0\u001b[0m,\n", - " parkingSensors: \u001b[32mTrue\u001b[0m,\n", - " laneAssist: \u001b[32mTrue\u001b[0m\n", - " },\n", - " performance: {\n", - " engine: \u001b[32m\"Mercedes-Benz OM615\"\u001b[0m,\n", - " horsepower: \u001b[32m350.0\u001b[0m,\n", - " topSpeed: \u001b[32m220.0\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.0\u001b[0m\n", - " }\n", - "}\n" - ] - } - ], + "outputs": [], "source": [ - "complex_car = {\"type\": \"object\", \"properties\": {\"car\": {\"type\": \"object\", \"properties\": {\"make\": {\"type\": \"string\"}, \"model\": {\"type\": \"string\"}, \"year\": {\"type\": \"number\"}, \"colors\": {\"type\": \"array\", \"items\": {\"type\": \"string\"}}, \"features\": {\"type\": \"object\", \"properties\": {\"audio\": {\"type\": \"object\", \"properties\": {\"brand\": {\"type\": \"string\"}, \"speakers\": {\"type\": \"number\"}, \"hasBluetooth\": {\"type\": \"boolean\"}}}, \"safety\": {\"type\": \"object\", \"properties\": {\"airbags\": {\"type\": \"number\"}, \"parkingSensors\": {\"type\": \"boolean\"}, \"laneAssist\": {\"type\": \"boolean\"}}}, \"performance\": {\"type\": \"object\", \"properties\": {\"engine\": {\"type\": \"string\"}, \"horsepower\": {\"type\": \"number\"}, \"topSpeed\": {\"type\": \"number\"}}}}}}}, \"owner\": {\"type\": \"object\", \"properties\": {\"firstName\": {\"type\": \"string\"}, \"lastName\": {\"type\": \"string\"}, \"age\": {\"type\": \"number\"}}}}}\n", + "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\": \"array\", \"items\": {\"type\": \"string\"}},\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", @@ -230,7 +412,7 @@ "print(\"Generating...\")\n", "output = builder()\n", "\n", - "highlight_values(output)\n" + "highlight_values(output)" ] } ], @@ -250,7 +432,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.10" + "version": "3.9.16" }, "orig_nbformat": 4 }, diff --git a/jsonformer/__init__.py b/jsonformer/__init__.py index 495551f..7f02ee5 100644 --- a/jsonformer/__init__.py +++ b/jsonformer/__init__.py @@ -1,2 +1,2 @@ -from jsonformer.main import Jsonformer -from jsonformer.format import highlight_values \ No newline at end of file +from prob_jsonformer.main import prob_jsonformer +from prob_jsonformer.format import highlight_values diff --git a/jsonformer/main.py b/jsonformer/main.py index 9c13471..4c6686a 100644 --- a/jsonformer/main.py +++ b/jsonformer/main.py @@ -1,6 +1,6 @@ from typing import List, Union, Dict, Any -from jsonformer.logits_processors import ( +from prob_jsonformer.logits_processors import ( NumberStoppingCriteria, OutputNumbersTokens, StringStoppingCriteria, @@ -12,7 +12,7 @@ import json GENERATION_MARKER = "|GENERATION|" -class Jsonformer: +class prob_jsonformer: value: Dict[str, Any] = {} def __init__( @@ -80,9 +80,11 @@ class Jsonformer: if iterations > 3: raise ValueError("Failed to generate a valid number") - return self.generate_number(temperature=self.temperature * 1.3, iterations=iterations+1) + return self.generate_number( + temperature=self.temperature * 1.3, iterations=iterations + 1 + ) - def generate_boolean(self) -> bool: + def generate_boolean(self, prob=False) -> bool: prompt = self.get_prompt() self.debug("[generate_boolean]", prompt, is_prompt=True) @@ -90,13 +92,17 @@ class Jsonformer: output = self.model.forward(input_tensor.to(self.model.device)) logits = output.logits[0, -1] - # todo: this assumes that "true" and "false" are both tokenized to a single token + # TODO: this assumes that "true" and "false" are both tokenized to a single token # this is probably not true for all tokenizers # this can be fixed by looking at only the first token of both "true" and "false" + # FIXME: consdier " True", " true", "\ntrue" etc true_token_id = self.tokenizer.convert_tokens_to_ids("true") false_token_id = self.tokenizer.convert_tokens_to_ids("false") - result = logits[true_token_id] > logits[false_token_id] + if prob: + result = dict(true=logits[true_token_id], false=logits[false_token_id]) + else: + result = logits[true_token_id] > logits[false_token_id] self.debug("[generate_boolean]", result) @@ -199,7 +205,6 @@ class Jsonformer: output = self.model.forward(input_tensor.to(self.model.device)) logits = output.logits[0, -1] - top_indices = logits.topk(30).indices sorted_token_ids = top_indices[logits[top_indices].argsort(descending=True)] @@ -208,10 +213,10 @@ class Jsonformer: for token_id in sorted_token_ids: decoded_token = self.tokenizer.decode(token_id) - if ',' in decoded_token: + if "," in decoded_token: found_comma = True break - if ']' in decoded_token: + if "]" in decoded_token: found_close_bracket = True break diff --git a/pyproject.toml b/pyproject.toml index 7014fed..0c1e899 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,13 +1,15 @@ [tool.poetry] -name = "jsonformer" +name = "prob_jsonformer" version = "0.12.0" description = "" authors = ["1rgs "] readme = "README.md" [tool.poetry.dependencies] -python = "^3.8" +python = "^3.9,<4.0" termcolor = "^2.3.0" +jaxtyping = "^0.2.28" +pandas = "^2.2.2" [tool.poetry.group.dev.dependencies] ipykernel = "^6.22.0" @@ -23,4 +25,4 @@ build-backend = "poetry.core.masonry.api" [virtualenvs] create = true -in-project = true \ No newline at end of file +in-project = true