test/amis2
1
0
Fork 0
mirror of https://gitee.com/baidu/amis.git synced 2024-11-29 18:48:45 +08:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
吴多益 2024-02-22 14:45:20 +08:00 committed by GitHub
parent 7e03085faf
commit dd34bf9883
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 12 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
scripts/bot/.gitignore vendored
View File

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

14
scripts/bot/gen_embedding.py
View File

@ -2,7 +2,7 @@ import sys
import os import os
import glob import glob
import uuid import uuid
import pickle import json
from embedding import get_embedding from embedding import get_embedding
from split_markdown import split_markdown from split_markdown import split_markdown
from vector_store import get_client from vector_store import get_client
@ -21,11 +21,11 @@ text_blocks_by_id = {}
embedding_cache = {} embedding_cache = {}
embedding_cache_file = os.path.join( 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): if os.path.exists(embedding_cache_file):
with open(embedding_cache_file, 'rb') as f: with open(embedding_cache_file, 'rb') as f:
embedding_cache = pickle.load(f) embedding_cache = json.load(f)
def get_embedding_with_cache(text): def get_embedding_with_cache(text):
@ -65,8 +65,8 @@ for filename in glob.iglob(doc_dir + '**/*.md', recursive=True):
) )
with open(os.path.join(os.path.dirname(__file__), 'text.pickle'), 'wb') as f: with open(os.path.join(os.path.dirname(__file__), 'text.json'), 'w') as f:
pickle.dump(text_blocks_by_id, f, pickle.HIGHEST_PROTOCOL) json.dump(text_blocks_by_id, f)
with open(embedding_cache_file, 'wb') as f: with open(embedding_cache_file, 'w') as f:
pickle.dump(embedding_cache, f, pickle.HIGHEST_PROTOCOL) json.dump(embedding_cache, f)

6
scripts/bot/gui.py
View File

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