上一篇文章简单介绍了向量数据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进行对话内容操作。此过程不在此说明。