使用LlamaIndex进行递归检索和查询

最新推荐文章于 2024-09-28 17:54:53 发布

qq_37836323

最新推荐文章于 2024-09-28 17:54:53 发布

阅读量317

点赞数 3

文章标签： python 开发语言

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

版权

介绍

本文将介绍如何使用LlamaIndex库来实现递归检索和查询。我们将通过一个示例展示如何使用递归检索模块处理分层数据。递归检索的概念不仅是探索最相关的节点，还会探索节点关系以执行额外的检索/查询引擎。

示例介绍

在本示例中，我们将处理一个包含文本和嵌入式结构化表格的维基百科关于亿万富翁的文章（PDF格式）。我们首先在每个表格上创建一个Pandas查询引擎，然后通过IndexNode将每个表格表示为节点，并将其存储在向量存储中。

环境设置

我们需要安装相关的Python库，包括llama-index、camelot等：

%pip install llama-index-embeddings-openai
%pip install llama-index-readers-file pymupdf
%pip install llama-index-llms-openai

代码示例

加载文档和表格

首先，我们使用PyMuPDFReader读取文档的主要文本，并使用camelot提取文档中的一些结构化表格：

import camelot
from llama_index.core import VectorStoreIndex
from llama_index.core.query_engine import PandasQueryEngine
from llama_index.core.schema import IndexNode
from llama_index.llms.openai import OpenAI
from llama_index.readers.file import PyMuPDFReader
from typing import List

file_path = "billionaires_page.pdf"

# 初始化PDF读取器
reader = PyMuPDFReader()
docs = reader.load(file_path)

# 使用camelot解析表格
def get_tables(path: str, pages: List[int]):
    table_dfs = []
    for page in pages:
        table_list = camelot.read_pdf(path, pages=str(page))
        table_df = table_list[0].df
        table_df = (
            table_df.rename(columns=table_df.iloc[0])
            .drop(table_df.index[0])
            .reset_index(drop=True)
        )
        table_dfs.append(table_df)
    return table_dfs

table_dfs = get_tables(file_path, pages=[3, 25])

创建Pandas查询引擎

我们在每个结构化表格上创建一个Pandas查询引擎：

llm = OpenAI(model="gpt-4")

df_query_engines = [
    PandasQueryEngine(table_df, llm=llm) for table_df in table_dfs
]

response = df_query_engines[0].query(
    "What's the net worth of the second richest billionaire in 2023?"
)
print(str(response))  # 输出：$180 billion

response = df_query_engines[1].query(
    "How many billionaires were there in 2009?"
)
print(str(response))  # 输出：793

构建向量索引

我们在文档和附加的IndexNode对象上构建向量索引，并将这些节点链接到表格：

from llama_index.core import Settings

doc_nodes = Settings.node_parser.get_nodes_from_documents(docs)

summaries = [
    "This node provides information about the world's richest billionaires in 2023",
    "This node provides information on the number of billionaires and their combined net worth from 2000 to 2023."
]

df_nodes = [
    IndexNode(text=summary, index_id=f"pandas{idx}")
    for idx, summary in enumerate(summaries)
]

df_id_query_engine_mapping = {
    f"pandas{idx}": df_query_engine
    for idx, df_query_engine in enumerate(df_query_engines)
}

vector_index = VectorStoreIndex(doc_nodes + df_nodes)
vector_retriever = vector_index.as_retriever(similarity_top_k=1)

使用递归检索器

我们定义一个RecursiveRetriever对象来递归检索/查询节点：

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

recursive_retriever = RecursiveRetriever(
    "vector",
    retriever_dict={"vector": vector_retriever},
    query_engine_dict=df_id_query_engine_mapping,
    verbose=True,
)

response_synthesizer = get_response_synthesizer(response_mode="compact")

query_engine = RetrieverQueryEngine.from_args(
    recursive_retriever, response_synthesizer=response_synthesizer
)

response = query_engine.query(
    "What's the net worth of the second richest billionaire in 2023?"
)
print(str(response))  # 输出：$180 billion