使用LLamaIndex和Gemini构建高级搜索引擎

2024年06月04日 由 alex 发表 62 0

简介

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


什么是 LlamaIndex?

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


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


什么是 RAG?

在检索增强生成(简称 RAG)中,有两个主要组件: 检索器和生成器。


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


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


检索器的重要性


1


让我们来了解一下 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)


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


2


为了建立自定义检索器来执行混合搜索,我们需要创建两个索引。第一个矢量索引可以执行矢量搜索,第二个关键词索引可以执行关键词搜索。为了创建索引,我们需要存储上下文和节点文档,以及嵌入模型和 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?")


3


结论

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


文章来源:https://medium.com/@erkajalkumari/building-advanced-search-engines-with-llamaindex-and-gemini-9593105dc08c
欢迎关注ATYUN官方公众号
商务合作及内容投稿请联系邮箱:bd@atyun.com
评论 登录
热门职位
Maluuba
20000~40000/月
Cisco
25000~30000/月 深圳市
PilotAILabs
30000~60000/年 深圳市
写评论取消
回复取消