mv

prob_jsonformer/logits_processors.py

@@ -25,7 +25,7 @@ class StringStoppingCriteria(StoppingCriteria):
         result = '"' in last_token
 
         if self.max_length is not None:
-            # because of tokens this wont work pefectly
+            # because of tokens this wont work pefectly, we might go 0-10 chars over
             gen_ids = input_ids[0][self.prompt_length :]
             o = self.tokenizer.decode(gen_ids, skip_special_tokens=True)
             str_l = len(o)

prob_jsonformer/main.py

@@ -7,7 +7,7 @@ from prob_jsonformer.logits_processors import (
     OutputIntegersTokens,
     StringStoppingCriteria,
 )
-from prob_jsonformer.choice_tree import choice_tree
+from prob_jsonformer.prob_choice_tree import prob_choice_tree
 from prob_jsonformer.type_prefixes import get_prefix_tokens_for_types
 
 from termcolor import cprint
@@ -193,7 +193,9 @@ class Jsonformer:
         choices_tokens = self.tokenizer(choices).input_ids
         choices_tokens = [torch.tensor(c) for c in choices_tokens]
 
-        r = list(choice_tree(self.model, self.tokenizer, input_ids, choices_tokens))
+        r = list(
+            prob_choice_tree(self.model, self.tokenizer, input_ids, choices_tokens)
+        )
         return r
 
     def generate_p_integer(self, range_min: float, range_max: float) -> float:

prob_jsonformer/prob_choice_tree.py

@@ -21,7 +21,7 @@ def get_valid_next_choices(choices_tokens, current_tokens):
     return torch.LongTensor(next_choices)
 
 
-def _choice_tree(
+def _prob_choice_tree(
     model: AutoModelForCausalLM,
     tokenizer: AutoTokenizer,
     input_ids: Int[Tensor, "seq"],
@@ -48,7 +48,7 @@ def _choice_tree(
         for i in range(len(next_choices)):
             next_choice = next_choices[i]
             next_prob = prob * probs[i].item()
-            yield from choice_tree(
+            yield from prob_choice_tree(
                 model=model,
                 tokenizer=tokenizer,
                 choices_tokens=choices_tokens,
@@ -59,12 +59,12 @@ def _choice_tree(
             )
 
 
-def choice_tree(
+def prob_choice_tree(
     *args,
     **kwargs,
 ):
     choice_json = list(
-        _choice_tree(
+        _prob_choice_tree(
             *args,
             **kwargs,
         )