Chroma向量数据库使用案例

​​​​​基于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向量数据部分代码示例

引用

import chromadb
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 中添加文档与向量

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

检索向量数据库

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

• beg模型代码示例

引入

from sentence_transformers import SentenceTransformer

创建模型

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

内容向量化embedding_fn

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

完整源码：

import chromadb
from chromadb.config import Settings
from pdfminer.high_level import extract_pages
from pdfminer.layout import LTTextContainer
from sentence_transformers import SentenceTransformer
 
class MyVectorDB:
    '''
    私有化向量数据库
    \n作者：George
    \n时间：2024年3月7日
    '''
 
 
    def __init__(self, collection_name):
        chroma_client = chromadb.Client(Settings(allow_reset=True))
        model = SentenceTransformer('BAAI/bge-large-zh-v1.5')
 
        # 为了演示，实际不需要每次 reset()
        chroma_client.reset()
 
        # 创建一个 collection
        self.collection = chroma_client.get_or_create_collection(name=collection_name)
        self.bge_model = model
    
 
    def add_documents(self, filename, page_numbers=None, min_line_length=1, metadata={}):
        paragraphs = self.extract_text_from_pdf(filename, page_numbers, min_line_length)
        '''向 collection 中添加文档与向量'''
        self.collection.add(
            embeddings=self.embedding_fn(paragraphs),  # 每个文档的向量
            documents=paragraphs,  # 文档的原文
            ids=[f"id{i}" for i in range(len(paragraphs))]  # 每个文档的 id
        )
 
    def search(self, query, top_n):
        '''检索向量数据库'''
        results = self.collection.query(
            query_embeddings=self.embedding_fn([query]),
            n_results=top_n
        )
        return results
     
    def embedding_fn(self, paragraphs):
        '''文本向量化'''
        doc_vecs = [
            self.bge_model.encode(doc, normalize_embeddings=True).tolist()
            for doc in paragraphs
        ]
        return doc_vecs
     
    def extract_text_from_pdf(self, filename, page_numbers=None, min_line_length=1):
        '''从 PDF 文件中（按指定页码）提取文字'''
        paragraphs = []
        buffer = ''
        full_text = ''
        # 提取全部文本
        for i, page_layout in enumerate(extract_pages(filename)):
            # 如果指定了页码范围，跳过范围外的页
            if page_numbers is not None and i not in page_numbers:
                continue
            for element in page_layout:
                if isinstance(element, LTTextContainer):
                    full_text += element.get_text() + '\n'
        # 按空行分隔，将文本重新组织成段落
        lines = full_text.split('\n')
        for text in lines:
            if len(text) >= min_line_length:
                buffer += (' '+text) if not text.endswith('-') else text.strip('-')
            elif buffer:
                paragraphs.append(buffer)
                buffer = ''
        if buffer:
            paragraphs.append(buffer)
        return paragraphs
     
 
if "__main__" == __name__:
    # 创建一个向量数据库对象
    vector_db = MyVectorDB("demo")
    # 向向量数据库中添加文档
    vector_db.add_documents("llama2.pdf", page_numbers=[
                                    2, 3], min_line_length=10)
 
    user_query = "Llama 2有多少参数"
    results = vector_db.search(user_query, 2)
 
    for para in results['documents'][0]:
        print(para+"\n") 

 总结：

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