From 6a38b4bb64484336f1efb9487a65f4aa51875c0b Mon Sep 17 00:00:00 2001 From: Arman Rafiee Date: Tue, 11 Mar 2025 16:32:50 +0100 Subject: [PATCH 1/7] Update logits_processors.py Imported LogitsProcessor rather than the no longer used LogitsWarper --- jsonformer/logits_processors.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/jsonformer/logits_processors.py b/jsonformer/logits_processors.py index db288d3..f29a4b9 100644 --- a/jsonformer/logits_processors.py +++ b/jsonformer/logits_processors.py @@ -1,5 +1,5 @@ from typing import List -from transformers import PreTrainedTokenizer, LogitsWarper, StoppingCriteria +from transformers import PreTrainedTokenizer, LogitsProcessor, StoppingCriteria import torch class StringStoppingCriteria(StoppingCriteria): @@ -61,7 +61,7 @@ def __call__( return False -class OutputNumbersTokens(LogitsWarper): +class OutputNumbersTokens(LogitsProcessor): def __init__(self, tokenizer: PreTrainedTokenizer, prompt: str): self.tokenizer = tokenizer self.tokenized_prompt = tokenizer(prompt, return_tensors="pt") From c3cd47292572ee8b8602957fe00863b99626090d Mon Sep 17 00:00:00 2001 From: Arman Rafiee Date: Tue, 11 Mar 2025 23:56:17 +0100 Subject: [PATCH 2/7] Update main.py --- jsonformer/main.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/jsonformer/main.py b/jsonformer/main.py index 9c13471..f5dcfb1 100644 --- a/jsonformer/main.py +++ b/jsonformer/main.py @@ -23,10 +23,10 @@ def __init__( prompt: str, *, debug: bool = False, - max_array_length: int = 10, - max_number_tokens: int = 6, + max_array_length: int = 128, + max_number_tokens: int = 2048, temperature: float = 1.0, - max_string_token_length: int = 10, + max_string_token_length: int = 1024, ): self.model = model self.tokenizer = tokenizer From ec91cecfcd80bfeb64d92ce5451ace744a69a8a2 Mon Sep 17 00:00:00 2001 From: Arman Rafiee Date: Tue, 11 Mar 2025 23:57:36 +0100 Subject: [PATCH 3/7] Update pyproject.toml --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 7014fed..2ae34a0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,8 +1,8 @@ [tool.poetry] name = "jsonformer" -version = "0.12.0" +version = "0.12.1" description = "" -authors = ["1rgs "] +authors = ["1rgs ", "armand0e "] readme = "README.md" [tool.poetry.dependencies] @@ -23,4 +23,4 @@ build-backend = "poetry.core.masonry.api" [virtualenvs] create = true -in-project = true \ No newline at end of file +in-project = true From 04fab7d09a807d9ce41db828c424637264fb9efd Mon Sep 17 00:00:00 2001 From: Arman Rafiee Date: Thu, 13 Mar 2025 03:02:06 -0400 Subject: [PATCH 4/7] Restored anyOf support in Jsonformer --- jsonformer/main.py | 196 ++++++++++++++------------------------------- 1 file changed, 58 insertions(+), 138 deletions(-) diff --git a/jsonformer/main.py b/jsonformer/main.py index f5dcfb1..f4ca94a 100644 --- a/jsonformer/main.py +++ b/jsonformer/main.py @@ -1,211 +1,141 @@ from typing import List, Union, Dict, Any - -from jsonformer.logits_processors import ( - NumberStoppingCriteria, - OutputNumbersTokens, - StringStoppingCriteria, -) +from jsonformer.logits_processors import NumberStoppingCriteria, OutputNumbersTokens, StringStoppingCriteria from termcolor import cprint from transformers import PreTrainedModel, PreTrainedTokenizer import json - -GENERATION_MARKER = "|GENERATION|" - +GENERATION_MARKER = '|GENERATION|' class Jsonformer: value: Dict[str, Any] = {} - def __init__( - self, - model: PreTrainedModel, - tokenizer: PreTrainedTokenizer, - json_schema: Dict[str, Any], - prompt: str, - *, - debug: bool = False, - max_array_length: int = 128, - max_number_tokens: int = 2048, - temperature: float = 1.0, - max_string_token_length: int = 1024, - ): + def __init__(self, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, json_schema: Dict[str, Any], prompt: str, *, debug: bool=False, max_array_length: int=128, max_number_tokens: int=2048, temperature: float=1.0, max_string_token_length: int=1024): self.model = model self.tokenizer = tokenizer self.json_schema = json_schema self.prompt = prompt - self.number_logit_processor = OutputNumbersTokens(self.tokenizer, self.prompt) - - self.generation_marker = "|GENERATION|" + self.generation_marker = '|GENERATION|' self.debug_on = debug self.max_array_length = max_array_length - self.max_number_tokens = max_number_tokens self.temperature = temperature self.max_string_token_length = max_string_token_length - def debug(self, caller: str, value: str, is_prompt: bool = False): + def debug(self, caller: str, value: str, is_prompt: bool=False): if self.debug_on: if is_prompt: - cprint(caller, "green", end=" ") - cprint(value, "yellow") + cprint(caller, 'green', end=' ') + cprint(value, 'yellow') else: - cprint(caller, "green", end=" ") - cprint(value, "blue") + cprint(caller, 'green', end=' ') + cprint(value, 'blue') - def generate_number(self, temperature: Union[float, None] = None, iterations=0): + def generate_number(self, temperature: Union[float, None]=None, iterations=0): prompt = self.get_prompt() - self.debug("[generate_number]", prompt, is_prompt=True) - input_tokens = self.tokenizer.encode(prompt, return_tensors="pt").to( - self.model.device - ) - response = self.model.generate( - input_tokens, - max_new_tokens=self.max_number_tokens, - num_return_sequences=1, - logits_processor=[self.number_logit_processor], - stopping_criteria=[ - NumberStoppingCriteria(self.tokenizer, len(input_tokens[0])) - ], - temperature=temperature or self.temperature, - pad_token_id=self.tokenizer.eos_token_id, - ) + self.debug('[generate_number]', prompt, is_prompt=True) + input_tokens = self.tokenizer.encode(prompt, return_tensors='pt').to(self.model.device) + response = self.model.generate(input_tokens, max_new_tokens=self.max_number_tokens, num_return_sequences=1, logits_processor=[self.number_logit_processor], stopping_criteria=[NumberStoppingCriteria(self.tokenizer, len(input_tokens[0]))], temperature=temperature or self.temperature, pad_token_id=self.tokenizer.eos_token_id) response = self.tokenizer.decode(response[0], skip_special_tokens=True) - - response = response[len(prompt) :] - response = response.strip().rstrip(".") - self.debug("[generate_number]", response) + response = response[len(prompt):] + response = response.strip().rstrip('.') + self.debug('[generate_number]', response) try: return float(response) except ValueError: if iterations > 3: - raise ValueError("Failed to generate a valid number") - - return self.generate_number(temperature=self.temperature * 1.3, iterations=iterations+1) + raise ValueError('Failed to generate a valid number') + return self.generate_number(temperature=self.temperature * 1.3, iterations=iterations + 1) def generate_boolean(self) -> bool: prompt = self.get_prompt() - self.debug("[generate_boolean]", prompt, is_prompt=True) - - input_tensor = self.tokenizer.encode(prompt, return_tensors="pt") + 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_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] - - self.debug("[generate_boolean]", result) - + self.debug('[generate_boolean]', result) return result.item() def generate_string(self) -> str: 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( - input_tokens, - max_new_tokens=self.max_string_token_length, - num_return_sequences=1, - temperature=self.temperature, - stopping_criteria=[ - StringStoppingCriteria(self.tokenizer, len(input_tokens[0])) - ], - pad_token_id=self.tokenizer.eos_token_id, - ) - - # 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]) :] + 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(input_tokens, max_new_tokens=self.max_string_token_length, num_return_sequences=1, temperature=self.temperature, stopping_criteria=[StringStoppingCriteria(self.tokenizer, len(input_tokens[0]))], pad_token_id=self.tokenizer.eos_token_id) + 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] - response = self.tokenizer.decode(response, skip_special_tokens=True) - - self.debug("[generate_string]", "|" + response + "|") - + self.debug('[generate_string]', '|' + response + '|') if response.count('"') < 1: return response - return response.split('"')[0].strip() - def generate_object( - self, properties: Dict[str, Any], obj: Dict[str, Any] - ) -> Dict[str, Any]: - for key, schema in properties.items(): - self.debug("[generate_object] generating value for", key) + def generate_object(self, properties: Dict[str, Any], obj: Dict[str, Any]) -> Dict[str, Any]: + for (key, schema) in properties.items(): + self.debug('[generate_object] generating value for', key) obj[key] = self.generate_value(schema, obj, key) return obj - def generate_value( - self, - schema: Dict[str, Any], - obj: Union[Dict[str, Any], List[Any]], - key: Union[str, None] = None, - ) -> Any: - schema_type = schema["type"] - if schema_type == "number": + def generate_value(self, schema: Dict[str, Any], obj: Union[Dict[str, Any], List[Any]], key: Union[str, None]=None) -> Any: + if 'anyOf' in schema: + options = [option for option in schema['anyOf'] if 'type' in option] + if options: + schema = options[0] + else: + raise ValueError('No valid type in anyOf for key: ' + str(key)) + schema_type = schema['type'] + if schema_type == 'number': if key: obj[key] = self.generation_marker else: obj.append(self.generation_marker) return self.generate_number() - elif schema_type == "boolean": + elif schema_type == 'boolean': if key: obj[key] = self.generation_marker else: obj.append(self.generation_marker) return self.generate_boolean() - elif schema_type == "string": + elif schema_type == 'string': if key: obj[key] = self.generation_marker else: obj.append(self.generation_marker) return self.generate_string() - elif schema_type == "array": + elif schema_type == 'array': new_array = [] - obj[key] = new_array - return self.generate_array(schema["items"], new_array) - elif schema_type == "object": + if key: + obj[key] = new_array + else: + obj.append(new_array) + return self.generate_array(schema['items'], new_array) + elif schema_type == 'object': new_obj = {} if key: obj[key] = new_obj else: obj.append(new_obj) - return self.generate_object(schema["properties"], new_obj) + return self.generate_object(schema['properties'], new_obj) else: - raise ValueError(f"Unsupported schema type: {schema_type}") + raise ValueError(f'Unsupported schema type: {schema_type}') 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 - obj.append(self.generation_marker) input_prompt = self.get_prompt() obj.pop() - input_tensor = self.tokenizer.encode(input_prompt, return_tensors="pt") + input_tensor = self.tokenizer.encode(input_prompt, return_tensors='pt') 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)] - found_comma = False found_close_bracket = False - for token_id in sorted_token_ids: decoded_token = self.tokenizer.decode(token_id) if ',' in decoded_token: @@ -214,32 +144,22 @@ def generate_array(self, item_schema: Dict[str, Any], obj: Dict[str, Any]) -> li if ']' in decoded_token: found_close_bracket = True break - if found_close_bracket or not found_comma: break - return obj def get_prompt(self): - template = """{prompt}\nOutput result in the following JSON schema format:\n{schema}\nResult: {progress}""" + template = '{prompt}\nOutput result in the following JSON schema format:\n{schema}\nResult: {progress}' progress = json.dumps(self.value) gen_marker_index = progress.find(f'"{self.generation_marker}"') if gen_marker_index != -1: progress = progress[:gen_marker_index] else: - raise ValueError("Failed to find generation marker") - - prompt = template.format( - prompt=self.prompt, - schema=json.dumps(self.json_schema), - progress=progress, - ) - + raise ValueError('Failed to find generation marker') + prompt = template.format(prompt=self.prompt, schema=json.dumps(self.json_schema), progress=progress) return prompt def __call__(self) -> Dict[str, Any]: self.value = {} - generated_data = self.generate_object( - self.json_schema["properties"], self.value - ) - return generated_data + generated_data = self.generate_object(self.json_schema['properties'], self.value) + return generated_data \ No newline at end of file From 41e050a691180ed3869e9771ea73e9552233b5da Mon Sep 17 00:00:00 2001 From: Arman Rafiee Date: Mon, 17 Mar 2025 12:47:17 -0400 Subject: [PATCH 5/7] Update main.py --- jsonformer/main.py | 196 +++++++++++++++++++++++++++++++-------------- 1 file changed, 138 insertions(+), 58 deletions(-) diff --git a/jsonformer/main.py b/jsonformer/main.py index f4ca94a..7e2fa3b 100644 --- a/jsonformer/main.py +++ b/jsonformer/main.py @@ -1,141 +1,211 @@ from typing import List, Union, Dict, Any -from jsonformer.logits_processors import NumberStoppingCriteria, OutputNumbersTokens, StringStoppingCriteria + +from jsonformer.logits_processors import ( + NumberStoppingCriteria, + OutputNumbersTokens, + StringStoppingCriteria, +) from termcolor import cprint from transformers import PreTrainedModel, PreTrainedTokenizer import json -GENERATION_MARKER = '|GENERATION|' + +GENERATION_MARKER = "|GENERATION|" + class Jsonformer: value: Dict[str, Any] = {} - def __init__(self, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, json_schema: Dict[str, Any], prompt: str, *, debug: bool=False, max_array_length: int=128, max_number_tokens: int=2048, temperature: float=1.0, max_string_token_length: int=1024): + def __init__( + self, + model: PreTrainedModel, + tokenizer: PreTrainedTokenizer, + json_schema: Dict[str, Any], + prompt: str, + *, + debug: bool = False, + max_array_length: int = 20, + max_number_tokens: int = 4096, + temperature: float = 0.1, + max_string_token_length: int = 1024, + ): self.model = model self.tokenizer = tokenizer self.json_schema = json_schema self.prompt = prompt + self.number_logit_processor = OutputNumbersTokens(self.tokenizer, self.prompt) - self.generation_marker = '|GENERATION|' + + self.generation_marker = "|GENERATION|" self.debug_on = debug self.max_array_length = max_array_length + self.max_number_tokens = max_number_tokens self.temperature = temperature self.max_string_token_length = max_string_token_length - def debug(self, caller: str, value: str, is_prompt: bool=False): + def debug(self, caller: str, value: str, is_prompt: bool = False): if self.debug_on: if is_prompt: - cprint(caller, 'green', end=' ') - cprint(value, 'yellow') + cprint(caller, "green", end=" ") + cprint(value, "yellow") else: - cprint(caller, 'green', end=' ') - cprint(value, 'blue') + cprint(caller, "green", end=" ") + cprint(value, "blue") - def generate_number(self, temperature: Union[float, None]=None, iterations=0): + def generate_number(self, temperature: Union[float, None] = None, iterations=0): prompt = self.get_prompt() - self.debug('[generate_number]', prompt, is_prompt=True) - input_tokens = self.tokenizer.encode(prompt, return_tensors='pt').to(self.model.device) - response = self.model.generate(input_tokens, max_new_tokens=self.max_number_tokens, num_return_sequences=1, logits_processor=[self.number_logit_processor], stopping_criteria=[NumberStoppingCriteria(self.tokenizer, len(input_tokens[0]))], temperature=temperature or self.temperature, pad_token_id=self.tokenizer.eos_token_id) + self.debug("[generate_number]", prompt, is_prompt=True) + input_tokens = self.tokenizer.encode(prompt, return_tensors="pt").to( + self.model.device + ) + response = self.model.generate( + input_tokens, + max_new_tokens=self.max_number_tokens, + num_return_sequences=1, + logits_processor=[self.number_logit_processor], + stopping_criteria=[ + NumberStoppingCriteria(self.tokenizer, len(input_tokens[0])) + ], + temperature=temperature or self.temperature, + pad_token_id=self.tokenizer.eos_token_id, + ) response = self.tokenizer.decode(response[0], skip_special_tokens=True) - response = response[len(prompt):] - response = response.strip().rstrip('.') - self.debug('[generate_number]', response) + + response = response[len(prompt) :] + response = response.strip().rstrip(".") + self.debug("[generate_number]", response) try: return float(response) except ValueError: if iterations > 3: - raise ValueError('Failed to generate a valid number') - return self.generate_number(temperature=self.temperature * 1.3, iterations=iterations + 1) + raise ValueError("Failed to generate a valid number") + + return self.generate_number(temperature=self.temperature * 1.3, iterations=iterations+1) def generate_boolean(self) -> bool: prompt = self.get_prompt() - self.debug('[generate_boolean]', prompt, is_prompt=True) - input_tensor = self.tokenizer.encode(prompt, return_tensors='pt') + 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] - true_token_id = self.tokenizer.convert_tokens_to_ids('true') - false_token_id = self.tokenizer.convert_tokens_to_ids('false') + + # 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") + result = logits[true_token_id] > logits[false_token_id] - self.debug('[generate_boolean]', result) + + self.debug("[generate_boolean]", result) + return result.item() def generate_string(self) -> str: 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(input_tokens, max_new_tokens=self.max_string_token_length, num_return_sequences=1, temperature=self.temperature, stopping_criteria=[StringStoppingCriteria(self.tokenizer, len(input_tokens[0]))], pad_token_id=self.tokenizer.eos_token_id) - 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]):] + 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( + input_tokens, + max_new_tokens=self.max_string_token_length, + num_return_sequences=1, + temperature=self.temperature, + stopping_criteria=[ + StringStoppingCriteria(self.tokenizer, len(input_tokens[0])) + ], + pad_token_id=self.tokenizer.eos_token_id, + ) + + # 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] + response = self.tokenizer.decode(response, skip_special_tokens=True) - self.debug('[generate_string]', '|' + response + '|') + + self.debug("[generate_string]", "|" + response + "|") + if response.count('"') < 1: return response + return response.split('"')[0].strip() - def generate_object(self, properties: Dict[str, Any], obj: Dict[str, Any]) -> Dict[str, Any]: - for (key, schema) in properties.items(): - self.debug('[generate_object] generating value for', key) + def generate_object( + self, properties: Dict[str, Any], obj: Dict[str, Any] + ) -> Dict[str, Any]: + for key, schema in properties.items(): + self.debug("[generate_object] generating value for", key) obj[key] = self.generate_value(schema, obj, key) return obj - def generate_value(self, schema: Dict[str, Any], obj: Union[Dict[str, Any], List[Any]], key: Union[str, None]=None) -> Any: - if 'anyOf' in schema: - options = [option for option in schema['anyOf'] if 'type' in option] - if options: - schema = options[0] - else: - raise ValueError('No valid type in anyOf for key: ' + str(key)) - schema_type = schema['type'] - if schema_type == 'number': + def generate_value( + self, + schema: Dict[str, Any], + obj: Union[Dict[str, Any], List[Any]], + key: Union[str, None] = None, + ) -> Any: + schema_type = schema["type"] + if schema_type == "number": if key: obj[key] = self.generation_marker else: obj.append(self.generation_marker) return self.generate_number() - elif schema_type == 'boolean': + elif schema_type == "boolean": if key: obj[key] = self.generation_marker else: obj.append(self.generation_marker) return self.generate_boolean() - elif schema_type == 'string': + elif schema_type == "string": if key: obj[key] = self.generation_marker else: obj.append(self.generation_marker) return self.generate_string() - elif schema_type == 'array': + elif schema_type == "array": new_array = [] - if key: - obj[key] = new_array - else: - obj.append(new_array) - return self.generate_array(schema['items'], new_array) - elif schema_type == 'object': + obj[key] = new_array + return self.generate_array(schema["items"], new_array) + elif schema_type == "object": new_obj = {} if key: obj[key] = new_obj else: obj.append(new_obj) - return self.generate_object(schema['properties'], new_obj) + return self.generate_object(schema["properties"], new_obj) else: - raise ValueError(f'Unsupported schema type: {schema_type}') + raise ValueError(f"Unsupported schema type: {schema_type}") 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 + obj.append(self.generation_marker) input_prompt = self.get_prompt() obj.pop() - input_tensor = self.tokenizer.encode(input_prompt, return_tensors='pt') + input_tensor = self.tokenizer.encode(input_prompt, return_tensors="pt") 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)] + found_comma = False found_close_bracket = False + for token_id in sorted_token_ids: decoded_token = self.tokenizer.decode(token_id) if ',' in decoded_token: @@ -144,22 +214,32 @@ def generate_array(self, item_schema: Dict[str, Any], obj: Dict[str, Any]) -> li if ']' in decoded_token: found_close_bracket = True break + if found_close_bracket or not found_comma: break + return obj def get_prompt(self): - template = '{prompt}\nOutput result in the following JSON schema format:\n{schema}\nResult: {progress}' + template = """{prompt}\nOutput result in the following JSON schema format:\n{schema}\nResult: {progress}""" progress = json.dumps(self.value) gen_marker_index = progress.find(f'"{self.generation_marker}"') if gen_marker_index != -1: progress = progress[:gen_marker_index] else: - raise ValueError('Failed to find generation marker') - prompt = template.format(prompt=self.prompt, schema=json.dumps(self.json_schema), progress=progress) + raise ValueError("Failed to find generation marker") + + prompt = template.format( + prompt=self.prompt, + schema=json.dumps(self.json_schema), + progress=progress, + ) + return prompt def __call__(self) -> Dict[str, Any]: self.value = {} - generated_data = self.generate_object(self.json_schema['properties'], self.value) - return generated_data \ No newline at end of file + generated_data = self.generate_object( + self.json_schema["properties"], self.value + ) + return generated_data From 442f1323ed8a03705a8b77f6e06f625dfced9710 Mon Sep 17 00:00:00 2001 From: Arman Rafiee Date: Mon, 17 Mar 2025 17:12:09 -0400 Subject: [PATCH 6/7] Update pyproject.toml --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 2ae34a0..666e755 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "jsonformer" -version = "0.12.1" +version = "0.12.2" description = "" authors = ["1rgs ", "armand0e "] readme = "README.md" From fa3763aa9e9603e55b2caa44f1cf9b0ecd3764f7 Mon Sep 17 00:00:00 2001 From: Arman Rafiee Date: Mon, 17 Mar 2025 17:36:33 -0400 Subject: [PATCH 7/7] Update main.py --- jsonformer/main.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/jsonformer/main.py b/jsonformer/main.py index 7e2fa3b..84f95b8 100644 --- a/jsonformer/main.py +++ b/jsonformer/main.py @@ -23,10 +23,10 @@ def __init__( prompt: str, *, debug: bool = False, - max_array_length: int = 20, + max_array_length: int = 256, max_number_tokens: int = 4096, - temperature: float = 0.1, - max_string_token_length: int = 1024, + temperature: float = 1.0, + max_string_token_length: int = 2048, ): self.model = model self.tokenizer = tokenizer @@ -80,7 +80,9 @@ def generate_number(self, temperature: Union[float, None] = None, iterations=0): 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: prompt = self.get_prompt() @@ -199,7 +201,6 @@ def generate_array(self, item_schema: Dict[str, Any], obj: Dict[str, Any]) -> li 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 +209,10 @@ def generate_array(self, item_schema: Dict[str, Any], obj: Dict[str, Any]) -> li 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