使用jieba估计实时传输模式的token计数

presets.py

@@ -34,7 +34,7 @@ pre code {
 standard_error_msg = "☹️发生了错误：" # 错误信息的标准前缀
 error_retrieve_prompt = "连接超时，无法获取对话。请检查网络连接，或者API-Key是否有效。" # 获取对话时发生错误
 summarize_prompt = "请总结以上对话，不超过100字。" # 总结对话时的 prompt
-max_token_streaming = 2000 # 流式对话时的最大 token 数
+max_token_streaming = 400 # 流式对话时的最大 token 数
 timeout_streaming = 5 # 流式对话时的超时时间
 max_token_all = 3500 # 非流式对话时的最大 token 数
 timeout_all = 200 # 非流式对话时的超时时间

requirements.txt

@@ -1,3 +1,4 @@
 gradio
 mdtex2html
 pypinyin
+jieba

utils.py

@@ -12,6 +12,7 @@ import csv
 import mdtex2html
 from pypinyin import lazy_pinyin
 from presets import *
+import jieba
 
 if TYPE_CHECKING:
     from typing import TypedDict
@@ -45,6 +46,10 @@ def postprocess(
             )
         return y
 
+def count_words(input_str):
+    words = jieba.lcut(input_str)
+    return len(words)
+
 def parse_text(text):
     lines = text.split("\n")
     lines = [line for line in lines if line != ""]
@@ -89,7 +94,7 @@ def construct_assistant(text):
     return construct_text("assistant", text)
 
 def construct_token_message(token, stream=False):
-    extra = "【仅包含回答的计数】 " if stream else ""
+    extra = "【粗略计数（因为实时传输回答）】 " if stream else ""
     return f"{extra}Token 计数: {token}"
 
 def get_response(openai_api_key, system_prompt, history, temperature, top_p, stream):
@@ -125,6 +130,10 @@ def stream_predict(openai_api_key, system_prompt, history, inputs, chatbot, prev
     counter = 0
     status_text = "OK"
     history.append(construct_user(inputs))
+    if len(previous_token_count) == 0:
+        rough_user_token_count = count_words(inputs) + count_words(system_prompt)
+    else:
+        rough_user_token_count = count_words(inputs)
     try:
         response = get_response(openai_api_key, system_prompt, history, temperature, top_p, True)
     except requests.exceptions.ConnectTimeout:
@@ -148,7 +157,7 @@ def stream_predict(openai_api_key, system_prompt, history, inputs, chatbot, prev
             # decode each line as response data is in bytes
             if chunklength > 6 and "delta" in chunk['choices'][0]:
                 finish_reason = chunk['choices'][0]['finish_reason']
-                status_text = construct_token_message(sum(previous_token_count)+token_counter, stream=True)
+                status_text = construct_token_message(sum(previous_token_count)+token_counter+rough_user_token_count, stream=True)
                 if finish_reason == "stop":
                     yield get_return_value()
                     break