AI agent 大模型实战之 RAG检索本地知识库

运行demo：

下载大模型

from modelscope import snapshot_download
model_dir = snapshot_download(
model_id='BAAI/bge-large-zh-v1.5',
cache_dir = '/Volumes/c/workspace/ai-agent-test/rag/model'
)

安装 向量数据库chromadb

client = chromadb.PersistentClient("./chroma_v2") #文件数据库
collections = client.get_or_create_collection(
name="my_db",
metadata={
"hnsw:space": "cosine"
}
)

读取本地txt  txt  =>  embeding

path_list = list(Path("docs").glob("*.txt"))
text_list = []
for path in path_list:
text = path.read_text(encoding="utf-8")
text_list.append(text)
embeddings = model.encode(text_list) #进行向量嵌入
collections.add(
embeddings=embeddings.tolist(),
documents=text_list,
metadatas=[{"id":i} for i,_ in enumerate(text_list)],
ids=[f"doc_{i}" for i,_ in enumerate(text_list)],
)
print(f"数据库中的数据量:{collections.count()}")

查询 embeding   =>  txt

query=["2023.12.1 做了什么"]
query_embedding = model.encode(query)
data = collections.query(query_embedding.tolist(),n_results=5)