From 0ea8c5a407361ae1155aa2a9dada5b192f0332a9 Mon Sep 17 00:00:00 2001
From: wassname <git@wassname.org>
Date: Sat, 11 May 2024 15:02:28 +0800
Subject: [PATCH] mv

---
 prob_jsonformer/logits_processors.py                    | 2 +-
 prob_jsonformer/main.py                                 | 6 ++++--
 prob_jsonformer/{choice_tree.py => prob_choice_tree.py} | 8 ++++----
 3 files changed, 9 insertions(+), 7 deletions(-)
 rename prob_jsonformer/{choice_tree.py => prob_choice_tree.py} (95%)

diff --git a/prob_jsonformer/logits_processors.py b/prob_jsonformer/logits_processors.py
index de88e37..f45e613 100644
--- a/prob_jsonformer/logits_processors.py
+++ b/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)
diff --git a/prob_jsonformer/main.py b/prob_jsonformer/main.py
index 19000e7..cfad70d 100644
--- a/prob_jsonformer/main.py
+++ b/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:
diff --git a/prob_jsonformer/choice_tree.py b/prob_jsonformer/prob_choice_tree.py
similarity index 95%
rename from prob_jsonformer/choice_tree.py
rename to prob_jsonformer/prob_choice_tree.py
index 127aa24..053492b 100644
--- a/prob_jsonformer/choice_tree.py
+++ b/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,
         )