chroma db读取PDF实操

上一篇文章简单介绍了向量数据chroma db,今天一起来学习下具体如何将pdf的内容写入chroma db，然后再根据openai来根据pdf的内容进行回复。

第一步：先进行pdf文档加载

from pdfminer.high_level import extract_pages
from pdfminer.layout import LTTextContainer


def pdf_extract(filename, min_line_length=1):

    paragraphs = []
    buffer = ''
    full_text = ''

    for i, page_layout in enumerate(extract_pages(filename)):
        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

第二步：初始化向量数据库并将向量化文档入库：

class MyVectorDBConnector:
    def __init__(self, collection_name, embedding_fn):
        chroma_client = chromadb.Client(Settings(allow_reset=True))

        chroma_client.reset()

        # 创建一个collection
        self.collection = chroma_client.get_or_create_collection(
            name = collection_name)
        self.embedding_fn = embedding_fn

    def add_documents(self, documents):
        '''向 collection 中添加文档与向量'''
        self.collection.add(
            embeddings=self.embedding_fn(documents),
            documents=documents,
            ids=[f"id{i}" for i in range(len(documents))]
        )

最后一步就是可以检索查询内容了：

    def search(self, query_text, top_n=3):
        '''检索向量数据库'''
        results = self.collection.query(
            query_embeddings=self.embedding_fn([query_text]),
            n_results=top_n
        )
        return results

以上步骤就是简单的将pdf文档内容保存到向量数据库chroma并检索过程，如果想要进行内容对话，我们可以将检索出来的内容再调调gpt进行对话内容操作。此过程不在此说明。