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进行对话内容操作。此过程不在此说明。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值