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)