重新排序器和两阶段检索的强大功能用于检索增强生成

当涉及自然语言处理（NLP）和信息检索时，高效、准确地检索相关信息的能力至关重要。随着该领域的不断发展，正在开发新技术和方法来提高检索系统的性能，特别是在 检索增强生成 （抹布）。其中一种技术，称为带有重新排序器的两阶段检索，已成为解决传统检索方法固有局限性的强大解决方案。

在本文中，我们讨论了两阶段检索和重新排序器的复杂性，探讨了它们的基本原理、实施策略以及它们在提高 RAG 系统的准确性和效率方面所带来的好处。我们还将提供实际示例和代码片段来说明概念并促进对这种尖端技术的更深入理解。

了解检索增强生成 (RAG)

在深入探讨两阶段检索和重新排序的细节之前，让我们简要回顾一下检索增强生成 (RAG) 的概念。 RAG 是一种通过为大型语言模型 (LLM) 提供对外部信息源（例如数据库或文档集合）的访问来扩展大型语言模型 (LLM) 的知识和功能的技术。更多内容请参阅文章“深入探讨法学硕士的检索增强生成“。

典型的 RAG 过程包括以下步骤：

1. 询问：用户向系统提出问题或提供指令。
2. 恢复：系统查询矢量数据库或文档集合以查找与用户查询相关的信息。
3. 提高：检索到的信息与用户的原始查询或指令相结合。
4. 信号生成：语言模型处理增强输入并生成响应，利用外部信息来增强其输出的准确性和全面性。

虽然 RAG 已被证明是一种强大的技术，但它也并非没有挑战。关键问题之一在于检索阶段，传统的检索方法可能无法识别最相关的文档，导致语言模型的响应不理想或不准确。

两阶段检索和重新排序的需要

传统的检索方法，例如基于关键字匹配或向量空间模型的检索方法，通常很难捕获查询和文档之间细微的语义关系。这种限制可能会导致检索到的文档仅具有表面相关性或错过可以显着提高生成响应质量的关键信息。

为了应对这一挑战，研究人员和从业者已转向使用重新排序器进行两阶段检索。这种方法涉及两个步骤：

1. 初始检索：在第一阶段，使用快速有效的检索方法（例如向量空间模型或基于关键字的搜索）检索相对较大的潜在相关文档集。
2. 重新排名：在第二阶段，采用更复杂的重新排序模型根据最初检索到的文档与查询的相关性对其进行重新排序，从而有效地将最相关的文档置于列表的顶部。

重新排序模型通常是神经网络或基于变压器的架构，经过专门训练以评估文档与给定查询的相关性。通过利用先进的自然语言理解功能，重新排序器可以捕获查询和文档之间的语义细微差别和上下文关系，从而实现更准确和相关的排名。

两阶段检索和重新排序的好处

在 RAG 系统的背景下，采用带有重新排序器的两阶段检索可带来几个显着的好处：

1. 提高精度：通过对最初检索到的文档进行重新排序并将最相关的文档提升到顶部，系统可以为语言模型提供更准确和精确的信息，从而生成更高质量的响应。
2. 缓解域外问题：用于传统检索的嵌入模型通常是在通用文本语料库上进行训练的，这可能无法充分捕获特定领域的语言和语义。另一方面，重新排序模型可以针对特定领域的数据进行训练，从而减轻“域外”问题并提高专业领域内检索到的文档的相关性。
3. 可扩展性：两阶段方法通过在初始阶段利用快速且轻量级的检索方法来实现高效扩展，同时为较小的文档子集保留计算量更大的重新排序过程。
4. 灵活性：重新排序模型可以独立于初始检索方法进行交换或更新，从而为系统不断变化的需求提供灵活性和适应性。

ColBERT：高效且有效的后期交互

重排序领域的杰出模型之一是 ColBERT（BERT 上的情境化后期交互）。 ColBERT 是一种文档重排序模型，它利用 BERT 的深度语言理解能力，同时引入一种称为“后期交互”的新颖交互机制。

ColBERT：通过 BERT 上的情境化后期交互实现高效且有效的段落搜索

ColBERT 中的后期交互机制可以通过单独处理查询和文档直到检索过程的最后阶段来实现高效、精确的检索。具体来说，ColBERT 使用 BERT 对查询和文档进行独立编码，然后采用轻量级但功能强大的交互步骤来建模它们的细粒度相似性。通过延迟但保留这种细粒度的交互，ColBERT 可以利用深度语言模型的表达能力，同时获得离线预先计算文档表示的能力，从而大大加快查询处理速度。

ColBERT 的后期交互架构具有多种优势，包括提高计算效率、文档集合大小的可扩展性以及对现实场景的实际适用性。此外，ColBERT 还通过去噪监督和残差压缩（在 ColBERTv2 中）等技术得到了进一步增强，这些技术改进了训练过程并减少了模型的空间占用，同时保持了高检索效率。

此代码片段演示了如何配置和使用 jina-colbert-v1-en 模型来索引文档集合，利用其有效处理长上下文的能力。

使用 Reranker 实施两阶段检索

现在我们已经了解了两阶段检索和重新排序背后的原理，让我们在 RAG 系统的背景下探索它们的实际实现。我们将利用流行的库和框架来演示这些技术的集成。

设置环境

在深入研究代码之前，让我们先设置开发环境。我们将使用 Python 和几个流行的 NLP 库，包括 Hugging Face Transformers、Sentence Transformers 和 LanceDB。

# Install required libraries
!pip install datasets huggingface_hub sentence_transformers lancedb

资料准备

出于演示目的，我们将使用 Hugging Face Datasets 中的“ai-arxiv-chunked”数据集，其中包含 400 多篇关于机器学习、自然语言处理和大型语言模型的 ArXiv 论文。

from datasets import load_dataset

dataset = load_dataset("jamescalam/ai-arxiv-chunked", split="train")

<pre>

接下来，我们将预处理数据并将其分割成更小的块，以方便高效的检索和处理。

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

def chunk_text(text, chunk_size=512, overlap=64):
tokens = tokenizer.encode(text, return_tensors="pt", truncation=True)
chunks = tokens.split(chunk_size - overlap)
texts = [tokenizer.decode(chunk) for chunk in chunks]
return texts

chunked_data = []
for doc in dataset:
text = doc["chunk"]
chunked_texts = chunk_text(text)
chunked_data.extend(chunked_texts)

在初始检索阶段，我们将使用 Sentence Transformer 模型将文档和查询编码为密集向量表示，然后使用 LanceDB 等向量数据库执行近似最近邻搜索。

from sentence_transformers import SentenceTransformer
from lancedb import lancedb

# Load Sentence Transformer model
model = SentenceTransformer('all-MiniLM-L6-v2')

# Create LanceDB vector store
db = lancedb.lancedb('/path/to/store')
db.create_collection('docs', vector_dimension=model.get_sentence_embedding_dimension())

# Index documents
for text in chunked_data:
vector = model.encode(text).tolist()
db.insert_document('docs', vector, text)

from sentence_transformers import SentenceTransformer
from lancedb import lancedb

# Load Sentence Transformer model
model = SentenceTransformer('all-MiniLM-L6-v2')

# Create LanceDB vector store
db = lancedb.lancedb('/path/to/store')
db.create_collection('docs', vector_dimension=model.get_sentence_embedding_dimension())

# Index documents
for text in chunked_data:
vector = model.encode(text).tolist()
db.insert_document('docs', vector, text)

通过对文档进行索引，我们可以通过查找给定查询向量的最近邻居来执行初始检索。

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

def chunk_text(text, chunk_size=512, overlap=64):
tokens = tokenizer.encode(text, return_tensors="pt", truncation=True)
chunks = tokens.split(chunk_size - overlap)
texts = [tokenizer.decode(chunk) for chunk in chunks]
return texts

chunked_data = []
for doc in dataset:
text = doc["chunk"]
chunked_texts = chunk_text(text)
chunked_data.extend(chunked_texts)

重新排名

初始检索后，我们将采用重新排序模型根据检索到的文档与查询的相关性对它们进行重新排序。在此示例中，我们将使用 ColBERT 重新排序器，这是一种快速、准确的基于转换器的模型，专为文档排序而设计。

from lancedb.rerankers import ColbertReranker

reranker = ColbertReranker()

# Rerank initial documents
reranked_docs = reranker.rerank(query, initial_docs)

这个 reranked_docs 列表现在包含根据 ColBERT 重新排序器确定的与查询的相关性重新排序的文档。

增强和生成

有了重新排序的相关文档，我们就可以进入 RAG 管道的增强和生成阶段。我们将使用 Hugging Face Transformers 库中的语言模型来生成最终响应。

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

tokenizer = AutoTokenizer.from_pretrained("t5-base")
model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")

# Augment query with reranked documents
augmented_query = query + " " + " ".join(reranked_docs[:3])

# Generate response from language model
input_ids = tokenizer.encode(augmented_query, return_tensors="pt")
output_ids = model.generate(input_ids, max_length=500)
response = tokenizer.decode(output_ids[0], skip_special_tokens=True)

print(response)