利用Gemini构建处理各种PDF文档的Document AI管道

照片由 Matt Noble 在 Unsplash 拍摄

自动化文档处理是ChatGPT革命中的最大赢家之一，因为大型语言模型（LLMs）能够应对各种主题和任务，而无需特定领域的标注训练资料，这意味着它们可以在不依赖特定领域训练资料的情况下工作。这使得构建可以处理、解析和自动理解任何文档的AI应用程序变得更加容易。尽管使用大型语言模型的简单方法仍然受制于非文本上下文，比如图表、图片和表格，但在本文中，我们将特别关注如何解决这些问题，特别是PDF文件。

从基本层面来说，PDF只是字符、图像和线条及其精确坐标的集合。它们没有内在的“文本”结构，也不是为了作为文本进行处理而构建的，而是仅用于查看。这就是为什么处理这类文档会很困难，因为仅基于文本的方法无法捕捉所有布局和视觉元素，从而使这些文档的完整性和信息量大打折扣，导致上下文和信息的丢失。

一种绕过这种“仅限文本”限制的方法是在将文档输入LLM之前，通过检测表格、图片和图形布局进行大量的预处理。表格可以解析为Markdown或JSON格式，图片和图形可以用其图注来表示，文本则可以原样喂入LLM。然而，这种方法需要定制的模型，并仍会有一些信息丢失，我们能否做得更好呢？

最近大多数大型模型已经变得多模态，这意味着它们可以处理文本、代码、图像等多种模态。这为我们提供了一个更简单的解决方案，即一个模型可以一次性完成所有任务。因此，我们可以直接将页面作为图像输入，而不需要描述图像或解析表格。我们的流程可以加载PDF，将每一页提取为图像，然后通过大语言模型将这些图像分割成块并进行索引。如果检索到某个块，则整个页面都会包含在大语言模型的任务上下文中。接下来，我们将详细说明如何在实际操作中实现这个方案。

我们正在实施的流程是一个两步过程。首先，我们将每一页分成重要的部分并对每个部分进行总结。其次，我们对这些部分进行一次索引，每次收到请求时搜索这些部分，并将每个检索到的部分的完整上下文包含在LLM的上下文中。

我们将每一页提取为图像，并将其传递给多模态LLM进行分割。像Gemini这样的模型能够轻松理解和处理页面的布局。

• 表格 被当作一个整体处理。
• 图形 形成另一个部分。
• 文本块 被拆分成单独的部分。
• …等等

对于每个元素，LLM 会生成一个摘要文本，这个摘要文本可以嵌入并索引到向量数据库中去。

在这篇教程里，我们仅使用文本嵌入，仅是为了简化，但更好的做法是直接使用图像嵌入。

数据库中的每一项包括：

• 摘录段落的简要总结；
• 所在页码；
• 页面完整图像的链接，提供更多上下文。

此模式允许进行局部级别的搜索（在块级别），同时保持追踪上下文（通过链接回整个页面）。例如，如果搜索查询检索到一个项目，代理程序可以包含整个页面的图像，以便为大语言模型提供完整的布局和更多的背景信息，从而提高响应质量。

通过提供完整的图片，所有的视觉线索和重要的布局信息（如图片、标题、列表点等）以及相邻的项目（如表格、段落等）在生成回复时都可供LLM参考。

我们将把每个步骤都作为独立且可重用的代理来实现。

第一个代理主要负责解析、分块和生成摘要。这涉及将文档分割成重要的部分，然后为每个部分生成摘要。只需对每个PDF文件运行一次该代理，以预处理文档。

第二个代理负责索引、搜索和检索操作。这包括将块的嵌入向量插入向量数据库，以便进行高效搜索。索引每份文档仅执行一次，而搜索则可以根据不同的查询需求重复进行。

我们为这两个代理程序使用Gemini，它是一个具备强大视觉理解能力的多模态AI模型。

第一个处理程序负责将每一页拆分成有意义的部分，并按照以下步骤总结每个部分。

步骤 1：将 PDF 页面转换成图像

我们使用 pdf2image 库这个工具。然后将图像编码为 Base64 编码，以便更容易地将它们添加到 LLM（大型语言模型）请求中。

这里就是实现了。

从 document_ai_agents.document_utils 导入 extract_images_from_pdf
从 document_ai_agents.image_utils 导入 pil_image_to_base64_jpeg
从 pathlib 导入 Path

类 DocumentParsingAgent：
    @classmethod
    def get_images(cls, state):
        """
        此函数从 PDF 文件中提取每一页并将其转换为 Base64 编码的 JPEG 图片。
        """
        assert Path(state.document_path).is_file(), "文件不存在或不是一个有效的文件路径。"
        # 从 PDF 中提取图片
        images = extract_images_from_pdf(state.document_path)
        assert images, "没有找到图片"
        # 将图片转换成 Base64 编码的 JPEG 格式
        pages_as_base64_jpeg_images = [pil_image_to_base64_jpeg(x) for x in images]
        # 返回包含所有页面的 Base64 编码 JPEG 图片的字典。
        return {"pages_as_base64_jpeg_images": pages_as_base64_jpeg_images}

从PDF中提取图片: 将PDF中的每页提取为PIL图像。

pil_image_to_base64_jpeg：将图像转码为Base64编码的JPEG格式的字符串。

第二步：分段和摘要

每个图像随后会被送到LLM进行分割和摘要。我们使用结构化结果以确保我们需要得到预期格式的预测。

    从 pydantic 导入 BaseModel, Field  
    从 typing 导入 Literal  
    导入 json  
    导入 google.generativeai 作为 genai  
    从 langchain_core.documents 导入 Document  

    class DetectedLayoutItem(BaseModel):  
        """  
        页面上每个检测到的布局元素的模式，用于描述布局元素的类型和摘要。  
        """  
        element_type: Literal["Table", "Figure", "Image", "Text-block"] = Field(  
            ...,   
            description="检测到的项目的类型。例如：Table, Figure, Image, Text-block."  
        )  
        summary: str = Field(..., description="布局元素的详细描述。")  

    class LayoutElements(BaseModel):  
        """  
        页面上布局元素的列表的模式。  
        """  
        layout_items: list[DetectedLayoutItem] = []  

    class FindLayoutItemsInput(BaseModel):  
        """  
        处理单页的输入模型。  
        """  
        document_path: str  
        base64_jpeg: str  
        page_number: int  

    class DocumentParsingAgent:  
        def __init__(self, model_name="gemini-1.5-flash-002"):  
            """  
            使用适当的模式初始化 LLM，以便正确处理布局元素。  
            """  
            layout_elements_schema = prepare_schema_for_gemini(LayoutElements)  
            self.model_name = model_name  
            self.model = genai.GenerativeModel(  
                self.model_name,  
                generation_config={  
                    "response_mime_type": "application/json",  
                    "response_schema": layout_elements_schema,  
                },  
            )  
        def find_layout_items(self, state: FindLayoutItemsInput):  
            """  
            将页面图像发送给 LLM 进行分割和描述和总结。  
            """  
            messages = [  
                f"请按照以下格式总结 PDF 页面上的所有相关布局元素：{LayoutElements.schema_json()}。"  
                f"表格应至少有两列和两行，并且表格中的内容应该能够反映页面的真实情况。"  
                f"坐标应与每个布局元素重叠，换句话说，每个布局元素应该被准确地定位和描述。",
                {"mime_type": "image/jpeg", "data": state.base64_jpeg},  
            ]  
            # 将提示发送给 LLM  
            result = self.model.generate_content(messages)  
            data = json.loads(result.text)  

            # 将 JSON 输出转换成文档形式  
            documents = [  
                Document(  
                    page_content=item["summary"],  
                    metadata={  
                        "page_number": state.page_number,  
                        "element_type": item["element_type"],  
                        "document_path": state.document_path,  
                    },  
                )  
                for item in data["layout_items"]  
            ]  
            return {"documents": documents}

LayoutElements 定义了输出的结构，包括各种布局项目类型（如表、图等）及其摘要。

第三步：页面的并行处理

页面并行处理以加快速度。以下方法会生成一个任务列表，以便一次处理所有页面图像，因为处理主要依赖于IO。

from langgraph.types import Send   

class 文档解析器:  
    @classmethod  
    def 继续处理布局项(cls, state):  
        """  
        生成任务以并行处理每一页。  
        """  
        return [  
            Send(  
                "find_layout_items",  
                FindLayoutItemsInput(  # 查找布局项输入对象
                    base64_jpeg=base64_jpeg,  
                    page_number=i,  
                    document_path=state.document_path,  
                ),  
            )  
            for i, base64_jpeg in enumerate(state['pages_as_base64_jpeg_images'])  
        ]

每个页面都单独发送到 find_layout_items 函数进行处理。

整个工作流程

代理程序的工作流程是使用一个 StateGraph 构建的，将图像提取和布局检测步骤整合成一个统一的流程。

    从langgraph.graph导入StateGraph, START和END  

    class DocumentParsingAgent:  
        def build_agent(self):  
            """  
            使用状态图来构建代理的工作流程。  
            """  
            builder = StateGraph(DocumentLayoutParsingState)  

            # 向构建器添加一个名为'get_images'的节点，该节点执行self.get_images方法
            builder.add_node("get_images", self.get_images)  
            # 向构建器添加一个名为'find_layout_items'的节点，该节点执行self.find_layout_items方法
            builder.add_node("find_layout_items", self.find_layout_items)  
            # 定义图的流程并连接节点
            builder.add_edge(START, "get_images")  # 从START节点连接到'get_images'节点
            builder.add_conditional_edges("get_images", self.continue_to_find_layout_items)  # 添加从'get_images'到self.continue_to_find_layout_items的条件边
            builder.add_edge("find_layout_items", END)  # 从'find_layout_items'节点连接到END节点

            # 定义图的流程和节点连接后，编译状态图。
            self.graph = builder.compile()

为了在样本PDF上运行代理，我们这样做。

    if __name__ == "__main__":
        _状态 = DocumentLayoutParsingState(
            document_path="path/to/document.pdf"
        )
        代理对象 = DocumentParsingAgent()

        # 步骤 1：从 PDF 中提取图像
        提取的图像 = 代理对象.get_images(_状态)
        _状态.pages_as_base64_images = 提取的图像["pages_as_base64_images"]

        # 步骤 2：处理第一张页面（作为示例）
        布局结果 = 代理对象.find_layout_items(
            FindLayoutItemsInput(
                base64_jpeg=_状态.pages_as_base64_images[0],
                page_number=0,
                document_path=_状态.document_path,
            )
        )
        # 显示结果如下
        for 项 in 布局结果["documents"]:
            print(项.page_content)
            print(项.metadata["element_type"])

这产生了PDF的解析、分段和总结的表示，这将成为我们接下来要构建的第二个代理程序的输入。

这个代理程序负责索引和检索的工作。它将前一个代理的文档保存到矢量数据库中，并用这些结果进行检索。这可以分成索引和检索两个步骤。

步骤 1：生成拆分文档的索引

我们使用生成的摘要，并将它们转化为向量，然后保存到ChromaDB数据库中。

    class DocumentRAGAgent:  
        def index_documents(self, state: DocumentRAGState):  
            """  
            将解析文档索引到向量存储中。  
            """  
            assert state.documents, "文档列表不应为空"  
            # 检查文档是否已索引  
            if self.vector_store.get(where={"document_path": state.document_path})["ids"]:  
                logger.info(  
                    "该文件的文档已索引，跳过此步骤"  
                )  
                return  # 如果已索引则跳过此步骤  
            # 将解析文档加入向量存储  
            self.vector_store.add_documents(state.documents)  
            logger.info(f"已为 {state.document_path} 索引了 {len(state.documents)} 个文档")

index_documents 方法将文档摘要存储到向量库中。我们会保留文档路径和页码等元数据，以便后续使用。

步骤2：处理问题和回答

当用户提问时，智能助手会在向量数据库中查找最相关的片段。它检索这些片段的概要及其相关页面的图片，以便更好地理解上下文。

    class DocumentRAGAgent:  
        def answer_question(self, state: DocumentRAGState):  
            """  
            根据用户的问题，从相关文档中检索信息并生成回答。  
            """  
            # 根据查询，检索与问题相关的前K个文档  
            relevant_documents: list[Document] = self.retriever.invoke(state.question)  

            # 检索相应的页面图像（避免重复的图像）  
            images = list(  
                set(  
                    [  
                        state.pages_as_base64_jpeg_images[doc.metadata["page_number"]]  
                        for doc in relevant_documents  
                    ]  
                )  
            )  
            logger.info(f"正在回答问题：{state.question}")  
            # 结合图像、相关文档内容和问题  
            messages = (  
                [{"mime_type": "image/jpeg", "data": base64_jpeg} for base64_jpeg in images]  
                + [doc.page_content for doc in relevant_documents]  
                + [  
                    f"只使用提供的图像和文本信息回答这个问题：{state.question}",  
                ]  
            )  
            # 用大模型生成回答  
            response = self.model.generate_content(messages)  
            return {"response": response.text, "relevant_documents": relevant_documents}

检索器（Retriever）查询向量存储库以找到与用户问题最相关的片段。然后我们为Gemini构建上下文，结合文本片段和图片生成响应。

代理工作流程的全部

代理的工作流程分为两个阶段，一个是索引阶段，另一个是提问与回答阶段。

    class DocumentRAGAgent:  
        def build_agent(self):  
            """  
            构建RAG代理的工作流程。  
            """  
            builder = StateGraph(DocumentRAGState)  
            # 添加用于索引和回答问题的节点  
            builder.add_node("index_documents", self.index_documents)  
            builder.add_node("answer_question", self.answer_question)  
            # 定义这个工作流  
            builder.add_edge(START, "index_documents")  
            builder.add_edge("index_documents", "answer_question")  
            builder.add_edge("answer_question", END)  
            self.graph = builder.compile()

示例

     if __name__ == "__main__":  
        from pathlib import Path  

      # 导入第一个代理来解析文档  
        from document_ai_agents.document_parsing_agent import (  
            DocumentLayoutParsingState,  
            DocumentParsingAgent,  
        )  
        # 步骤 1：使用第一个代理解析文档  
        state1 = DocumentLayoutParsingState(  
            document_path=str(Path(__file__).parents[1] / "data" / "docs.pdf")  
        )  
        agent1 = DocumentParsingAgent()  
        result1 = agent1.graph.invoke(state1)  
        # 步骤 2：设置第二个代理以进行检索和回答  
        state2 = DocumentRAGState(  
            question="本论文中致谢了谁？",  
            document_path=str(Path(__file__).parents[1] / "data" / "docs.pdf"),  
            pages_as_base64_jpeg_images=result1["pages_as_base64_jpeg_images"],  
            documents=result1["documents"],  
        )  
        agent2 = DocumentRAGAgent()  
        # 对文档进行索引  
        agent2.graph.invoke(state2)  
        # 回答第一个问题，  
        result2 = agent2.graph.invoke(state2)  
        print(result2["response"])  
        # 回答第二个问题，  
        state3 = DocumentRAGState(  
            question="在使用M-RCNN对PubLayNet进行微调时，宏平均值是多少？",  
            document_path=str(Path(__file__).parents[1] / "data" / "docs.pdf"),  
            pages_as_base64_jpeg_images=result1["pages_as_base64_jpeg_images"],  
            documents=result1["documents"],  
        )  
        result3 = agent2.graph.invoke(state3)  
        print(result3["response"])

有了这个实现，文档处理、检索和问答的流程就完成了。

示例：使用文档 AI 管道流程

让我们通过这个文档的实际例子来走一遍，说明使用文档 LLM & Adaptation.pdf，这是一套包含文字、公式和图表的 39 张幻灯片（CC BY 4.0）。

• 执行时间：解析39页的文档耗时29秒。
• 结果：一号代理生成了一个包含段落摘要和每页的base64编码JPEG格式图片的索引文件。

请解释LoRA，并写出相关的方程式。

获取的页面：

来源：LLM & Adaptation.pdf：许可证：CC-BY

这张图片由作者制作。

大型语言模型能够利用视觉信息，生成连贯且正确的基于文档的回复，并包含方程式和图表。

在这次快速教程里，我们展示了如何利用最近的大型语言模型的多模态特性来，将您的文档AI处理管道提升一步，从而有望提高您从信息提取或检索增强生成流程中得到的输出质量。

我们建立了一个更强大的文档分段步骤，能够检测并总结段落、表格和图表等重要项目，然后利用这一步骤的结果查询项目和页面集合，以使用Gemini提供更相关和精确的答案。下一步，您可以尝试在您的用例和文档上使用它，尝试使用可扩展向量数据库，并将这些代理部署到您的AI应用程序中。

在这里可以找到完整的代码和示例：https://github.com/CVxTz/document_ai_agents

谢谢你的阅读！ 😃