准备检索器#

检索器是 RAG 系统中最重要的组件之一，它负责从知识库中检索出和用户查询最相关的 k 篇文档。在 FlexRAG 中，有三种类型的检索器，分别是基于网络的检索器 WebRetriever 、基于 API 的检索器 APIBasedRetriever 以及 FlexRetriever。这三种检索器之间的关系如下图所示：

FlexRAG 中三种类型检索器的区别如下：

基于网络的检索器 WebRetriever ：直接从互联网上进行实时信息检索。该类检索器可以获取实时的动态信息，比如突发新闻、实时事件以及最新发布的信息，非常适合被应用于静态知识库无法满足的场景。更多有关此类检索器的信息，您可以访问高级教程：从互联网上获取相关信息以获取更多信息。
基于 API 的检索器 APIBasedRetriever：通过 APIs 与外部检索系统进行交互。这类检索器充当了专有数据库、企业系统或第三方服务的桥梁，实现与现有数据基础架构的无缝集成。在 FlexRAG 中，我们提供了两种基于 API 的检索器：ElasticRetriever 和TypesenseRetriever。
FlexRetriever：是一款先进的本地检索器，支持 多字段 和 多索引 检索功能：它允许将每个文档解析为多个语义字段（例如标题、摘要、内容），并为每个字段构建专用索引。此外，FlexRetriever 支持跨多个索引的混合搜索，从而可以根据复杂的信息需求定制灵活、细粒度的检索策略。此外，FlexRetriever 同时支持稀疏和密集检索方法，适用于各种检索任务。FlexRetriever 还与 Hugging Face 生态系统完全兼容，因此可以通过 Hugging Face Hub 轻松发布、共享和重用检索器。这种集成使用户能够以最低配置贡献和利用社区构建的检索流程。

在本教程中，我们将会向您展示如何从 HuggingFace Hub 中载入检索器或构建您自己的 FlexRetriever 。

通过 HuggingFace Hub 加载 `FlexRetriever`#

FlexRAG 提供了多个基于各种知识库构建的预定义 FlexRetriever 。这些检索器可在 HuggingFace Hub 上获取，并可轻松加载到您的应用中使用。您可以在 FlexRAG 代码库中找到可用检索器的列表。

您可以使用 FlexRetriever 类中的 load_from_hub 函数加载预定义的检索器。例如，要加载基于 enwiki_2021_atlas 数据集构建的检索器，您可以运行以下代码：

from flexrag.retriever import FlexRetriever

retriever = FlexRetriever.load_from_hub(repo_id='FlexRAG/wiki2021_atlas_bm25s')
passages = retriever.search('What is the capital of France?')

您还可以指定 top_k 参数来检索给定查询的前 k 个段落。例如，要检索前 5 个段落，您可以运行以下代码：

passages = retriever.search('What is the capital of France?', top_k=5)

备注

在快速入门：部署和评估您的 RAG 助手中，我们也提供了一些实例展示如何使用预先构建的检索器。

准备构建您自己的 `FlexRetriever`#

除了使用预定义的检索器之外，您还可以基于知识库准备自己的 FlexRetriever 。本节将指导您完成使用知识库准备 FlexRetriever 的过程。

下载知识库#

在构建检索器之前，您需要准备好您的知识库。在这个例子中，我们使用了 DPR 项目构建的维基百科知识库。您可以通过下面的命令下载该知识库：

# Download the corpus
wget https://dl.fbaipublicfiles.com/dpr/wikipedia_split/psgs_w100.tsv.gz
# Unzip the corpus
gzip -d psgs_w100.tsv.gz

备注

您也许希望使用您自己构建的知识库，FlexRAG 支持以 行分割文件 格式保存的知识库（如 csv、jsonl 或 tsv）。这些文件中，每一行代表一个知识片段，每个知识片段包含多个字段（如 text、id等）。您可以将您的知识库保存在一个或多个文件中。如希望了解更多有关信息，您可以查看构建知识库一节以了解如何构建知识库。在这个示例中，维基百科知识库共有三个字段： id、title 以及 text。其中 text 字段保存了维基百科页面的文本片段，title 字段保存了当前页面的标题，id 字段则作为每个知识片段的唯一标识符。您可以通过下面的命令来查看知识库的前五行：

在这个示例中，维基百科知识库共有三个字段： id、title 以及 text。其中 text 字段保存了维基百科页面的文本片段，title 字段保存了当前页面的标题，id 字段则作为每个知识片段的唯一标识符。您可以通过下面的命令来查看知识库的前五行：

head -n 5 psgs_w100.tsv

该命令的输出应该如下所示：

id      text    title
<id-1>  <text-1>        <title-1>
<id-2>  <text-2>        <title-2>
...

向检索器中添加文档#

在准备好知识库之后，你可以将文档添加到检索器中。在 FlexRAG 中，你可以使用 add_passages 函数将文档添加到检索器中。add_passages 函数接受一个 Context 的迭代器作为输入，其中每个 Context 表示一条知识内容。Context 类包含以下字段：

context_id ：该知识条片段的唯一标识符；
data ：一个包含该知识片段所有信息的字典，如 title 、 text 等；
source` （可选）：该知识片段的来源；
metadata （可选）：该知识片段的元信息；

你可以使用 RAGCorpusDataset 类来加载知识库，并将其转换为 Context 的迭代器。RAGCorpusDataset 类接受一个配置对象作为输入，该对象指定了知识库的文件路径、需要保存的字段以及唯一标识符字段。下面的代码片段演示了如何使用 Wikipedia 知识库来准备检索器：

from flexrag.datasets import RAGCorpusDataset, RAGCorpusDatasetConfig
from flexrag.retriever import FlexRetriever, FlexRetrieverConfig

RETRIEVER_PATH = "<path_to_retriever>"  # path to save the retriever
CORPUS_PATH = ["psgs_w100.tsv"]


def add_passages():
    corpus = RAGCorpusDataset(
        RAGCorpusDatasetConfig(
            file_paths=CORPUS_PATH,
            saving_fields=["title", "text"],
            id_field="id",
        )
    )
    retriever = FlexRetriever(
        FlexRetrieverConfig(
            log_interval=100000,
            batch_size=4096,
            retriever_path=RETRIEVER_PATH,
        )
    )
    retriever.add_passages(passages=corpus)
    return

add_passages()

我们还提供了一个命令行工具，以便于准备检索器。你可以运行以下命令，将文档添加到检索器中：

CORPUS_PATH='psgs_w100.tsv'
CORPUS_FIELDS='[title,text]'
RETRIEVER_PATH="<path_to_retriever>"

python -m flexrag.entrypoints.prepare_retriever \
    file_paths=[$CORPUS_PATH] \
    saving_fields=$CORPUS_FIELDS \
    id_field='id' \
    retriever_type=flex \
    flex_config.retriever_path=$RETRIEVER_PATH \
    reinit=True

为检索器构建索引#

在使用检索器之前，你需要为知识库构建索引。对于 FlexRetriever ，你可以使用 add_index 方法来构建索引。通过指定 index_name 、 index_config 和 indexed_fields_config 参数，你可以为知识库创建一个索引。

from flexrag.retriever import FlexRetriever
from flexrag.retriever.index import MultiFieldIndexConfig, RetrieverIndexConfig

RETRIEVER_PATH = "<path_to_retriever>"  # path to the retriever


def add_bm25_index():
    retriever = FlexRetriever.load_from_local(RETRIEVER_PATH)
    retriever.add_index(
        index_name="bm25",
        index_config=RetrieverIndexConfig(
            index_type="bm25",
        ),
        indexed_fields_config=MultiFieldIndexConfig(
            indexed_fields=["title", "text"],
            # concatenate the `title` and `text` fields for indexing
            merge_method="concat",
        ),
    )
    return


add_bm25_index()

在此示例中，我们为知识库中的标题（ title ）和正文（ text ）字段创建了一个 BM25 索引。merge_method 参数用于指定在建立索引时如何合并这些字段。在本例中，我们将标题和正文字段拼接成一个单一字段用于索引。

你也可以通过指定 index_type=faiss 来构建一个稠密索引（dense index）。稠密索引通过计算查询与被检索文档之间的语义相似度，来找到最相关的文档。其中查询和文档分别通过查询编码器（query encoder）和段落编码器（passage encoder）进行编码，从而获得各自对应的稠密向量。你可以运行以下代码，以 Wikipedia 作为知识库来构建一个稠密索引：

from flexrag.models import EncoderConfig, HFEncoderConfig
from flexrag.retriever import FlexRetriever
from flexrag.retriever.index import (
    MultiFieldIndexConfig,
    RetrieverIndexConfig,
    FaissIndexConfig,
)

RETRIEVER_PATH = "<path_to_retriever>"  # path to the retriever


def add_faiss_index():
    retriever = FlexRetriever.load_from_local(RETRIEVER_PATH)
    retriever.add_index(
        index_name="contriever",
        index_config=RetrieverIndexConfig(
            index_type="faiss",  # specify the index type
            faiss_config=FaissIndexConfig(
                # let FaissIndex determine the index configuration automatically
                # you can also specify a specific index type like "Flat", "IVF", etc.
                index_type="auto",
                index_train_num=-1,  # use all available data for training
                query_encoder_config=EncoderConfig(
                    encoder_type="hf",  # specify using Hugging Face model
                    hf_config=HFEncoderConfig(
                        # specify the Contriever model
                        # you can also choose other models
                        model_path="facebook/contriever-msmarco",  
                        # use the first GPU for query encoding
                        # if you do not want to use GPU, set device_id to []
                        device_id=[4],  
                    ),
                ),
                passage_encoder_config=EncoderConfig(
                    encoder_type="hf",
                    hf_config=HFEncoderConfig(
                        model_path="facebook/contriever-msmarco",
                        device_id=[0, 1, 2, 3],  # use four GPUs for data parallelism
                    ),
                ),
            ),
        ),
        indexed_fields_config=MultiFieldIndexConfig(
            indexed_fields=["title", "text"],
            merge_method="concat",  # concatenate the `title` and `text` fields for indexing
        ),
    )
    return


add_faiss_index()

在上述代码中，我们为知识库中的标题（ title ）和正文（ text ）字段创建了一个 Faiss 索引。index_type 参数指定要构建的索引类型，此处设置为 faiss 。faiss_config 参数则用于指定 Faiss 索引的相关配置，包括查询编码器（query encoder）和段落编码器（passage encoder）的配置

备注

在上面的脚本中，我们将 device_id 这一参数指定为 [0,1,2,3]，这将指定使用编号为 0,1,2,3 的四块 GPU 来进行编码任务，FlexRAG 将自动应用数据并行来加速编码过程。如果您仅有一块 GPU，您可以将编号修改为 [0] 来使用该 GPU，如果您想使用 CPU 进行编码，您可以将编码修改为 [] 来使用CPU。

FlexRAG 还提供了一个用于准备检索器的命令行工具。你可以运行以下命令来构建检索器：

RETRIEVER_PATH="<path_to_retriever>"  # path to the retriever

python -m flexrag.entrypoints.add_index \
    retriever_path=$RETRIEVER_PATH \
    index_name=bm25 \
    rebuild=False \
    indexed_fields=["title","text"] \
    merge_method="concat"

在您自己的程序中使用检索器#

在完成检索器的准备后，您就可以在您的 RAG 应用或任务中使用这个检索器了。举例来说，您可以使用 FlexRetriever 来为您的每个查询检索最相关的 5 个相关文档：

from flexrag.retriever import FlexRetriever


retriever = FlexRetriever.load_from_local("<path_to_retriever>")
passages = retriever.search('What is the capital of France?', top_k=5)[0]
print(passages)

将检索器部署为服务#

FlexRAG 提供了一个用于将检索器部署为服务的入口。这在以下场景中非常有用：当你希望使用该检索器来微调你自己的 RAG 助手，或者当你希望在生产环境中演示该检索器的功能时。你可以通过运行以下命令来部署检索器：

python -m flexrag.entrypoints.serve_retriever \
    host='0.0.0.0' \
    port='3402' \
    retriever_path=<path_to_retriever> \
    used_indexes=['bm25']

部署检索器之后，你可以通过访问 http://:/search 来使用检索服务，或者访问 http://:/docs 查看相关文档说明。你可以向 /search 端点发送一个包含查询内容和 top-k 参数的 POST 请求。下面是一个使用检索服务的示例：

import requests

def search_retriever(query, top_k=5):
    url = "http://<host>:<port>/search"
    payload = {
        "queries": [query],
        "top_k": top_k,
    }
    response = requests.post(url, json=payload)
    return response.json()

将您的检索器上传到 HuggingFace Hub#

您可以通过将检索器上传到 HuggingFace Hub 上来与社区分享您的检索器。举例来说，您可以使用下面的代码来上传您刚刚构建的 FlexRetriever ：

from flexrag.retrievers import FlexRetriever


retriever = FlexRetriever.load_from_local("<path_to_retriever>")
retriever.save_to_hub(repo_id="<your-repo-id>", token="<your-hf-token>")

在这段代码中，您需要指定 repo_id 和 token 这两个参数来将检索器上传到相应的仓库中。您可以在 HuggingFace 账户设置中获取您的token。在完成上传后，您就可以通过 HuggingFace Hub 的仓库链接来分享您的检索器了。

重要

为了确保您分享的 FlexRetriever 能够被其他人使用，您需要确保您的配置中的 查询编码器 和 文档编码器 均 已配置 且 可被公开访问 。在上面的例子中，我们将 查询编码器 和 文档编码器 均设置为 facebook/contriever-msmarco ，而该编码器是 HuggingFace Hub 上已有的编码器，因此可以被访问。如果您使用了独有的模型作为编码器，推荐您将编码器一并上传到 HuggingFace Hub 上。

准备检索器

目录

准备检索器#

通过 HuggingFace Hub 加载 FlexRetriever#

准备构建您自己的 FlexRetriever#

下载知识库#

向检索器中添加文档#

为检索器构建索引#

在您自己的程序中使用检索器#

将检索器部署为服务#

将您的检索器上传到 HuggingFace Hub#

通过 HuggingFace Hub 加载 `FlexRetriever`#

准备构建您自己的 `FlexRetriever`#