简介

检索器是 RAG（检索增强生成）管道中最重要的部分。在本文中，你将使用 LlamaIndex 实现一个结合关键字和矢量搜索的自定义检索器。使用 Gemini LLM 与多个文档聊天是我们构建 RAG 管道的项目用例。在项目开始时，我们将首先了解一些关键组件，如服务和存储上下文，以构建这样一个应用程序。

什么是 LlamaIndex？

大型语言模型领域正在迅速扩大，每天都有显著改进。随着越来越多的模型快速发布，人们越来越需要将这些模型与自定义数据进行整合。这种整合为商业、企业和最终用户提供了更大的灵活性和更深入的数据连接。

LlamaIndex 最初名为 GPT-index，是专为 LLM 应用程序设计的数据框架。随着像 ChatGPT 这样构建自定义数据驱动聊天机器人的流行程度不断提高，像 LlamaIndex 这样的框架变得越来越有价值。LlamaIndex 的核心是提供各种数据连接器，以促进数据摄取。在本文中，我们将探讨如何将我们的数据作为上下文传递给 LLM，这一概念就是我们所说的检索增强生成（RAG）。

什么是 RAG？

在检索增强生成（简称 RAG）中，有两个主要组件： 检索器和生成器。

• 检索器可以是矢量数据库，它的任务是检索与用户查询相关的文档，并将其作为上下文传递给提示器。
• 生成器模型是一个大型语言模型，它的工作是将检索到的文档与提示一起，根据上下文生成有意义的响应。

这样，RAG 就成了通过自动 “少量 ”提示进行语境学习的最佳解决方案。

检索器的重要性

让我们来了解一下 RAG 管道中寻回器组件的重要性。

要开发自定义检索器，关键是要确定最适合我们需求的检索器类型。为此，我们将实施一种混合搜索（Hybrid Search），将关键词搜索和矢量搜索整合在一起。

矢量搜索根据相似性或语义搜索来识别用户查询的相关文档，而关键词搜索则根据术语出现的频率来查找文档。在为混合搜索构建自定义检索器时，一个重要的决定是选择使用 AND 还是 OR 操作：

• AND 运算： 这种方法会检索包含所有指定术语的文档，限制性更强，但能确保高相关性。可以将其视为关键字搜索和矢量搜索结果的交集。
• OR 运算： 这种方法会检索包含任意指定术语的文档，从而增加结果的广度，但可能会降低相关性。可以将其视为关键词搜索和矢量搜索结果的联合。

使用 LLamaIndex 创建自定义检索器

现在让我们使用 LlamaIndex 建立客户检索器。为此，我们需要遵循一定的步骤。

第一步：安装

要开始在 Google Colab 或 Jupyter Notebook 上执行代码，需要安装所需的库，在我们的例子中，我们将使用 LlamaIndex 来构建自定义检索器，Gemini 用于嵌入模型和 LLM 推断，PyPDF 用于数据连接器。

!pip install llama-index
!pip install llama-index-multi-modal-llms-gemini
!pip install llama-index-embeddings-gemini

第二步：设置 Google API 密钥

在本项目中，我们将利用 Google Gemini 作为大语言模型来生成响应，并利用 LlamaIndex 作为嵌入模型来转换数据并将其存储在向量数据库或内存存储中。

from getpass import getpass

GOOGLE_API_KEY = getpass("Enter your Google API:")

第 三 步：加载数据并创建文档节点

在 LlamaIndex 中，数据加载是通过 SimpleDirectoryLoader 来完成的。首先，你需要创建一个文件夹，然后将任何格式的数据上传到该数据文件夹。在我们的示例中，我将上传一个 PDF 文件到数据文件夹。文件加载完成后，会被解析为节点，以便将文件分割成更小的片段。节点是在 LlamaIndex 框架内定义的数据模式。

最新版本的 LlamaIndex 更新了代码结构，现在包括节点解析器、嵌入模型和设置中的 LLM 的定义。

from llama_index.core import SimpleDirectoryReader
from llama_index.core import Settings

documents = SimpleDirectoryReader('data').load_data()
nodes = Settings.node_parser.get_nodes_from_documents(documents)

第四步：设置嵌入模型和大型语言模型

Gemini 支持多种模型，包括 gemini-pro、gemini-1.0-pro、gemini-1.5 和 vision model 等。在本例中，我们将使用默认模型并提供 Google API 密钥。对于 Gemini 中的嵌入模型，我们目前使用的是 embedding-001。确保添加了有效的 API 密钥。

from llama_index.embeddings.gemini import GeminiEmbedding
from llama_index.llms.gemini import Gemini

Settings.embed_model = GeminiEmbedding(
    model_name="models/embedding-001", api_key=GOOGLE_API_KEY
)
Settings.llm = Gemini(api_key=GOOGLE_API_KEY)

第五步：定义存储上下文并存储数据

一旦数据被解析为节点，LlamaIndex 就会提供一个存储上下文，它为存储数据的向量嵌入提供默认的文档存储。该存储上下文将数据保存在内存中，以便日后编制索引。

from llama_index.core import StorageContext

storage_context = StorageContext.from_defaults()
storage_context.docstore.add_documents(nodes)

创建索引--关键字和索引

为了建立自定义检索器来执行混合搜索，我们需要创建两个索引。第一个矢量索引可以执行矢量搜索，第二个关键词索引可以执行关键词搜索。为了创建索引，我们需要存储上下文和节点文档，以及嵌入模型和 LLM 的默认设置。

from llama_index.core import SimpleKeywordTableIndex, VectorStoreIndex

vector_index = VectorStoreIndex(nodes, storage_context=storage_context)
keyword_index = SimpleKeywordTableIndex(nodes, storage_context=storage_context)

第六步：构建自定义检索器

要使用 LlamaIndex 为混合搜索构建自定义检索器，我们首先需要定义模式，特别是通过适当配置节点。对于检索器来说，矢量索引检索器和关键字检索器都是必需的。这样我们就可以执行混合搜索，将两种技术整合在一起，最大限度地减少幻觉。此外，我们还必须指定模式（AND 或 OR），这取决于我们想如何组合搜索结果。

节点配置完成后，我们将使用向量和关键词检索器查询每个节点 ID 的捆绑包。根据所选模式，我们定义并最终确定自定义检索器。

from llama_index.core import QueryBundle
from llama_index.core.schema import NodeWithScore

from llama_index.core.retrievers import (
    BaseRetriever,
    VectorIndexRetriever,
    KeywordTableSimpleRetriever,
)
from typing import List
class CustomRetriever(BaseRetriever):
    def __init__(
        self,
        vector_retriever: VectorIndexRetriever,
        keyword_retriever: KeywordTableSimpleRetriever,
        mode: str = "AND") -> None:
       
        self._vector_retriever = vector_retriever
        self._keyword_retriever = keyword_retriever
        if mode not in ("AND", "OR"):
            raise ValueError("Invalid mode.")
        self._mode = mode
        super().__init__()
    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        vector_nodes = self._vector_retriever.retrieve(query_bundle)
        keyword_nodes = self._keyword_retriever.retrieve(query_bundle)
        vector_ids = {n.node.node_id for n in vector_nodes}
        keyword_ids = {n.node.node_id for n in keyword_nodes}
        combined_dict = {n.node.node_id: n for n in vector_nodes}
        combined_dict.update({n.node.node_id: n for n in keyword_nodes})
        if self._mode == "AND":
            retrieve_ids = vector_ids.intersection(keyword_ids)
        else:
            retrieve_ids = vector_ids.union(keyword_ids)
        retrieve_nodes = [combined_dict[r_id] for r_id in retrieve_ids]
        return retrieve_nodes

第七步：定义检索器

既然已经定义了自定义检索器类，我们就需要实例化检索器并合成查询引擎。响应合成器用于根据用户查询和给定的文本块集从 LLM 生成响应。响应合成器的输出是一个响应对象，它将自定义检索器作为参数之一。

from llama_index.core import get_response_synthesizer
from llama_index.core.query_engine import RetrieverQueryEngine

vector_retriever = VectorIndexRetriever(index=vector_index, similarity_top_k=2)
keyword_retriever = KeywordTableSimpleRetriever(index=keyword_index)
# custom retriever => combine vector and keyword retriever
custom_retriever = CustomRetriever(vector_retriever, keyword_retriever)
# define response synthesizer
response_synthesizer = get_response_synthesizer()
custom_query_engine = RetrieverQueryEngine(
    retriever=custom_retriever,
    response_synthesizer=response_synthesizer,
)

第八步：运行自定义检索器查询引擎

最后，我们开发了自定义检索引擎，它能显著减少幻觉。为了测试其有效性，我们运行了用户查询，包括一个来自上下文内的提示和另一个来自上下文外的提示，然后对生成的响应进行评估。

query = "what does the data context contain?"
print(custom_query_engine.query(query))
print(custom_query_engine.query("what is science?")

结论

我们成功实现了一种自定义检索器，它通过使用 LlamaIndex 将矢量检索器和关键词检索器结合起来，在 Gemini LLM 和嵌入的支持下执行混合搜索。这种方法在一定程度上有效减少了典型 RAG 管道中的 LLM 幻觉。