使用 LlamaIndex 做一个智能体根据 PDF 回答问题

原创已于 2024-06-28 09:38:50 修改 · 1.9k 阅读

10 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能 #gpt #python

于 2024-06-26 17:26:25 首次发布

GPT 专栏收录该内容

4 篇文章

订阅专栏

什么是LlamaIndex

LlamaIndex 是一个为开发「上下文增强」的大语言模型应用的框架（也就是SDK）。上下文增强，泛指任何在私有或特定领域数据基础上应用大语言模型的情况。例如：

Question-Answering Chatbots (也就是 RAG)
Document Understanding and Extraction （文档理解与信息抽取）
Autonomous Agents that can perform research and take actions （智能体应用）

LlamaIndex 有 Python 和 Typescript 两个版本，Python 版的文档相对更完善。

Python 文档地址：LlamaIndex - LlamaIndex
Python API 接口文档：API Reference - LlamaIndex
TS 文档地址：What is LlamaIndex? | LlamaIndex.TS
TS API 接口文档：llamaindex | LlamaIndex.TS

安装

pip install llama-index

加载数据

通过 SimpleDirectoryReader 可以将本地的数据库加载，后期供大模型使用

支持的文件类型：

.csv - comma-separated values
.docx - Microsoft Word
.epub - EPUB ebook format
.hwp - Hangul Word Processor
.ipynb - Jupyter Notebook
.jpeg, .jpg - JPEG image
.mbox - MBOX email archive
.md - Markdown
.mp3, .mp4 - audio and video
.pdf - Portable Document Format
.png - Portable Network Graphics
.ppt, .pptm, .pptx - Microsoft PowerPoint

举个例子

from llama_index.core import SimpleDirectoryReader

reader = SimpleDirectoryReader(
        input_dir="./data", # 目标目录
        recursive=False, # 是否递归遍历子目录
        required_exts=[".pdf"] # (可选)只读取指定后缀的文件
    )
documents = reader.load_data()

为了方便展示，我们定义一个函数来展示对象的内部结构

import json
from pydantic.v1 import BaseModel

def show_json(data):
    """用于展示json数据"""
    if isinstance(data, str):
        obj = json.loads(data)
        print(json.dumps(obj, indent=4))
    elif isinstance(data, dict) or isinstance(data, list):
        print(json.dumps(data, indent=4))
    elif issubclass(type(data), BaseModel):
        print(json.dumps(data.dict(), indent=4, ensure_ascii=False))

我从网上下载了一些研报 pdf 做测试，现在打印一下 documents 是什么

show_json(documents[0])

输出：

{
    "id_": "8e087dc2-45c9-4d2b-a743-e0a20493fa2f",
    "embedding": null,
    "metadata": {
        "file_path": "D:\\Code\\Python\\Hello\\AI\\01\\data\\docstore.json",
        "file_name": "docstore.json",
        "file_type": "application/json",
        "file_size": 2,
        "creation_date": "2024-06-26",
        "last_modified_date": "2024-06-26"
    },
    "excluded_embed_metadata_keys": [
        "file_name",
        "file_type",
        "file_size",
        "creation_date",
        "last_modified_date",
        "last_accessed_date"
    ],
    "excluded_llm_metadata_keys": [
        "file_name",
        "file_type",
        "file_size",
        "creation_date",
        "last_modified_date",
        "last_accessed_date"
    ],
    "relationships": {},
    "text": "{}",
    "mimetype": "text/plain",
    "start_char_idx": null,
    "end_char_idx": null,
    "text_template": "{metadata_str}\n\n{content}",
    "metadata_template": "{key}: {value}",
    "metadata_seperator": "\n",
    "class_name": "Document"
}

正好是一个 text 为空的一段，这段对于大模型来说是没有有效信息的

正常的应该是这样

{
    "id_": "c392bb17-90ce-4c5f-89b6-abb2ae7c890a",
    "embedding": null,
    "metadata": {
        "page_label": "7",
        "file_name": "spyl.pdf",
        "file_path": "D:\\Code\\Python\\Hello\\AI\\01\\data\\spyl.pdf",
        "file_type": "application/pdf",
        "file_size": 743694,
        "creation_date": "2024-06-25",
        "last_modified_date": "2024-06-25"
    },
    "excluded_embed_metadata_keys": [
        "file_name",
        "file_type",
        "file_size",
        "creation_date",
        "last_modified_date",
        "last_accessed_date"
    ],
    "excluded_llm_metadata_keys": [
        "file_name",
        "file_type",
        "file_size",
        "creation_date",
        "last_modified_date",
        "last_accessed_date"
    ],
    "relationships": {},
    "text": " 行业简评报告 证券研究报告  \n \n \n请务必仔细阅读本报告最后部分的重要法律说明  \n5 据会稽山绍兴酒，今年 618，会稽山新品气泡黄酒 “一日一熏 ”在其抖音官方直播间  \n72小时卖出 1000万+，是抖音黄酒类第 2~10名总和的 100倍，且成交 80%以上都是年\n轻人。此外，根据数据统计， 6月12日-6月15日，气泡黄酒在抖音端销售额超过了啤\n酒热卖的前三销售额的总和，抢占部分啤酒市场。  \n资料来源： 酒业家  \n \n【5月规上白酒产量同比增长 4.9%】  \n国家统计局数据显示， 5月，中国规模以上企业白酒（折 65度，商品量）产量 36.1\n万千升，同比增长 4.9%；啤酒产量 353.5万千升，同比下降 4.5%；葡萄酒产量 0.8万千\n升， 同比下降 11.1%。1-5月， 中国规模以上企业累计白酒 （折 65度， 商品量） 产量 190.2\n万千升，同比增长 6.5%；啤酒产量 1504.5万千升，同比增长 0.7%；葡萄酒产量 4.9万\n千升，与上年同期持平。  \n资料来源： 酒业家  \n \n【伊利开卖无酒精饮料畅意 100%奶啤】  \n日前，伊利推出一款“酒精健康替代”新品 ——畅意 100%奶啤，产品定位为一款\n有情绪价值的无酒精饮料，适合消费者在商务宴请、好友聚餐、 KTV等场景下饮用，喝\n完还能开车。在酸甜好喝的同时更低负担，产品还添加有伊利自有菌株 K56，改善肠道\n环境，助力畅快嗨吃无负担。目前，畅意 100%奶啤将重点聚焦餐饮场景解辣解腻，还\n可以辐射到追剧、打游戏、下午茶等场景。 6月1日量产， 6月底铺货，将先在西南地\n区线下渠道发售，覆盖传统渠道、餐饮、便利店、校园场景等，定价 5元/瓶（ 300mL)。 \n资料来源： 食品板  \n \n4风险提示  \n宏观经济不及预期，行业竞争加剧，食品安全风险。",
    "mimetype": "text/plain",
    "start_char_idx": null,
    "end_char_idx": null,
    "text_template": "{metadata_str}\n\n{content}",
    "metadata_template": "{key}: {value}",
    "metadata_seperator": "\n",
    "class_name": "Document"
}

Process finished with exit code 0

我们可以对 SimpleDirectoryReader 读取到的数据进行进一步的处理

def gen_nodes():
    docs = SimpleDirectoryReader("./data", required_exts=[".pdf"], file_extractor={".pdf": PyMuPDFReader()}).load_data()
    node_parser = TokenTextSplitter(
        separator="\n\n",
        backup_separators=["\n", " 。", "！", "？", "；", "?", "!"],
        chunk_size=500,
        chunk_overlap=100
    )
    nodes = node_parser.get_nodes_from_documents(docs)
    return nodes

nodes = gen_nodes()

show_json(nodes[0])

在上面代码中我们通过 TokenTextSplitter 对文本进行切分，切分后每一个 chunk 的大小是 500 token, 上下chunk 可重复的部分不超过 100 token, 这样处理后，空的无效信息就会被过滤掉，长度不够的会合并一起，同时因为有交叠信息，避免搜索的时候漏掉一些关键信息。

输出：

{
    "id_": "6c5b4339-fdc6-447e-b46b-611c6ccd1859",
    "embedding": null,
    "metadata": {
        "file_path": "D:\\Code\\Python\\Hello\\AI\\01\\data\\spyl.pdf",
        "file_name": "spyl.pdf",
        "file_type": "application/pdf",
        "file_size": 743694,
        "creation_date": "2024-06-25",
        "last_modified_date": "2024-06-25",
        "total_pages": 8,
        "source": "1"
    },
    "excluded_embed_metadata_keys": [
        "file_name",
        "file_type",
        "file_size",
        "creation_date",
        "last_modified_date",
        "last_accessed_date"
    ],
    "excluded_llm_metadata_keys": [
        "file_name",
        "file_type",
        "file_size",
        "creation_date",
        "last_modified_date",
        "last_accessed_date"
    ],
    "relationships": {
        "1": {
            "node_id": "559905da-f246-4f84-bcde-5288d5059a13",
            "node_type": "4",
            "metadata": {
                "file_path": "D:\\Code\\Python\\Hello\\AI\\01\\data\\spyl.pdf",
                "file_name": "spyl.pdf",
                "file_type": "application/pdf",
                "file_size": 743694,
                "creation_date": "2024-06-25",
                "last_modified_date": "2024-06-25",
                "total_pages": 8,
                "source": "1"
            },
            "hash": "f22de76d5bdd9b79b9d34e322259bc405ca22b4f5f145934b6331206705c027f",
            "class_name": "RelatedNodeInfo"
        },
        "3": {
            "node_id": "ea84330b-a299-43c3-966b-c3224e2c4137",
            "node_type": "1",
            "metadata": {},
            "hash": "3028fef80d36c6d11e4a376eb0e09ca8afee25b24ac021f3bf246beb1d3b3b8c",
            "class_name": "RelatedNodeInfo"
        }
    },
    "text": "1 \n \n[Table_Rank] \n评级： 看好 \n \n[Table_Authors] \n \n赵瑞 \n分析师 \nSAC 执证编号：S0110522120001 \nzhaorui@sczq.com.cn \n \n \n \n[Table_Chart] \n市场指数走势（最近1 年） \n \n \n资料来源：聚源数据 \n相关研究 \n[Table_OtherReport] \n \n底部寻金，四维聚焦优选稳健和关注\n困境反转两类标的—— 2024 白酒行\n业年度策略报告 \n \n激荡三十年,白酒龙头企业兴衰沉浮的\n启示与展望 \n \n关注食品消费新品类，猪价保持弱势\n震荡 \n \n \n \n核心观点 \n[Table_Summary] \n \n⚫ \n酒水市场：据今日酒价数据，截至6 月23 日，24 年飞天原箱、散\n瓶批价分别为2420 元、2140 元，周环比分别下降180 元、150 元，\n飞天批价自6 月14 日阶段性低点企稳回升之后，再次出现快速下\n探。近期飞天批价波动主要由于在宏观经济仍处弱复苏、以及酒水",
    "mimetype": "text/plain",
    "start_char_idx": 2,
    "end_char_idx": 453,
    "text_template": "{metadata_str}\n\n{content}",
    "metadata_template": "{key}: {value}",
    "metadata_seperator": "\n",
    "class_name": "TextNode"
}

其他的 TextSplitter

SentenceSplitter：在切分指定长度的 chunk 同时尽量保证句子边界不被切断；
CodeSplitter：根据 AST（编译器的抽象句法树）切分代码，保证代码功能片段完整；
SemanticSplitterNodeParser：根据语义相关性对将文本切分为片段。

索引（Indexing）与检索（Retrieval）

基础概念：在「检索」相关的上下文中，「索引」即index，通常是指为了实现快速检索而设计的特定「数据结构」。

构建索引

可以直接在内存中构建一个索引

# 构建 index
index = VectorStoreIndex(nodes)

获取 retriever

# 获取 retriever
vector_retriever = index.as_retriever(
    similarity_top_k=2 # 返回前两个结果
)

检索

# 检索
results = vector_retriever.retrieve("白酒价格")

show_json(results[0]) # 有多个结果，这里展示第一个

输出

{
    "node": {
        "id_": "dc661086-ff9f-4560-8ac2-483b81223f80",
        "embedding": null,
        "metadata": {
            "file_path": "D:\\Code\\Python\\Hello\\AI\\01\\data\\spyl.pdf",
            "file_name": "spyl.pdf",
            "file_type": "application/pdf",
            "file_size": 743694,
            "creation_date": "2024-06-25",
            "last_modified_date": "2024-06-25",
            "total_pages": 8,
            "source": "3"
        },
        "excluded_embed_metadata_keys": [
            "file_name",
            "file_type",
            "file_size",
            "creation_date",
            "last_modified_date",
            "last_accessed_date"
        ],
        "excluded_llm_metadata_keys": [
            "file_name",
            "file_type",
            "file_size",
            "creation_date",
            "last_modified_date",
            "last_accessed_date"
        ],
        "relationships": {
            "1": {
                "node_id": "39abd1b8-3353-4632-895a-981cd7f609ad",
                "node_type": "4",
                "metadata": {
                    "file_path": "D:\\Code\\Python\\Hello\\AI\\01\\data\\spyl.pdf",
                    "file_name": "spyl.pdf",
                    "file_type": "application/pdf",
                    "file_size": 743694,
                    "creation_date": "2024-06-25",
                    "last_modified_date": "2024-06-25",
                    "total_pages": 8,
                    "source": "3"
                },
                "hash": "f81084b2e6112000a1faa64be0e87691c8203ac9815168ec92debfbcfcc38206",
                "class_name": "RelatedNodeInfo"
            },
            "2": {
                "node_id": "fa922781-5a5d-4df5-92a8-fb8389adc2d6",
                "node_type": "1",
                "metadata": {
                    "file_path": "D:\\Code\\Python\\Hello\\AI\\01\\data\\spyl.pdf",
                    "file_name": "spyl.pdf",
                    "file_type": "application/pdf",
                    "file_size": 743694,
                    "creation_date": "2024-06-25",
                    "last_modified_date": "2024-06-25",
                    "total_pages": 8,
                    "source": "3"
                },
                "hash": "ff2511e76d50da6f4387b69836bb5bff8f5686ea58dbdc691975d4f1cc3dc9ae",
                "class_name": "RelatedNodeInfo"
            },
            "3": {
                "node_id": "466cb2c0-3cfb-4e84-8c67-f8b052d007a7",
                "node_type": "1",
                "metadata": {},
                "hash": "43b2a8d2583cf7feaa14902ec3df94cc28c691bdf37bc67c5daf34da80e610d4",
                "class_name": "RelatedNodeInfo"
            }
        },
        "text": "2.1 白酒价格 \n截至6 月23 日，飞天原箱、飞天散瓶批价分别为2420 元、2140 元，周环比分别下\n降180 元、150 元；普五批价965 元、1573 批价880 元，周环比基本平稳。 \n \n图2 飞天茅台、普五、1573 批发价格走势 \n \n-5.01%\n-1.30%\n-7%\n-6%\n-5%\n-4%\n-3%\n-2%\n-1%\n0%\n1%\n2%\n3%\n电子(申万)\n通信(申万)\n建筑装饰(申万)\n石油石化(申万)\n汽车(申万)\n银行(申万)\n交通运输(申万)\n非银金融(申万)\n环保(申万)\n有色金属(申万)\n机械设备(申万)\n公用事业(申万)\n计算机(申万)\n基础化工(申万)\n建筑材料(申万)\n国防军工(申万)\n煤炭(申万)\n医药生物(申万)\n电力设备(申万)\n美容护理(申万)\n钢铁(申万)\n家用电器(申万)\n农林牧渔(申万)\n纺织服饰(申万)\n轻工制造(申万)\n社会服务(申万)\n综合(申万)",
        "mimetype": "text/plain",
        "start_char_idx": 330,
        "end_char_idx": 738,
        "text_template": "{metadata_str}\n\n{content}",
        "metadata_template": "{key}: {value}",
        "metadata_seperator": "\n",
        "class_name": "TextNode"
    },
    "score": 0.7363739872707532,
    "class_name": "NodeWithScore"
}

Process finished with exit code 0

排在第一个的往往是最匹配的搜索结果，当然向量检索，关键词检索各有各的优缺点，实际生产使用一般两者结合，会结合两者分别匹配出来的结果进行重排序然后输出。

使用自定义的 Vector Store，以 `Chroma` 为例

import chromadb
from chromadb.config import Settings

from llama_index.vector_stores.chroma import ChromaVectorStore

def gen_index():
    chroma_client = chromadb.PersistentClient(settings=Settings(allow_reset=True))
    try:
        chroma_collection = chroma_client.create_collection("demo")
        vector_store = ChromaVectorStore(chroma_collection)
        storage_context = StorageContext.from_defaults(vector_store=vector_store)
        nodes = gen_nodes()
        index = VectorStoreIndex(nodes, storage_context=storage_context)
        storage_context.persist("./data")
    except:
        chroma_collection = chroma_client.get_collection("demo")
        vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
        # index = VectorStoreIndex.from_vector_store(vector_store)
        storage_context = StorageContext.from_defaults(persist_dir="./data", vector_store=vector_store)
        index = load_index_from_storage(storage_context)
    return index

以上代码把通过 chromadb 创建了一个 demo 的 collection，并通过 ChromaVectorStore 实例化了一个向量存储，index 创建的时候除了传入 nodes,还需要传入 storage_context ,用于将内存中的数据实例化到数据库以及 ./data 目录中，下次运行就不用再去切割 pdf，而是通过 load_index_from_storage(storage_context) 从盘中加载到内存中，速度要快很多。

这个过程我们封装成函数 gen_index 这样更方便调用

通过 index 检索：

def query(question):
    index = gen_index()
    vector_retriever = index.as_retriever()
    res = vector_retriever.retrieve(question)
    for node in res:
        show_json(node)

上面代码封装了检索，但是这个结果是以向量搜索得到的排序，可能不是最优的，可以对搜索后的结果进一步排序处理，这里使用 SentenceTransformerRerank 进行重排序

def query(question):
    index = gen_index()
    vector_retriever = index.as_retriever()
    res = vector_retriever.retrieve(question)
    postprocessor = SentenceTransformerRerank(
        model="BAAI/bge-reranker-large", top_n=2
    )
    res = postprocessor.postprocess_nodes(res, query_str=question)
    for node in res:
        show_json(node)

Chat Engine

接下来只需要简单几行代码就可以实现问答机器人了

def chat():
    index = gen_index()
    chat_engine = index.as_chat_engine(llm=OpenAI(temperature=0, model="gpt-3.5-turbo"))
    while 1:
        query = input("请输入查询：")
        if query == "exit":
            break
        query = f"你是证券分析方面的专家，根据用户的问题回答，使用中文回答。用户问题：{query}"
        streaming_response = chat_engine.stream_chat(query) # 流式输出
        for token in streaming_response.response_gen:
            print(token, end="")


if __name__ == "__main__":
    chat()

以上输出不同的模型输出可能会有差异，但表达的信息基本是从 pdf 中出来的。

---------------------补充--------------------------

上面的 Chat Engine 的例子过于简单，忽略了重排序，这样的 chat engine 在一些情况下是回答不好问题的，虽然看上去能用。

接下来我们把重排序的逻辑加上再进行 chat

def chat():
    index = gen_index()

    # 定义检索后排序模型
    reranker = SentenceTransformerRerank(
        model="BAAI/bge-reranker-large", top_n=2
    )
    # 定义 RAG Fusion 检索器
    fusion_retriever = QueryFusionRetriever(
        [index.as_retriever()],
        similarity_top_k=5, # 检索召回 top k 结果
        num_queries=3, # 生成 query 数
        use_async=True
    )
    # 构建单轮 query engine
    query_engine = RetrieverQueryEngine(
        retriever=fusion_retriever,
        node_postprocessors=[reranker]
    )

....

上面代码修正了 query_engine 的构建方式，加入了重排序模型

接下来构建 chat engine

....
    text_qa_template = """
    你是证券分析方面的专家，结合给出的参考知识，使用中文回答。\n
    参考内容如下：\n
    ---------------------\n
    {context_str}\n
    ---------------------\n
    问题: {query_str}\n
    回答: 
    """
    text_qa_template = PromptTemplate(text_qa_template)
    # chat_engine = index.as_chat_engine(
    #     llm=OpenAI(temperature=0, model="gpt-3.5-turbo"),
    #     text_qa_template=text_qa_template,
    #     verbose=True
    # )
    chat_engine = CondenseQuestionChatEngine.from_defaults(
        query_engine=query_engine,
        condense_question_prompt=text_qa_template
    )
    # 这里为何查询的时候输入是英文？
    while 1:
        query = input("\n请输入查询：")
        if query == "exit":
            break
        streaming_response = chat_engine.stream_chat(query)
        for token in streaming_response.response_gen:
            print(token, end="")
            if token == '。':
                print('\n')

上面我们修改了 chat engine 的构建方式，不是简单粗暴直接 index.as_chat_engine 了

--------------------补充 end---------------------------

如果有疑问，欢迎私信探讨