diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..ebfde9a --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,3 @@ +{ + "editor.formatOnSave": true, +} \ No newline at end of file diff --git a/example.ipynb b/example.ipynb index bf84fbb..825b797 100644 --- a/example.ipynb +++ b/example.ipynb @@ -37,13 +37,37 @@ "execution_count": 2, "metadata": {}, "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n", + "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Generating...\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n", + "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n" + ] + }, { "name": "stdout", "output_type": "stream", "text": [ - "Generating...\n", "{\n", - " temperature: \u001b[32m15.0\u001b[0m\n", + " name: \u001b[32m\"Mickey Mouse\"\u001b[0m,\n", + " age: \u001b[32m30.0\u001b[0m,\n", + " description: \u001b[32m\"A mouse like object\"\u001b[0m\n", "}\n" ] } @@ -52,18 +76,21 @@ "from jsonformer.format import highlight_values\n", "from jsonformer.main import Jsonformer\n", "\n", - "weather_schema = {\n", + "mickey_mouse_schema = {\n", " \"type\": \"object\",\n", " \"properties\": {\n", - " \"temperature\": {\"type\": \"number\"},\n", + " \"name\": {\"type\": \"string\"},\n", + " \"age\": {\"type\": \"number\"},\n", + " \"description\": {\"type\": \"string\"},\n", " },\n", "}\n", "\n", "builder = Jsonformer(\n", " model=model,\n", " tokenizer=tokenizer,\n", - " json_schema=weather_schema,\n", - " prompt=\"today's temperature in celsius is\",\n", + " json_schema=mickey_mouse_schema,\n", + " prompt=\"write a description about mickey mouse\",\n", + " max_string_token_length=100,\n", ")\n", "\n", "print(\"Generating...\")\n", @@ -77,14 +104,38 @@ "execution_count": 3, "metadata": {}, "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n", + "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n", + "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n", + "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Generating...\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n", + "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n" + ] + }, { "name": "stdout", "output_type": "stream", "text": [ - "Generating...\n", "{\n", " make: \u001b[32m\"audi\"\u001b[0m,\n", - " model: \u001b[32m\"model a4\"\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", @@ -122,10 +173,95 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 4, "metadata": {}, - "outputs": [], - "source": [] + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n", + "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Generating...\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n", + "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n", + "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n", + "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n", + "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n", + "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n", + "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n", + "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n", + "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n", + "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n", + "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n", + "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "{\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" + ] + } + ], + "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", + "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)\n" + ] } ], "metadata": { diff --git a/jsonformer/example.py b/jsonformer/example.py deleted file mode 100644 index 9f5c451..0000000 --- a/jsonformer/example.py +++ /dev/null @@ -1,48 +0,0 @@ -from jsonformer.format import highlight_values -from jsonformer.main import Jsonformer -from transformers import ( - AutoModelForCausalLM, - AutoTokenizer, -) - - -weather_schema = { - "type": "object", - "properties": { - "temperature": {"type": "number"}, - "humidity": { - "type": "number", - }, - "wind_speed": { - "type": "object", - "properties": { - "value": {"type": "number"}, - "unit": {"type": "string"}, - }, - }, - }, -} - -print("Loading model and tokenizer...") -model_name = "databricks/dolly-v2-12b" - -model = AutoModelForCausalLM.from_pretrained( - model_name, use_cache=True, device_map="auto" -) -tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, use_cache=True) -print("Loaded model and tokenizer") - - -builder = Jsonformer( - model=model, - tokenizer=tokenizer, - json_schema=weather_schema, - prompt="generate the weather", -) - -print("Generating...") -output = builder() - -highlight_values( - output, -) diff --git a/jsonformer/logits_processors.py b/jsonformer/logits_processors.py index 3d10051..682e906 100644 --- a/jsonformer/logits_processors.py +++ b/jsonformer/logits_processors.py @@ -3,6 +3,27 @@ from transformers import PreTrainedTokenizer, LogitsWarper, StoppingCriteria import torch +class StringStoppingCriteria(StoppingCriteria): + def __init__(self, tokenizer: PreTrainedTokenizer, prompt_length: int): + self.tokenizer = tokenizer + self.prompt_length = prompt_length + + def __call__( + self, + input_ids: torch.LongTensor, + _, + ) -> bool: + if len(input_ids[0]) <= self.prompt_length: + return False + + last_token_id = input_ids[0][-1] + last_token = self.tokenizer.decode(last_token_id, skip_special_tokens=True) + + result = '"' in last_token + + return result + + class NumberStoppingCriteria(StoppingCriteria): def __init__( self, diff --git a/jsonformer/main.py b/jsonformer/main.py index 359132a..f709a3b 100644 --- a/jsonformer/main.py +++ b/jsonformer/main.py @@ -1,6 +1,11 @@ from typing import List, Union, Dict, Any -from jsonformer.logits_processors import NumberStoppingCriteria, OutputNumbersTokens +from jsonformer.logits_processors import ( + NumberStoppingCriteria, + OutputNumbersTokens, + StringStoppingCriteria, +) +from termcolor import cprint from transformers import PreTrainedModel, PreTrainedTokenizer import json @@ -38,13 +43,18 @@ class Jsonformer: self.temperature = temperature self.max_string_token_length = max_string_token_length - def debug(self, *args, **kwargs): + def debug(self, caller: str, value: str, is_prompt: bool = False): if self.debug_on: - print(*args, **kwargs) + if is_prompt: + cprint(caller, "green", end=" ") + cprint(value, "yellow") + else: + cprint(caller, "green", end=" ") + cprint(value, "blue") def generate_number(self, temperature: Union[float, None] = None, iterations=0): prompt = self.get_prompt() - self.debug("[generate_number] prompt", prompt) + self.debug("[generate_number]", prompt, is_prompt=True) input_tokens = self.tokenizer.encode(prompt, return_tensors="pt").to( self.model.device ) @@ -60,10 +70,10 @@ class Jsonformer: pad_token_id=self.tokenizer.eos_token_id, ) response = self.tokenizer.decode(response[0], skip_special_tokens=True) - self.debug("[generate_number] response", response) + response = response[len(prompt) :] response = response.strip().rstrip(".") - + self.debug("[generate_number]", response) try: return float(response) except ValueError: @@ -74,50 +84,67 @@ class Jsonformer: def generate_boolean(self) -> bool: prompt = self.get_prompt() - self.debug("[generate_boolean] prompt", prompt) + self.debug("[generate_boolean]", prompt, is_prompt=True) input_tensor = self.tokenizer.encode(prompt, return_tensors="pt") 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 + # this is probably not true for all tokenizers + # this can be fixed by looking at only the first token of both "true" and "false" true_token_id = self.tokenizer.convert_tokens_to_ids("true") false_token_id = self.tokenizer.convert_tokens_to_ids("false") - true_logits = logits[true_token_id] - false_logits = logits[false_token_id] + result = logits[true_token_id] > logits[false_token_id] - if true_logits > false_logits: - return True - elif false_logits > true_logits: - return False - else: - print("Failed to generate a valid boolean value") - return None + self.debug("[generate_boolean]", result) + + return result.item() def generate_string(self) -> str: - prompt = self.get_prompt() - self.debug("[generate_string] prompt", prompt) + prompt = self.get_prompt() + '"' + self.debug("[generate_string]", prompt, is_prompt=True) + input_tokens = self.tokenizer.encode(prompt, return_tensors="pt").to( + self.model.device + ) + response = self.model.generate( - self.tokenizer.encode(prompt, return_tensors="pt").to(self.model.device), + input_tokens, max_new_tokens=self.max_string_token_length, num_return_sequences=1, temperature=self.temperature, - pad_token_id=self.tokenizer.eos_token_id, + stopping_criteria=[ + StringStoppingCriteria(self.tokenizer, len(input_tokens[0])) + ], ) - response = self.tokenizer.decode(response[0], skip_special_tokens=True) - response = response[len(prompt) :].strip() + # Some models output the prompt as part of the response + # This removes the prompt from the response if it is present + if ( + len(response[0]) >= len(input_tokens[0]) + and (response[0][: len(input_tokens[0])] == input_tokens).all() + ): + response = response[0][len(input_tokens[0]) :] + if response.shape[0] == 1: + response = response[0] - self.debug("[generate_string] response", response) - split = response.split('"') - assert len(split) >= 2 - return split[1] + response = self.tokenizer.decode(response, skip_special_tokens=True) + + self.debug("[generate_string]", "|" + response + "|") + + if response.count('"') < 1: + raise ValueError( + "Failed to generate a valid string, try increasing max_string_token_length" + ) + + return response.split('"')[0].strip() def generate_object( self, properties: Dict[str, Any], obj: Dict[str, Any] ) -> Dict[str, Any]: - # self.debug("[generate_object] properties", properties) for key, schema in properties.items(): + self.debug("[generate_object] generating value for", key) obj[key] = self.generate_value(schema, obj, key) return obj @@ -162,6 +189,7 @@ class Jsonformer: def generate_array(self, item_schema: Dict[str, Any], obj: Dict[str, Any]) -> list: for _ in range(self.max_array_length): + # forces array to have at least one element element = self.generate_value(item_schema, obj) obj[-1] = element diff --git a/poetry.lock b/poetry.lock index 80fc3e0..4ba502a 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry and should not be changed by hand. +# This file is automatically @generated by Poetry 1.4.2 and should not be changed by hand. [[package]] name = "accelerate" @@ -1597,6 +1597,10 @@ category = "dev" optional = false python-versions = ">=3.8.0" files = [ + {file = "torch-2.0.0-1-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:c9090bda7d2eeeecd74f51b721420dbeb44f838d4536cc1b284e879417e3064a"}, + {file = "torch-2.0.0-1-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:bd42db2a48a20574d2c33489e120e9f32789c4dc13c514b0c44272972d14a2d7"}, + {file = "torch-2.0.0-1-cp38-cp38-manylinux2014_aarch64.whl", hash = "sha256:8969aa8375bcbc0c2993e7ede0a7f889df9515f18b9b548433f412affed478d9"}, + {file = "torch-2.0.0-1-cp39-cp39-manylinux2014_aarch64.whl", hash = "sha256:ab2da16567cb55b67ae39e32d520d68ec736191d88ac79526ca5874754c32203"}, {file = "torch-2.0.0-cp310-cp310-manylinux1_x86_64.whl", hash = "sha256:7a9319a67294ef02459a19738bbfa8727bb5307b822dadd708bc2ccf6c901aca"}, {file = "torch-2.0.0-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:9f01fe1f6263f31bd04e1757946fd63ad531ae37f28bb2dbf66f5c826ee089f4"}, {file = "torch-2.0.0-cp310-cp310-win_amd64.whl", hash = "sha256:527f4ae68df7b8301ee6b1158ca56350282ea633686537b30dbb5d7b4a52622a"}, @@ -1775,6 +1779,15 @@ category = "dev" optional = false python-versions = "*" files = [ + {file = "triton-2.0.0-1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:38806ee9663f4b0f7cd64790e96c579374089e58f49aac4a6608121aa55e2505"}, + {file = "triton-2.0.0-1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:226941c7b8595219ddef59a1fdb821e8c744289a132415ddd584facedeb475b1"}, + {file = "triton-2.0.0-1-cp36-cp36m-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4c9fc8c89874bc48eb7e7b2107a9b8d2c0bf139778637be5bfccb09191685cfd"}, + {file = "triton-2.0.0-1-cp37-cp37m-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d2684b6a60b9f174f447f36f933e9a45f31db96cb723723ecd2dcfd1c57b778b"}, + {file = "triton-2.0.0-1-cp38-cp38-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9d4978298b74fcf59a75fe71e535c092b023088933b2f1df933ec32615e4beef"}, + {file = "triton-2.0.0-1-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:74f118c12b437fb2ca25e1a04759173b517582fcf4c7be11913316c764213656"}, + {file = "triton-2.0.0-1-pp37-pypy37_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9618815a8da1d9157514f08f855d9e9ff92e329cd81c0305003eb9ec25cc5add"}, + {file = "triton-2.0.0-1-pp38-pypy38_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1aca3303629cd3136375b82cb9921727f804e47ebee27b2677fef23005c3851a"}, + {file = "triton-2.0.0-1-pp39-pypy39_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e3e13aa8b527c9b642e3a9defcc0fbd8ffbe1c80d8ac8c15a01692478dc64d8a"}, {file = "triton-2.0.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8f05a7e64e4ca0565535e3d5d3405d7e49f9d308505bb7773d21fb26a4c008c2"}, {file = "triton-2.0.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bb4b99ca3c6844066e516658541d876c28a5f6e3a852286bbc97ad57134827fd"}, {file = "triton-2.0.0-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:47b4d70dc92fb40af553b4460492c31dc7d3a114a979ffb7a5cdedb7eb546c08"},