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") diff --git a/jsonformer/main.py b/jsonformer/main.py index 9c13471..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 = 10, - max_number_tokens: int = 6, + max_array_length: int = 256, + max_number_tokens: int = 4096, temperature: float = 1.0, - max_string_token_length: int = 10, + 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 diff --git a/pyproject.toml b/pyproject.toml index 7014fed..666e755 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,8 +1,8 @@ [tool.poetry] name = "jsonformer" -version = "0.12.0" +version = "0.12.2" 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