From 9810790981a81ffe442f39f54a73c8961891a3bb Mon Sep 17 00:00:00 2001
From: wassname <git@wassname.org>
Date: Sat, 11 May 2024 10:54:29 +0800
Subject: [PATCH] add maxLength for str

---
 prob_jsonformer/logits_processors.py | 24 +++++++++++++++++++-----
 prob_jsonformer/main.py              |  8 +++++---
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/prob_jsonformer/logits_processors.py b/prob_jsonformer/logits_processors.py
index 1af968d..4e66d26 100644
--- a/prob_jsonformer/logits_processors.py
+++ b/prob_jsonformer/logits_processors.py
@@ -4,9 +4,13 @@ import torch
 
 
 class StringStoppingCriteria(StoppingCriteria):
-    def __init__(self, tokenizer: PreTrainedTokenizer, prompt_length: int):
+    def __init__(
+        self, tokenizer: PreTrainedTokenizer, prompt_length: int, max_length: int = None
+    ):
         self.tokenizer = tokenizer
         self.prompt_length = prompt_length
+        self.max_length = max_length
+        print(max_length, ", max_length")
 
     def __call__(
         self,
@@ -21,6 +25,12 @@ class StringStoppingCriteria(StoppingCriteria):
 
         result = '"' in last_token
 
+        if self.max_length is not None:
+            str_l = len(self.tokenizer.decode(input_ids[0], skip_special_tokens=True))
+            if str_l > self.max_length:
+                print("maxlen", str_l)
+                return True
+
         return result
 
 
@@ -52,7 +62,7 @@ class NumberStoppingCriteria(StoppingCriteria):
             and len(decoded.replace(" ", "").split(".")[1]) > self.precision
         ):
             return True
-        
+
         if (
             len(decoded) > 1
             and "," in decoded
@@ -85,7 +95,8 @@ class OutputNumbersTokens(LogitsWarper):
                 or (
                     all(c.isdigit() or c == "." for c in token_str)
                     and token_str.count(".") <= 1
-                ) or (
+                )
+                or (
                     "," in token_str
                     and all(c.isdigit() or c == "." for c in token_str.split(",")[0])
                     and token_str.count(".") <= 1
@@ -99,6 +110,7 @@ class OutputNumbersTokens(LogitsWarper):
 
         return scores
 
+
 class IntegerStoppingCriteria(StoppingCriteria):
     def __init__(
         self,
@@ -128,7 +140,7 @@ class IntegerStoppingCriteria(StoppingCriteria):
             and any(c.isdigit() for c in decoded.split(",")[0])
         ):
             return True
-        
+
         if (
             len(decoded) > 1
             and any(c.isdigit() for c in decoded)
@@ -138,6 +150,7 @@ class IntegerStoppingCriteria(StoppingCriteria):
 
         return False
 
+
 class OutputIntegersTokens(LogitsWarper):
     def __init__(self, tokenizer: PreTrainedTokenizer, prompt: str):
         self.tokenizer = tokenizer
@@ -151,7 +164,8 @@ class OutputIntegersTokens(LogitsWarper):
             if (
                 token_str == ""
                 or all(c.isdigit() for c in token_str)
-                or "," in token_str and all(c.isdigit() for c in token_str.split(",")[0])
+                or "," in token_str
+                and all(c.isdigit() for c in token_str.split(",")[0])
             ):
                 self.allowed_mask[token_id] = True
 
diff --git a/prob_jsonformer/main.py b/prob_jsonformer/main.py
index d23b0ab..34d1d9b 100644
--- a/prob_jsonformer/main.py
+++ b/prob_jsonformer/main.py
@@ -144,7 +144,7 @@ class Jsonformer:
 
         return result.item()
 
-    def generate_string(self) -> str:
+    def generate_string(self, maxLength=None) -> str:
         prompt = self.get_prompt() + '"'
         self.debug("[generate_string]", prompt, is_prompt=True)
         input_tokens = self.tokenizer.encode(prompt, return_tensors="pt").to(
@@ -157,7 +157,7 @@ class Jsonformer:
             num_return_sequences=1,
             temperature=self.temperature,
             stopping_criteria=[
-                StringStoppingCriteria(self.tokenizer, len(input_tokens[0]))
+                StringStoppingCriteria(self.tokenizer, len(input_tokens[0]), maxLength)
             ],
             pad_token_id=self.tokenizer.eos_token_id,
         )
@@ -313,7 +313,9 @@ class Jsonformer:
                 obj[key] = self.generation_marker
             else:
                 obj.append(self.generation_marker)
-            return self.generate_string()
+            return self.generate_string(
+                schema["maxLength"] if "maxLength" in schema else None
+            )
         elif schema_type == "choice_probs":
             if key:
                 obj[key] = self.generation_marker