feat: add token estimation

GPT/query.py

@@ -11,7 +11,6 @@ API_KEY = util.read_file(r'.user\API-KEYS.txt').strip()
 
 openai.api_key = API_KEY
 
-
 SESSION_LANG = st.session_state['SESSION_LANGUAGE']
 _ = language.set_language()
 
@@ -61,21 +60,23 @@ def run(query, model, prompt_file, isQuestion, params, info_file=None):
     return all_response
 
 
-def run_stream(query, model, prompt_file, isQuestion, params, info_file=None):
+def get_stream_prompt(query, prompt_file, isQuestion, info_file=None):
-    client = None
+    openai.api_key = API_KEY
     if isQuestion:
         data = util.read_json(INFO.BRAIN_DATA)
-        results = GPT.toolkit.search_chunks(query, data, count=1)
+        result = GPT.toolkit.search_chunks(query, data, count=1)
-        for result in results:
+        my_info = util.read_file(info_file)
-            my_info = util.read_file(info_file)
+        prompt = util.read_file(prompt_file)
-            prompt = util.read_file(prompt_file)
+        prompt = prompt.replace('<<INFO>>', result[0]['content'])
-            prompt = prompt.replace('<<INFO>>', result['content'])
+        prompt = prompt.replace('<<QS>>', query)
-            prompt = prompt.replace('<<QS>>', query)
+        prompt = prompt.replace('<<MY-INFO>>', my_info)
-            prompt = prompt.replace('<<MY-INFO>>', my_info)
-            client = GPT.toolkit.gpt3_stream(API_KEY, prompt, model, params)
-
     else:
         chunk = textwrap.wrap(query, 10000)[0]
         prompt = util.read_file(prompt_file).replace('<<DATA>>', chunk)
-        client = GPT.toolkit.gpt3_stream(API_KEY, prompt, model, params)
+    return prompt
+
+
+def run_stream(query, model, prompt_file, isQuestion, params, info_file=None):
+    prompt = get_stream_prompt(query, prompt_file, isQuestion, info_file)
+    client = GPT.toolkit.gpt3_stream(API_KEY, prompt, model, params)
     return client

Seanium_Brain.py

@@ -136,7 +136,7 @@ with header:
 
 # main
 with body:
-    question = st.text_area(_('Ask Brain: '))
+    query = st.text_area(_('Ask Brain: '))
     col1, col2 = st.columns([1, 3])
     with col1:
         send = st.button(_('📩Send'))
@@ -144,13 +144,15 @@ with body:
         if os.path.exists(CURRENT_LOG_FILE):
             st_tool.download_as(_("📥download log"))
     # execute brain calculation
-    if not question == '' and send:
+    if not query == '':
-        st_tool.execute_brain(question,
+        st.markdown(f'Token estimation: `{st_tool.predict_token(query, prompt_core)}`')
-                              param,
+        if send:
-                              op,
+            st_tool.execute_brain(query,
-                              models,
+                                  param,
-                              prompt_core,
+                                  op,
-                              prompt_dictionary,
+                                  models,
-                              _('question'),
+                                  prompt_core,
-                              enable_stream,
+                                  prompt_dictionary,
-                              SESSION_LANG)
+                                  _('question'),
+                                  enable_stream,
+                                  SESSION_LANG)

streamlit_toolkit/tools.py

@@ -4,6 +4,7 @@ import json
 import streamlit as st
 import tkinter as tk
 from tkinter import filedialog
+from langchain.llms import OpenAI
 
 import modules.utilities as util
 import modules.INFO as INFO
@@ -18,6 +19,14 @@ SESSION_TIME = st.session_state['SESSION_TIME']
 CURRENT_LOG_FILE = f'{INFO.LOG_PATH}/log_{SESSION_TIME}.log'
 
 
+def predict_token(query: str, prompt_core: GPT.model.prompt_core) -> int:
+    """predict how many tokens to generate"""
+    llm = OpenAI()
+    token = llm.get_num_tokens(GPT.query.get_stream_prompt(query, prompt_file=prompt_core.question,
+                                                           isQuestion=True,
+                                                           info_file=prompt_core.my_info))
+    return token
+
 def create_log():
     if not os.path.exists(CURRENT_LOG_FILE):
         util.write_file(f'Session {SESSION_TIME}\n\n', CURRENT_LOG_FILE)