Merge pull request #59 from eren23/content-reset-turbo

Content reset turbo
geeks-of-data · Mar 23, 2023 · fd77e4f · fd77e4f
2 parents 23a6284 + 2132aa7
commit fd77e4f
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 8 deletions.
diff --git a/knowledgegpt/extractors/base_extractor.py b/knowledgegpt/extractors/base_extractor.py
@@ -71,8 +71,7 @@ def extract(self, query, max_tokens, load_index=False) -> tuple[str, str, list]:
             else:
                 raise Exception("You can not load index when you have df and embeddings")
         else:
-            # import pandas as pd
-            # import numpy as np
+
             import pickle
             self.prepare_df()
             self.set_embeddings()
@@ -87,7 +86,6 @@ def extract(self, query, max_tokens, load_index=False) -> tuple[str, str, list]:
                     print("Directory ", self.index_path, " already exists")
 
                 self.df.to_csv(self.index_path + "/df.csv")
-                # np.savez(self.index_path + "/embeddings.npy", self.embeddings)
                 with open(self.index_path + "/embeddings.pkl", "wb") as f:
                     pickle.dump(self.embeddings, f)
 

diff --git a/knowledgegpt/utils/utils_completion.py b/knowledgegpt/utils/utils_completion.py
@@ -3,6 +3,7 @@
 import openai
 import pandas as pd
 import numpy as np
+import tiktoken
 
 COMPLETIONS_API_PARAMS = {
     "temperature": 0.0,
@@ -26,7 +27,6 @@ def answer_query_with_context(
         model_lang: str = "en",
         is_turbo: str = False,
         messages: list = None,
-        is_first_time: bool = True,
         index_type: str = "basic",
         max_tokens=1000
 ) -> str:
@@ -45,10 +45,7 @@ def answer_query_with_context(
     :return: The answer to the query.
 
     """
-
-    COMPLETIONS_API_PARAMS["max_tokens"] = max_tokens
-    COMPLETIONS_API_PARAMS_TURBO["max_tokens"] = max_tokens
-
+
     if len(messages) < 3 or not is_turbo:
         prompt = construct_prompt(
             verbose=verbose,
@@ -67,16 +64,34 @@ def answer_query_with_context(
         if is_turbo:
             messages.append({"role": "user", "content": prompt})
 
+    encoding = encoding = tiktoken.get_encoding("gpt2")
+    if is_turbo:
+        messages_token_length = encoding.encode(str(messages))
+        if len(messages_token_length) > 4096:
+
+            del messages[2:4]
+
+
+
     if not verbose:
         print(prompt)
 
+
     if not is_turbo:
+        prompt_len = len(encoding.encode(prompt))
+        COMPLETIONS_API_PARAMS["max_tokens"] =  2000 - prompt_len
+
         response = openai.Completion.create(
             prompt=prompt,
             **COMPLETIONS_API_PARAMS
         )
     else:
+
+        messages_token_length = encoding.encode(str(messages))
+        COMPLETIONS_API_PARAMS_TURBO["max_tokens"] = 4096 - len(messages_token_length)
+
         response = openai.ChatCompletion.create(
+
             messages=messages,
             **COMPLETIONS_API_PARAMS_TURBO,
         )