在数字化浪潮的推动下，信息检索已成为我们日常生活中不可或缺的一部分。然而，随着数据量的爆炸式增长，如何快速精准地从海量文档中检索出有价值的信息，成为了一个巨大的挑战。本文将带您走进 Pinecone 向量数据库的世界，探索如何利用这一前沿技术，结合检索增强生成（RAG）模型，实现对自定义文档的高效查询。这不仅是技术的飞跃，更是对传统检索方式的一次深刻革新。
向量数据库
向量数据库（Vector DB）现在非常流行。像 GPT-4o、Llama-3、Claude3.5 等大语言模型（LLMs）正基于专业使用案例和行业特定数据准备被行业采用。检索增强生成（RAG）—— 在推理期间将输入增强为与输入提示相关的数据，对于这些使用案例来说是一个令人兴奋的范式。
向量数据库提供了一种快速查询大量数据以找到最相关文档块的方法。与传统数据库相比，向量数据库在查询大维文本嵌入方面是高效的。
文档嵌入和标记化
首先处理文档的第一步是将其分成块，并获取每个块的嵌入。对于嵌入模型，我们使用的是北京智源人工智能研究院的嵌入模型“BGE-M3”，它是 1024 维的。
接下来，我们定义分块中允许的最大 Token。每个 Token 约为一个词的 3/4。通常，正确的数量是一个探索点，只能通过反复试验找到。如果预期答案遍布文档中的小文本部分，则较小的分块尺寸效果更好。较大的分块则更适合于更长、更详细的描述。较大的分块会倾向于从一个或几个较长的部分检索数据，而较短的部分则会偏向于具有有限上下文的较短文本片段。
python
复制代码
def get_embedding(text, model="BAAI/bge-m3"):
# 替换换行符为空格
text = text.replace("\n", " ")
# 使用指定模型获取嵌入
return client.embeddings.create(input=[text], model=model).data[0].embedding
def tokenize(text, max_tokens) -> pd.DataFrame:
""" 将文本拆分为最大令牌数的块 """
# 加载适用于bge-m3模型的分词器
tokenizer = tiktoken.get_encoding("cl100k_base")
df = pd.DataFrame({'title': ['0'], 'text': [text]})
# 对文本进行分词，并保存令牌数量到新列中
df['n_tokens'] = df.text.apply(lambda x: len(tokenizer.encode(x)))
shortened = []
# 遍历数据框
for _, row in df.iterrows():
# 如果文本为空，跳到下一行
if row['text'] is None:
continue
# 如果令牌数大于最大令牌数，将文本拆分为多个块
if row['n_tokens'] > max_tokens:
shortened += split_into_many(row['text'], tokenizer, max_tokens)
else:
shortened.append(row['text'])
df = pd.DataFrame(shortened, columns=['text'])
df['n_tokens'] = df.text.apply(lambda x: len(tokenizer.encode(x)))
df['embeddings'] = df.text.apply(lambda x: get_embedding(x))
return df
def split_into_many(text: str, tokenizer: tiktoken.Encoding, max_tokens: int = 1024) -> list:
""" 将字符串拆分为指定数量令牌的多个字符串 """
# 将文本拆分为句子
sentences = text.split(' ')
# 获取每个句子的令牌数
n_tokens = [len(tokenizer.encode(" " + sentence)) for sentence in sentences]
chunks = []
tokens_so_far = 0
chunk = []
# 遍历句子和令牌数的元组
for sentence, token in zip(sentences, n_tokens):
chunk.append(sentence)
tokens_so_far += token + 1
# 如果令牌数加上当前句子的令牌数大于最大令牌数，将块添加到块列表中并重置
if tokens_so_far + token > max_tokens:
chunks.append(" ".join(chunk))
chunk = []
tokens_so_far = 0
return chunks
df = tokenize(text, 100)
这里是一个最大块大小为 100 个标记的示例数据框架：
创建向量数据库索引并上传数据
创建了数据框架后，就可以为嵌入创建索引，并将它们上传到 Pinecone 数据库中。这里注意，因为 Pinecone 不允许一次性添加大量的块，所以意味着需要一些批处理过程。
python
复制代码
def __init__(self, batch_size: int = 10) -> None:
self.batch_size = batch_size
# 从输入DataFrame中生成块
def to_batches(self, df: pd.DataFrame) -> Iterator[pd.DataFrame]:
splits = self.splits_num(df.shape[0])
if splits <= 1:
yield df
else:
for chunk in np.array_split(df, splits):
yield chunk
# 确定DataFrame包含多少块
def splits_num(self, elements: int) -> int:
return round(elements / self.batch_size)
__call__ = to_batches
df_batcher = BatchGenerator(300)
index_name = 'overview-of-nvidia'
# Check whether the index with the same name already exists - if so, delete it
if index_name in pc.list_indexes():
pc.delete_index(index_name)
pc.create_index(index_name, dimension=len(df['content_vector'][0]), spec=ServerlessSpec(
cloud="aws",
region="us-east-1"
))
index = pc.Index(name=index_name)
# 上传内容向量
for batch_df in df_batcher(df):
index.upsert(vectors=zip(batch_df.vector_id, batch_df.content_vector))
查询数据库
添加向量之后，只需要输入一个相同维度（1024）的查询向量，以找到具有最接近嵌入和相似度分数的行。下面的函数将这些行映射到原始数据框，并提供文本和余弦相似度。
ini
复制代码
content_mapped = dict(zip(df.vector_id,df.text))
# 使用指定命名空间中的标题查询文章并打印结果。
def query_article(query, top_k=3):
# 基于标题列创建向量嵌入
embedded_query = client.embeddings.create(input = [query], model="BAAI/bge-m3").data[0].embedding
# 使用标题向量作为参数传递的查询命名空间
query_result = index.query(vector=embedded_query,
top_k=3)
# 查询结果
print(f'\n查询结果： {query}')
if not query_result.matches:
print('没有查询结果')
matches = query_result.matches
ids = [res.id for res in matches]
scores = [res.score for res in matches]
df = pd.DataFrame({'id':ids,
'score':scores,
'content': [content_mapped[_id] for _id in ids],
})
counter = 0
for k,v in df.iterrows():
counter += 1
print(f'{v.content} (score = {v.score})')
print('\n')
return df
查询及结果如下：
Pinecone 并非唯一的向量数据库提供者，还有如 Weviate、Chroma 等，它们都采用类似的架构来创建索引、上传和查询文档。像 Langchain 和 Llamaindex 这样的框架抽象了大部分样板代码，有助于简化对 RAG 的采用。
结语
跟随本文的介绍，我们见证了 Pinecone 向量数据库在文档检索领域的强大潜力。通过 RAG 技术，我们不仅提升了检索的速度和准确性，更开启了对文档深层次理解和应用的新视角。这项技术的运用，不仅限于信息检索，更可能成为推动各行各业智能化转型的关键力量。在未来，我们有理由相信，向量数据库和 RAG 技术将在构建更加智能、高效的信息生态系统中发挥更加重要的作用。
