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数据库，也希望大家一起进步，有问题也可以评论相互学习！