feat: add estimation with max token

GPT/query.py

@@ -16,6 +16,7 @@ _ = language.set_language()
 
 
 def build(chunk_size=4000):
+    openai.api_key = API_KEY
     all_text = util.read_file(r'.user\input.txt')
 
     # split text into smaller chunk of 4000 char each

Seanium_Brain.py

@@ -64,12 +64,15 @@ with st.sidebar:
                                   "your prompt plus `max_tokens` cannot exceed the model's context length. Most "
                                   "models have a context length of 2048 tokens (except for the newest models, "
                                   "which support 4096)."))
-    chunk_size = st.slider(_('Chunk size'), 1500, 4500,
-                           value=util.read_json_at(INFO.BRAIN_MEMO, 'chunk_size', 4000),
-                           help=_("The number of tokens to consider at each step. The larger this is, the more "
-                                  "context the model has to work with, but the slower generation and expensive "
-                                  "will it be."))
-
+    col1, col2 = st.columns([3, 1])
+    with col1:
+        chunk_size = st.slider(_('Chunk size'), 1500, 4500,
+                               value=util.read_json_at(INFO.BRAIN_MEMO, 'chunk_size', 4000),
+                               help=_("The number of tokens to consider at each step. The larger this is, the more "
+                                      "context the model has to work with, but the slower generation and expensive "
+                                      "will it be."))
+    with col2:
+        update_brain = st.button(_('Update Brain'))
     with st.expander(label=_('Advanced Options')):
         top_p = st.slider(_('Top_P'), 0.0, 1.0, value=util.read_json_at(INFO.BRAIN_MEMO, 'top_p', 1.0),
                           help=_("An alternative to sampling with temperature, called nucleus sampling, where the "
@@ -144,8 +147,15 @@ with body:
         if os.path.exists(CURRENT_LOG_FILE):
             st_tool.download_as(_("📥download log"))
     # execute brain calculation
+    if update_brain:
+        st_tool.rebuild_brain(chunk_size)
     if not query == '':
-        st.markdown(f'Token estimation: `{st_tool.predict_token(query, prompt_core)}`')
+        if models.question_model == 'text-davinci-003' or 'text-davinci-003' in models.other_models:
+            max_model_token = 4000
+        else:
+            max_model_token = 2048
+
+        st.markdown(f'Token estimation: `{st_tool.predict_token(query, prompt_core)}/{max_model_token}`')
         if send:
             st_tool.execute_brain(query,
                                   param,

streamlit_toolkit/tools.py

@@ -27,6 +27,7 @@ def predict_token(query: str, prompt_core: GPT.model.prompt_core) -> int:
                                                            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)
@@ -259,6 +260,15 @@ def process_response_stream(query, target_model, prompt_file: str, params: GPT.m
     log(previous_chars, delimiter=f'{file_name.upper()}')
 
 
+def rebuild_brain(chunk_size: int):
+    msg = st.warning(_('Updating Brain...'), icon="⏳")
+    progress_bar = st.progress(0)
+    for idx, chunk_num in GPT.query.build(chunk_size):
+        progress_bar.progress((idx + 1) / chunk_num)
+    msg.success(_('Brain Updated!'), icon="👍")
+    time.sleep(2)
+
+
 def execute_brain(q, params: GPT.model.param,
                   op: GPT.model.Operation,
                   model: GPT.model.Model,
@@ -272,12 +282,7 @@ def execute_brain(q, params: GPT.model.param,
     log(f'\n\n\n\n[{str(time.ctime())}] - QUESTION: {q}')
 
     if mod.check_update.is_input_updated() or mod.check_update.is_param_updated(params.chunk_size, 'chunk_size'):
-        msg = st.warning(_('Updating Brain...'), icon="⏳")
-        progress_bar = st.progress(0)
-        for idx, chunk_num in GPT.query.build(params.chunk_size):
-            progress_bar.progress((idx + 1) / chunk_num)
-        msg.success(_('Brain Updated!'), icon="👍")
-        time.sleep(2)
+        rebuild_brain(params.chunk_size)
 
     # =================stream=================
     if stream: