Update chat_func.py

chat_func.py

@@ -7,7 +7,6 @@ import json
 import os
 import requests
 import urllib3
-import numpy as np
 
 from tqdm import tqdm
 import colorama
@@ -196,8 +195,16 @@ def stream_predict(
                 chatbot[-1] = (chatbot[-1][0], partial_words+display_append)
                 all_token_counts[-1] += 1
                 yield get_return_value()
-
-
+     return  all_token_counts          
+
+def main():
+    # 调用stream_predict函数获取结果...
+    with open("token_counts.txt", "a") as f:
+        for token_count in all_token_counts:
+            f.write(str(token_count) + "\n")
+if __name__ == '__main__':
+    main()
+    
 def predict_all(
     openai_api_key,
     system_prompt,
@@ -460,12 +467,12 @@ def reduce_token_size(
 user_id = os.getuid()
 
 # 将获取到的信息写入txt文件
-all_token_counts = 10
+#all_token_counts = 10
 
-def write_to_file(user_id, all_token_counts):
-    log_file = "user_info.txt"
-    log_str = f"{user_id}\t{all_token_counts}\n"
-    with open(log_file, 'a') as f:
-        f.write(log_str)
+#def write_to_file(user_id, all_token_counts):
+#    log_file = "user_info.txt"
+#    log_str = f"{user_id}\t{all_token_counts}\n"
+#    with open(log_file, 'a') as f:
+#        f.write(log_str)
 
-write_to_file(user_id,all_token_counts)
+#write_to_file(user_id,all_token_counts)