chore: 去掉 bot 示例中的 pickle 依赖 (baidu#9640)

scripts/bot/.gitignore

@@ -1,7 +1,7 @@
 db
 __pycache__
-text.pickle
-embedding.pickle
+text.json
+embedding.json
 .env
 m3e-base
 flagged

scripts/bot/gen_embedding.py

@@ -2,7 +2,7 @@
 import os
 import glob
 import uuid
-import pickle
+import json
 from embedding import get_embedding
 from split_markdown import split_markdown
 from vector_store import get_client
@@ -21,11 +21,11 @@
 embedding_cache = {}
 
 embedding_cache_file = os.path.join(
-    os.path.dirname(__file__), 'embedding.pickle')
+    os.path.dirname(__file__), 'embedding.json')
 
 if os.path.exists(embedding_cache_file):
     with open(embedding_cache_file, 'rb') as f:
-        embedding_cache = pickle.load(f)
+        embedding_cache = json.load(f)
 
 
 def get_embedding_with_cache(text):
@@ -65,8 +65,8 @@ def get_embedding_with_cache(text):
         )
 
 
-with open(os.path.join(os.path.dirname(__file__), 'text.pickle'), 'wb') as f:
-    pickle.dump(text_blocks_by_id, f, pickle.HIGHEST_PROTOCOL)
+with open(os.path.join(os.path.dirname(__file__), 'text.json'), 'w') as f:
+    json.dump(text_blocks_by_id, f)
 
-with open(embedding_cache_file, 'wb') as f:
-    pickle.dump(embedding_cache, f, pickle.HIGHEST_PROTOCOL)
+with open(embedding_cache_file, 'w') as f:
+    json.dump(embedding_cache, f)

scripts/bot/gui.py

@@ -3,7 +3,7 @@
 from embedding import get_embedding
 import gradio as gr
 import os
-import pickle
+import json
 from llm.wenxin import Wenxin, ModelName
 from dotenv import load_dotenv
 load_dotenv()
@@ -15,8 +15,8 @@
 wenxin = Wenxin()
 
 text_blocks_by_id = {}
-with open(os.path.join(os.path.dirname(__file__), 'text.pickle'), 'rb') as f:
-    text_blocks_by_id = pickle.load(f)
+with open(os.path.join(os.path.dirname(__file__), 'text.json'), 'rb') as f:
+    text_blocks_by_id = json.load(f)
 
 
 def get_prompt(context, query):