LlamaIndex自定义检索器:简单混合搜索教程
在本教程中,我们将展示如何定义一个非常简单的混合搜索版本!结合关键词查找检索和向量检索,使用“AND”和“OR”条件。
设置
如果你在Colab上打开此笔记本,你可能需要安装LlamaIndex。
!pip install llama-index
设置OpenAI API密钥:
import os
os.environ["OPENAI_API_KEY"] = "sk-..."
下载数据
!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)
定义自定义检索器
我们现在定义一个自定义检索器类,可以实现基本的关键词查找和语义搜索的混合搜索。
设置“AND”意味着我们取两个检索集的交集,设置“OR”意味着我们取两个检索集的并集。
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))
print(len(response.source_nodes))
相比之下,向量搜索会返回一个答案:
response = vector_query_engine.query(
"What did the author do during his time at Yale?"
)
print(str(response))
print(len(response.source_nodes))
通过这些步骤,你可以在LlamaIndex中定义和使用自定义检索器,实现高效的混合搜索。这种混合搜索方法结合了关键词查找和语义搜索的优势,能够更好地满足复杂的查询需求。