使用LlamaIndex实现简单的混合搜索引擎

最新推荐文章于 2024-09-28 20:13:38 发布

qq_37836323

最新推荐文章于 2024-09-28 20:13:38 发布

阅读量365

点赞数 3

文章标签：搜索引擎 python 人工智能

本文链接：https://blog.csdn.net/qq_29929123/article/details/140734670

版权

在本文中，我们将介绍如何使用LlamaIndex库构建一个简单的混合搜索引擎。这个引擎结合了关键词查找和向量检索的功能，通过设置"AND"和"OR"条件，可以实现更精确的搜索结果。以下是详细的步骤和示例代码。

设置环境

首先，我们需要安装LlamaIndex库。如果你在Colab上运行，可以使用以下命令安装：

!pip install llama-index

然后，设置环境变量：

import os

os.environ["OPENAI_API_KEY"] = "sk-..."  # 使用你的OpenAI API密钥

下载数据

接下来，我们下载Paul Graham的一篇文章作为示例数据：

!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

加载数据

将文档转换为节点，并插入到DocumentStore中：

from llama_index.core import SimpleDirectoryReader

# 加载文档
documents = SimpleDirectoryReader("./data/paul_graham").load_data()

from llama_index.core import Settings

nodes = Settings.node_parser.get_nodes_from_documents(documents)

from llama_index.core import StorageContext

# 初始化存储上下文（默认是内存中）
storage_context = StorageContext.from_defaults()
storage_context.docstore.add_documents(nodes)

构建索引

我们在相同的DocumentStore上构建向量索引和关键词索引：

from llama_index.core import SimpleKeywordTableIndex, VectorStoreIndex

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

定义自定义检索器

定义一个自定义检索器类，能够实现基本的混合搜索（关键词查找和语义搜索）：

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("无效的模式。")
        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[rid] for rid in retrieve_ids]
        return retrieve_nodes

插件检索器到查询引擎

将自定义检索器插件到查询引擎，并运行一些查询：

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 = CustomRetriever(vector_retriever, keyword_retriever)

# 定义响应合成器
response_synthesizer = get_response_synthesizer()

# 组装查询引擎
custom_query_engine = RetrieverQueryEngine(retriever=custom_retriever, response_synthesizer=response_synthesizer)
vector_query_engine = RetrieverQueryEngine(retriever=vector_retriever, response_synthesizer=response_synthesizer)
keyword_query_engine = RetrieverQueryEngine(retriever=keyword_retriever, response_synthesizer=response_synthesizer)

# 执行查询
response = custom_query_engine.query("What did the author do during his time at YC?")
print(response)  # 输出结果

response = custom_query_engine.query("What did the author do during his time at Yale?")
print(str(response))
len(response.source_nodes)