基于检索增强的大语言模型

Question

我想建立一个基于RAG的大型语言模型。我知道通常首先加载文档，然后将文档分块，向量化并存储在向量数据库中，最后检索向量数据库的内容。我现在不明白的是，文档在分块后存储在向量数据库中，所以检索的不是块的向量吗？怎么才能返回整个文档的向量呢？

我想要的最终的答案是基于检索到的块的原始文档信息来组合LLM，但不确定如何实现。另一件让我困惑的事情是，我想要做的事基于多个文档的，如果我检索这些块，我可能会得到多个来自于不同文档的分块。但其他组块不满足我提出的问题，请问如何处理这种情况？

Answer 1

引自免费微信小程序：皆我百晓生

要构建一个基于检索增强的大型语言模型（例如RAG），你需要执行以下几个步骤：

1. 文档预处理和分块：
   • 加载文档：将文本数据集加载到内存中。
   • 分块：将每个文档分割成较小的段落或句子。这可以按照固定长度（例如，每块256个单词）或按自然段落进行。

import re

def split_into_blocks(text, block_size):
    return [text[i:i+block_size] for i in range(0, len(text), block_size)]

documents = [...]  # List of documents as strings
block_size = 256
blocks = []
for doc in documents:
    blocks.extend(split_into_blocks(doc, block_size))

1. 向量化：
   • 使用预训练的嵌入模型（如BERT、RoBERTa等）将文本块转换为向量表示。
   • 这里可以使用transformers库中的encode_plus方法。

from transformers import BertTokenizerFast, BertModel

tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')

encoded_blocks = []
for block in blocks:
    inputs = tokenizer.encode_plus(block, add_special_tokens=True)
    encoded_block = model(**inputs)[0].mean(dim=1)  # Mean pooling over tokens
    encoded_blocks.append(encoded_block)

1. 存储向量数据库：
   • 将向量化后的块存储在一个可快速检索的数据结构中，如Faiss或Annoy。

from annoy import AnnoyIndex

index = AnnoyIndex(encoded_blocks[0].shape[0], 'euclidean')  # Choose an appropriate metric
for i, vector in enumerate(encoded_blocks):
    index.add_item(i, vector.cpu().numpy())  # Assuming you're using PyTorch; use .numpy() for TensorFlow
index.build(n_trees=10)  # Tune this parameter for better performance

1. 检索与组合：
   • 当用户输入查询时，计算查询的向量表示，然后从向量数据库中找到最相关的文档块。
   • 对于每个匹配的文档块，提取其原始文本，并将其传递给大型语言模型（LLM，如GPT-3）进行解码和答案生成。
   • 可能会出现来自不同文档的块组合在一起的情况。为了处理这个问题，一种可能的方法是在组合时考虑块之间的语义连贯性，或者在生成答案时添加一些上下文提示。

query = "Your question here"
query_encoded = encode_query(query, tokenizer, model)

# Perform retrieval from the index
relevant_indices = index.get_nns_by_vector(query_encoded.cpu().numpy(), k=10)  # k is the number of top matches to retrieve

# Retrieve original text and feed it into LLM
retrieved_text = [blocks[i] for i in relevant_indices]
llm_input = "Question: " + query + "\n" + "\n".join(retrieved_text)
answer = generate_answer(llm_input, lmm_model)  # Replace with your LLM generation function

注意：这里的代码片段仅作为示例，实际应用中可能需要根据具体需求进行调整。例如，编码和检索部分可能需要优化，以适应不同的模型和数据集规模。同时，generate_answer函数需要你自己实现，它应该调用一个预训练的语言模型（如Hugging Face Transformers中的GPT-3接口）来生成答案。如果你有多个文档来源，那么在检索阶段，确保你对每个来源都创建了一个独立的索引，以便正确地检索和组合结果。