Update logits_processors.py

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")

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
 

pyproject.toml

@@ -1,8 +1,8 @@
 [tool.poetry]
 name = "jsonformer"
-version = "0.12.0"
+version = "0.12.2"
 description = ""
-authors = ["1rgs <rgsduke@gmail.com>"]
+authors = ["1rgs <rgsduke@gmail.com>", "armand0e <aran.rafiee@gmail.com>"]
 readme = "README.md"
 
 [tool.poetry.dependencies]
@@ -23,4 +23,4 @@ build-backend = "poetry.core.masonry.api"
 
 [virtualenvs]
 create = true
-in-project = true
+in-project = true