Chroma向量数据库使用案例

 

转载自：https://blog.csdn.net/xzq_qzx_/article/details/136535125

​​​​​基于ChromaDB与BGEToward-VectorModel的本地私有化向量检索之路

安装基础环境

包括Python安装、pip安装及镜像设置等网上提供很多资料，这里就不介绍了。

• 安装chromaDB环境安装

pip install chromadb

 

• 安装pdf解析库

pip install pdfminer.six

 

• 安装模型库

pip install sentence_transformers

 

• 下载bge-large-zh-v1.5向量模型

如果能访问huggingface执行程序时自动下载，如果不能访问huggingface，请点击以下网盘链接进行下载：

链接:百度网盘 请输入提取码 提取码: fpej

手动下载模型需解压到项目工程目录，即与MyVectorDb.py在同一目录。

实践

在最后的环节，我们即将把理论化为实践。我们将用Python编写出一套基于chromadb向量数据库和bge-large-zh-v1.5向量模型实现本地向量检索的代码。

• chromadb向量数据部分代码示例

引用

1.  
   import chromadb
2.  
   from chromadb.config import Settings

创建数据库对象

chroma_client = chromadb.Client(Settings(allow_reset=True))

 

创建一个 collection，即创建一个名为“demo”的数据库

collection = chroma_client.get_or_create_collection(name="demo")

 

向 collection 中添加文档与向量

1.  
   collection.add(
2.  
   embeddings= embedding_fn(documents), # 每个文档的向量，这里调用了向量模型，将文档转化成向量
3.  
   documents=documents, # 文档的原文
4.  
   ids=[f"id{i}" for i in range(len(documents))] # 每个文档的 id
5.  
   )

检索向量数据库

1.  
   collection.query(
2.  
   query_embeddings=embedding_fn([query]), # 检索内容的向量，这里调用了向量模型，将文档转化成向量
3.  
   n_results=top_n
4.  
   )

• beg模型代码示例

引入

from sentence_transformers import SentenceTransformer

 

创建模型

model = SentenceTransformer('BAAI/bge-large-zh-v1.5')

 

内容向量化embedding_fn

1.  
   doc_vecs = [
2.  
   model.encode(doc, normalize_embeddings=True).tolist()
3.  
   for doc in documents
4.  
   ]

完整源码：

1.  
   import chromadb
2.  
   from chromadb.config import Settings
3.  
   from pdfminer.high_level import extract_pages
4.  
   from pdfminer.layout import LTTextContainer
5.  
   from sentence_transformers import SentenceTransformer
6.  
    
7.  
   class MyVectorDB:
8.  
   '''
9.  
   私有化向量数据库
10.  
    \n作者：George
11.  
    \n时间：2024年3月7日
12.  
    '''
13.  
     
14.  
     
15.  
    def __init__(self, collection_name):
16.  
    chroma_client = chromadb.Client(Settings(allow_reset=True))
17.  
    model = SentenceTransformer('BAAI/bge-large-zh-v1.5')
18.  
     
19.  
    # 为了演示，实际不需要每次 reset()
20.  
    chroma_client.reset()
21.  
     
22.  
    # 创建一个 collection
23.  
    self.collection = chroma_client.get_or_create_collection(name=collection_name)
24.  
    self.bge_model = model
25.  
     
26.  
     
27.  
    def add_documents(self, filename, page_numbers=None, min_line_length=1, metadata={}):
28.  
    paragraphs = self.extract_text_from_pdf(filename, page_numbers, min_line_length)
29.  
    '''向 collection 中添加文档与向量'''
30.  
    self.collection.add(
31.  
    embeddings=self.embedding_fn(paragraphs), # 每个文档的向量
32.  
    documents=paragraphs, # 文档的原文
33.  
    ids=[f"id{i}" for i in range(len(paragraphs))] # 每个文档的 id
34.  
    )
35.  
     
36.  
    def search(self, query, top_n):
37.  
    '''检索向量数据库'''
38.  
    results = self.collection.query(
39.  
    query_embeddings=self.embedding_fn([query]),
40.  
    n_results=top_n
41.  
    )
42.  
    return results
43.  
     
44.  
    def embedding_fn(self, paragraphs):
45.  
    '''文本向量化'''
46.  
    doc_vecs = [
47.  
    self.bge_model.encode(doc, normalize_embeddings=True).tolist()
48.  
    for doc in paragraphs
49.  
    ]
50.  
    return doc_vecs
51.  
     
52.  
    def extract_text_from_pdf(self, filename, page_numbers=None, min_line_length=1):
53.  
    '''从 PDF 文件中（按指定页码）提取文字'''
54.  
    paragraphs = []
55.  
    buffer = ''
56.  
    full_text = ''
57.  
    # 提取全部文本
58.  
    for i, page_layout in enumerate(extract_pages(filename)):
59.  
    # 如果指定了页码范围，跳过范围外的页
60.  
    if page_numbers is not None and i not in page_numbers:
61.  
    continue
62.  
    for element in page_layout:
63.  
    if isinstance(element, LTTextContainer):
64.  
    full_text += element.get_text() + '\n'
65.  
    # 按空行分隔，将文本重新组织成段落
66.  
    lines = full_text.split('\n')
67.  
    for text in lines:
68.  
    if len(text) >= min_line_length:
69.  
    buffer += (' '+text) if not text.endswith('-') else text.strip('-')
70.  
    elif buffer:
71.  
    paragraphs.append(buffer)
72.  
    buffer = ''
73.  
    if buffer:
74.  
    paragraphs.append(buffer)
75.  
    return paragraphs
76.  
     
77.  
     
78.  
    if "__main__" == __name__:
79.  
    # 创建一个向量数据库对象
80.  
    vector_db = MyVectorDB("demo")
81.  
    # 向向量数据库中添加文档
82.  
    vector_db.add_documents("llama2.pdf", page_numbers=[
83.  
    2, 3], min_line_length=10)
84.  
     
85.  
    user_query = "Llama 2有多少参数"
86.  
    results = vector_db.search(user_query, 2)
87.  
     
88.  
    for para in results['documents'][0]:
89.  
    print(para+"\n")

 总结：

这只是一个简单的演示样例，方便大家进一步理解和操作Chroma数据库，也希望大家一起进步，有问题也可以评论相互学习！