llamaindex 自定义检索器的混合搜索

自定义检索器的混合搜索

概述

在本教程中，我们将展示如何定义一个非常简单的混合搜索版本！结合关键词查找检索和向量检索，使用“AND”和“OR”条件。

安装依赖

首先，我们需要安装 LlamaIndex：

!pip install llama-index

设置 OpenAI API 密钥：

import os

os.environ["OPENAI_API_KEY"] = "sk-..."

下载数据

下载 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'

加载数据

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

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)

定义向量索引和关键词表索引

在相同的文档存储上构建向量索引和关键词索引：

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):
    """Custom retriever that performs both semantic search and hybrid search."""

    def __init__(
        self,
        vector_retriever: VectorIndexRetriever,
        keyword_retriever: KeywordTableSimpleRetriever,
        mode: str = "AND",
    ) -> None:
        """Init params."""

        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]:
        """Retrieve nodes given query."""

        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)

相比之下，向量搜索会返回一个答案：

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

拓展内容

混合搜索的优势：

1. 多样性：通过结合关键词查找和语义搜索，可以提高检索结果的多样性和全面性。
2. 鲁棒性：不同的检索策略可能在不同的场景下表现更好，混合搜索可以提高系统的鲁棒性。
3. 精确性：通过“AND”和“OR”条件，可以更精确地控制检索结果的准确性。

实际应用场景：

1. 问答系统：在问答系统中，混合搜索可以提供更准确和全面的答案。
2. 信息检索：在信息检索系统中，混合搜索可以提高检索结果的质量。
3. 知识图谱：在知识图谱中，混合搜索可以提供更丰富的知识关联和推理。

通过这些详细的讲解和示例，学生们可以更好地理解和掌握自定义检索器的混合搜索的定义与实现方法，从而在实际项目中高效地应用。