Doing token checking client side

app.py

@@ -22,6 +22,10 @@ env = Environment(loader=FileSystemLoader(proj_dir / 'templates'))
 template = env.get_template('template.j2')
 template_html = env.get_template('template_html.j2')
 
+# Initialize tokenizer
+tokenizer = AutoTokenizer.from_pretrained('inception-mbzuai/jais-13b-chat')
+
+
 def check_endpoint_status():
     # Replace with the actual API URL and headers
     api_url = os.getenv("ENDPOINT_URL")
@@ -50,7 +54,7 @@ def add_text(history, text):
 
 
 def bot(history, system_prompt=""):
-    top_k = 3
+    top_k = 5
     query = history[-1][0]
 
     logger.warning('Retrieving documents...')
@@ -60,8 +64,22 @@ def bot(history, system_prompt=""):
     document_time = document_start - perf_counter()
     logger.warning(f'Finished Retrieving documents in {round(document_time, 2)} seconds...')
 
+    
+    # Function to count tokens
+    def count_tokens(text):
+        return len(tokenizer.encode(text))
+    
     # Create Prompt
     prompt = template.render(documents=documents, query=query)
+    
+    # Check if the prompt is too long
+    token_count = count_tokens(prompt)
+    while token_count > 2048:
+        # Shorten your documents here. This is just a placeholder for the logic you'd use.
+        documents.pop()  # Remove the last document
+        prompt = template.render(documents=documents, query=query)  # Re-render the prompt
+        token_count = count_tokens(prompt)  # Re-count tokens
+
     prompt_html = template_html.render(documents=documents, query=query)
     logger.warning(prompt)