这篇文章分为三个部分：
- 查询扩展和重构
- 混合检索：密集和稀疏方法
- 带重新排序的多阶段检索
查询扩展和重构
RAG 系统面临的挑战之一是用户的查询可能与知识库中使用的术语不匹配。如果使用一个好的模型来生成嵌入，这不会成为问题，因为查询的上下文很重要。但是，你永远无法知道特定查询是否会出现这种情况。
查询扩展和重构可以通过生成查询的多个版本来弥补这一缺陷。其前提是，即使同一查询存在多个变体，至少其中一种变体能够帮助 RAG 检索出最相关的文档。
要进行查询扩展，您需要一个能够生成输入变体的模型。BART 就是一个例子。让我们看看如何使用它进行查询扩展：
from transformers import BartForConditionalGeneration, BartTokenizer
# Load BART model and tokenizer
tokenizer = BartTokenizer.from_pretrained("facebook/bart-large")
model = BartForConditionalGeneration.from_pretrained("facebook/bart-large")
def reformulate_query(query, n=2):
inputs = tokenizer(query, return_tensors="pt")
outputs = model.generate(
**inputs,
max_length=64,
num_beams=10,
num_return_sequences=n,
temperature=1.5, # High temperature for diversity
top_k=50,
do_sample=True
)
# Decode the outputs one by one
reformulations = [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]
all_queries = [query] + reformulations
return all_queries
# Generate reformulations from an example query
query = "How do transformer-based systems process natural language?"
reformulated_queries = reformulate_query(query)
print(f"Original Query: {query}")
print("Reformulated Queries:")
for i, q in enumerate(reformulated_queries[1:], 1):
print(f"{i}. {q}")
在此代码中，您将加载一个预先训练好的 BART 模型和分词器。它被创建为一个BartForConditionalGeneration对象，这是一个用于文本生成的序列到序列模型。与使用 Hugging Face Transformers 库中的模型一样，您将输入分词，并在函数中传递给模型reformulate_query()。您要求模型n仅针对一个输入生成输出。
要创建更多变体，您可以将温度设置为略高于 1，甚至可以尝试更高的值。使用 BART 的生成实际上是要求模型读取您的输入并记住其在“隐藏状态”下的含义，然后将隐藏状态解码回文本，并包含可能的变体。多种变体是使用集束搜索创建的，您可以根据需要添加更多生成参数。
使用标记器将多个输出逐一解码为文本。然后，您可以在代码末尾将它们打印出来。运行此代码，您可能会看到：
Original Query: How do transformer-based systems process natural language?
Reformulated Queries:
1. How do transformer-based systems process natural language?
2. How do transformer-based systems work in natural language?
原始查询中的模糊性越大，您得到的变化就越大。
混合检索：密集和稀疏方法
RAG 的理念是利用知识库中最相关的文档来补充查询的上下文。这些附加信息可以帮助模型生成更好的响应。您可以使用不同的方法来查找相关文档。
密集向量检索是指将知识库中的文档表示为一个高维向量。该向量中的所有维度都很重要，并且没有明确指出每个维度代表什么的具体含义。通常，密集向量是一个看起来随机的浮点数向量。
然而，稀疏向量包含许多零。它通常维度更高，并且是一个整数向量。例如，独热向量 (one-hot)，其中每个位置代表词汇表中的一个单词，并且只有当该单词出现在文档中时，其值才为 1。
稠密向量和稀疏向量并没有孰优孰劣之分。如果使用嵌入模型生成稠密向量，它通常擅长捕捉语义相似性。而稀疏向量通常擅长捕捉关键词。对稀疏向量进行操作可能会消耗大量内存，但可以使用 Okapi BM25 等技术来减少工作量。
在下面的代码中，你需要安装计算 BM25 分数的库。你可以使用 pip 来完成此操作：
pip install rank-bm25
让我们看看如何结合稀疏向量和密集向量来构建检索系统：
from rank_bm25 import BM25Okapi
from transformers import AutoTokenizer, AutoModel
import faiss
import numpy as np
import torch
dense_tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
dense_model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
def generate_embedding(text):
"""Generate dense vector using mean pooling"""
inputs = dense_tokenizer(text, padding=True, truncation=True, return_tensors="pt", max_length=512)
with torch.no_grad():
outputs = dense_model(**inputs)
attention_mask = inputs['attention_mask']
embeddings = outputs.last_hidden_state
expanded_mask = attention_mask.unsqueeze(-1).expand(embeddings.shape).float()
sum_embeddings = torch.sum(embeddings * expanded_mask, axis=1)
sum_mask = torch.clamp(expanded_mask.sum(axis=1), min=1e-9)
mean_embeddings = sum_embeddings / sum_mask
return mean_embeddings.cpu().numpy()
# Sample document collection
documents = [
"Transformers use self-attention mechanisms to process input sequences in "
"parallel, making them efficient for long sequences.",
"The attention mechanism in transformers allows the model to focus on different "
"parts of the input sequence when generating each output element.",
"Transformer models have a fixed context length determined by the positional "
"encoding and self-attention mechanisms.",
"To handle sequences longer than the context length, transformers can use "
"techniques like sliding windows or hierarchical processing.",
"Recurrent Neural Networks (RNNs) process sequences sequentially, which can be "
"inefficient for long sequences.",
"Long Short-Term Memory (LSTM) networks are a type of RNN designed to handle "
"long-term dependencies in sequences.",
"The Transformer architecture was introduced in the paper 'Attention Is All "
"You Need' by Vaswani et al.",
"BERT (Bidirectional Encoder Representations from Transformers) is a "
"transformer-based model designed for understanding the context of words.",
"GPT (Generative Pre-trained Transformer) is a transformer-based model designed "
"for natural language generation.",
"Transformer-XL extends the context length of transformers by using a "
"segment-level recurrence mechanism."
]
# Prepare for sparse retrieval (BM25)
tokenized_corpus = [doc.lower().split() for doc in documents]
bm25 = BM25Okapi(tokenized_corpus)
# Prepare for dense retrieval (FAISS)
document_embeddings = generate_embedding(documents)
dimension = document_embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(document_embeddings)
def hybrid_retrieval(query, k=3, alpha=0.5):
"""Hybrid retrieval: Use both the BM25 and L2 index on FAISS"""
# Sparse score of each document with BM25
tokenized_query = query.lower().split()
bm25_scores = bm25.get_scores(tokenized_query)
# Normalize BM25 scores to [0,1] unless all elements are zero
if max(bm25_scores) > 0:
bm25_scores = bm25_scores / max(bm25_scores)
# Sort all documents according to L2 distance to query
query_embedding = generate_embedding(query)
distances, indices = index.search(query_embedding, len(documents))
# Dense score: 1/distance as similarity metric, then normalize to [0,1]
eps = 1e-5 # a small value to prevent division by zero
dense_scores = 1 / (eps + np.array(distances[0]))
dense_scores = dense_scores / max(dense_scores)
# Combine scores = affine combination of sparse and dense scores
combined_scores = alpha * dense_scores + (1 - alpha) * bm25_scores
# Get top-k documents
top_indices = np.argsort(combined_scores)[::-1][:k]
results = [(documents[idx], combined_scores[idx]) for idx in top_indices]
return results
# Retrieve documents using hybrid retrieval
query = "How do transformers handle long sequences?"
results = hybrid_retrieval(query)
print(f"Query: {query}")
for i, (doc, score) in enumerate(results):
print(f"Document {i+1} (Score: {score:.4f}):")
print(doc)
print()
运行此程序时，您将看到：
Query: How do transformers handle long sequences?
Document 1 (Score: 0.7924):
Transformers use self-attention mechanisms to process input sequences in parallel, making them efficient for long sequences.
Document 2 (Score: 0.7458):
Long Short-Term Memory (LSTM) networks are a type of RNN designed to handle long-term dependencies in sequences.
Document 3 (Score: 0.7131):
To handle sequences longer than the context length, transformers can use techniques like sliding windows or hierarchical processing.
首先，您需要为文档集的所有文档创建一个 Okapi BM25 索引。Okapi BM25 是一种基于 TF-IDF 的评分方法，这意味着它通过检查精确单词的交集来比较两段文本。因此，大小写并不重要。因此，您需要使用 BM25 将文档转换为小写。
然后，使用预先训练的句子 Transformer 模型为文档集合生成密集向量。将这些密集向量存储在 FAISS 索引中，以便使用 L2 距离进行高效的相似性搜索。
这段代码的关键部分在于函数hybrid_retrieval()。准备好 Okapi BM25 和 FAISS 索引后，您需要查找与查询字符串最匹配的文档。获得的 BM25 分数是与每个文档对应的 TF-IDF 分数。您还计算了每个文档与 FAISS 的 L2 距离度量。然后，该距离被转换为分数，以匹配 BM25 的分数：分数越高，匹配度越高。为了确保两种方法具有可比性，您需要将分数归一化到 [0, 1] 范围内。
根据您的选择，您可以通过更改参数 来更加注重密集检索或稀疏检索alpha。然后使用组合得分来查找要返回的前 k 个文档。正如您从上面的输出中看到的那样。
这种混合方法通常比单独使用任何一种方法效果更好，特别是对于语义理解和特定术语都很重要且复杂的查询。
带重新排序的多阶段检索
如果您有一个完美的模型来评估文档与查询的相关性，那么一个简单的检索系统就足够了。然而，没有哪个模型是完美的。事实上，通常质量更高的模型计算量也更大。这就是多阶段检索的用武之地。
混合检索擅长快速挑选文档。尤其是在使用快速模型的情况下，您可以轻松计算大量文档的得分。然而，这种挑选并不总是正确的。您可以使用速度较慢但更准确的模型重新计算得分。这一次，不会考虑所有文档，而只考虑混合检索挑选出的文档。只要第一阶段使用的模型大致正确，第二阶段计算更密集的模型就能为您提供准确的选择。
这就是多阶段检索技术的意义所在。让我们看看如何实现它：
# Load pre-trained model and tokenizer for re-ranking
reranker_tokenizer = AutoTokenizer.from_pretrained("cross-encoder/ms-marco-MiniLM-L-6-v2")
reranker_model = AutoModelForSequenceClassification.from_pretrained("cross-encoder/ms-marco-MiniLM-L-6-v2")
def rerank(query, documents, top_k=3):
"""Sort documents by the reranker model and select top-k"""
# Prepare inputs for the re-ranker
pairs = [[query, doc] for doc in documents]
features = reranker_tokenizer(pairs, padding=True, truncation=True, return_tensors="pt")
# Get re-ranking scores
with torch.no_grad():
scores = reranker_model(**features).logits.squeeze(-1).cpu().numpy()
# Sort documents by score, then pick top-k
ranked_indices = np.argsort(scores)[::-1][:top_k]
reranked_docs = [(documents[idx], float(scores[idx])) for idx in ranked_indices]
return reranked_docs
def multi_stage_retrieval(query, documents, initial_k=5, final_k=3):
"""Multi-stage retrieval: Hybrid retrievel to shortlist documents, then pick with a reranker"""
# Stage 1: Initial retrieval using hybrid method
initial_results = hybrid_retrieval(query, k=initial_k)
initial_docs = [doc for doc, _ in initial_results]
# Stage 2: Re-ranking
reranked_results = rerank(query, initial_docs, top_k=final_k)
return reranked_results
# Example query
query = "How do transformers handle long sequences?"
results = multi_stage_retrieval(query, documents)
print(f"Query: {query}")
print("Re-ranked Results:")
for i, (doc, score) in enumerate(results):
print(f"Document {i+1} (Score: {score:.4f}):")
print(doc)
print()
此代码建立在上一节的基础之上。它使用hybrid_retrieval()与之前相同的功能。
在函数中multi_stage_retrieval()，首先使用混合检索获取文档列表。然后使用重排序模型对这些文档进行重排序。
重排序模型是一种交叉编码器模型，它是一种可用于排序任务的Transformer模型。简单来说，它以两个序列作为输入，并以 的格式连接起来[CLS] query [SEP] document [SEP]。该模型的输出是一个分数，显示文档与查询的相关性。这是一个速度较慢的模型，但比L2距离或BM25更准确。
在函数 中rerank()，您将对查询和混合检索筛选出的文档运行重排序模型。然后，您将根据重排序模型提供的分数挑选出前 k 个文档。函数中的参数initial_k和允许您控制召回率（检索所有相关文档）和准确率（确保检索到的文档相关）之间的权衡。较大的和 会提高召回率，但需要更多的重排序计算，而较小的 和 则侧重于最相关的文档。final_kmulti_stage_retrieval()initial_kfinal_k
