vectorize OutputNumbersTokens logit warper

jsonformer/logits_processors.py

@@ -2,7 +2,6 @@ from typing import List
 from transformers import PreTrainedTokenizer, LogitsWarper, StoppingCriteria
 import torch
 
-
 class StringStoppingCriteria(StoppingCriteria):
     def __init__(self, tokenizer: PreTrainedTokenizer, prompt_length: int):
         self.tokenizer = tokenizer
@@ -62,27 +61,24 @@ class NumberStoppingCriteria(StoppingCriteria):
 
         return False
 
-
 class OutputNumbersTokens(LogitsWarper):
     def __init__(self, tokenizer: PreTrainedTokenizer, prompt: str):
-        self.whitelist_tokens = []
-        self.tokenized_prompt = tokenizer(prompt, return_tensors="pt")
         self.tokenizer = tokenizer
+        self.tokenized_prompt = tokenizer(prompt, return_tensors="pt")
+        vocab_size = len(tokenizer)
+        self.allowed_mask = torch.zeros(vocab_size, dtype=torch.bool)
 
         for _, token_id in tokenizer.get_vocab().items():
-            token_str = tokenizer.decode(token_id)
-            token_str = token_str.strip()
+            token_str = tokenizer.decode(token_id).strip()
 
             if token_str == "" or (
                 all(c.isdigit() or c == "." for c in token_str)
                 and token_str.count(".") <= 1
             ):
-                self.whitelist_tokens.append(token_id)
+                self.allowed_mask[token_id] = True
 
-    def __call__(self, input_ids, scores):
-        input_ids = input_ids[:, len(self.tokenized_prompt["input_ids"][0]) :]
+    def __call__(self, _, scores):
+        mask = self.allowed_mask.expand_as(scores)
+        scores[~mask] = -float("inf")
 
-        scores[
-            :, [i for i in range(len(scores[0])) if i not in self.whitelist_tokens]
-        ] = -float("inf")
         return scores