reset token count

neulab · Apr 15, 2024 · 40255a1 · 40255a1
1 parent 0951e37
commit 40255a1
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/prompt2model/utils/api_tools.py b/prompt2model/utils/api_tools.py
@@ -96,12 +96,11 @@ def generate_one_completion(
             An OpenAI-like response object if there were no errors in generation.
             In case of API-specific error, Exception object is captured and returned.
         """
-        # num_prompt_tokens = count_tokens_from_string(prompt)
-        # if self.max_tokens:
-        #     max_tokens = self.max_tokens - num_prompt_tokens - token_buffer
-        # else:
-        #     max_tokens = 3 * num_prompt_tokens
-        max_tokens = self.max_tokens
+        num_prompt_tokens = count_tokens_from_string(prompt)
+        if self.max_tokens:
+            max_tokens = self.max_tokens - num_prompt_tokens - token_buffer
+        else:
+            max_tokens = 3 * num_prompt_tokens
 
         response = completion(  # completion gets the key from os.getenv
             model=self.model_name,