释放双子座的力量：为 PDF 创建通用文档 Ai 管道 - 表格

表格、图像、图形或方程式不再是问题！提供完整代码

Photo by Matt Noble on Unsplash

自动化文档处理是ChatGPT革命中最大的受益者之一，因为大语言模型能够在零样本环境下处理广泛的主题和任务，这意味着无需领域内的标记训练数据。这使得构建人工智能驱动的应用以处理、解析和自动理解任意文档变得更加容易。尽管使用大语言模型的简单方法仍然受到非文本上下文（如图形、图像和表格）的限制，但这正是我们将在本博客文章中尝试解决的问题，特别关注PDF文件。

从基本层面来看，PDF文件只是字符、图像和线条的集合，以及它们的确切坐标。它们没有内在的“文本”结构，并不是为了被处理为文本而构建的，而只是为了按原样查看。这使得处理它们变得困难，因为仅文本的方法无法捕捉这些类型文档中的所有布局和视觉元素，导致上下文和信息的显著丧失。

绕过这种“仅文本”限制的一种方法是通过检测表格、图像和布局对文档进行重度预处理，然后再将其输入到大语言模型中。表格可以解析为Markdown或JSON，图像和图形可以通过其标题表示，文本可以按原样输入。然而，这种方法需要自定义模型，并仍会导致一些信息的丧失，那么我们能做得更好吗？

多模态大语言模型

最近的大型模型现在是多模态的，这意味着它们可以处理多种模态，例如文本、代码和图像。这为我们的问题提供了一种更简单的解决方案，其中一个模型可以同时完成所有任务。因此，我们可以直接将页面作为图像输入并按原样处理，而不是为图像添加说明和解析表格。

我们的管道将能够：

1. 加载PDF文件
2. 将每一页提取为图像
3. 使用大语言模型进行分块
4. 对每个分块进行索引

如果检索到一个分块，则完整页面将包含在大语言模型的上下文中以执行任务。接下来，我们将详细说明如何在实践中实现这一点。

管道

我们正在实施的管道是一个两步过程。

1. 首先，我们将每一页分割成重要的分块，并对每个分块进行摘要。
2. 其次，我们对分块进行一次索引，然后在每次收到请求时搜索这些分块，并在大语言模型上下文中包含每个检索到的分块的完整上下文。

第一步：页面分割和摘要

我们将页面提取为图像，并将每个图像传递给多模态大语言模型进行分割。像Gemini这样的模型可以轻松理解和处理页面布局：

• 表格被识别为一个分块。
• 图形形成另一个分块。
• 文本块被分割成单独的分块。
• …

对于每个元素，大语言模型生成一个摘要，可以嵌入并索引到向量数据库中。

第2步：嵌入和上下文检索

在本教程中，为了简化，我们将仅使用文本嵌入，但一个改进是直接使用视觉嵌入。

数据库中的每个条目包括：

• 分块的摘要。
• 找到该条目的页面编号。
• 链接到完整页面的图像表示，以提供额外的上下文。

该模式允许进行局部级别搜索（在分块级别），同时保持上下文跟踪（通过链接回完整页面）。例如，如果搜索查询检索到一个项目，代理可以包含整个页面图像，以向大语言模型提供完整布局和额外上下文，从而最大化响应质量。

通过提供完整图像，所有视觉提示和重要布局信息（如图像、标题、项目符号和相邻项目，如表格和段落）在生成响应时都可以提供给大语言模型。

Agents

我们将每个步骤实现为一个独立的、可重用的代理：

Parsing, Chunking, and Summarization Agent

第一个代理用于解析、分块和摘要。这涉及将文档分割成重要的块，然后为每个块生成摘要。此代理每个PDF只需运行一次以预处理文档。

Indexing, Search, and Retrieval Agent

第二个代理管理索引、搜索和检索。这包括将块的嵌入插入向量数据库以实现高效搜索。每个文档执行一次索引，而搜索可以根据不同的查询重复进行任意次数。

对于这两个代理，我们使用Gemini，一个具有强大视觉理解能力的多模态大语言模型。

解析与分块代理

第一个代理负责将每一页分割成有意义的块，并对每个块进行摘要，按照以下步骤进行：

步骤 1：将 PDF 页面提取为图像

我们使用 pdf2image 库。然后将图像编码为 Base64 格式，以简化将其添加到 LLM 请求中。

以下是实现：

from document_ai_agents.document_utils import extract_images_from_pdf
from document_ai_agents.image_utils import pil_image_to_base64_jpeg
from pathlib import Path

class DocumentParsingAgent:
    @classmethod
    def get_images(cls, state):
        """
        将 PDF 的页面提取为 Base64编码的JPEG图像。
        """
        assert Path(state.document_path).is_file(), "文件不存在"

        images = extract_images_from_pdf(state.document_path)
        assert images, "未提取到图像"

        pages_as_base64_jpeg_images = [pil_image_to_base64_jpeg(x) for x in images]
        return {"pages_as_base64_jpeg_images": pages_as_base64_jpeg_images}

extract_images_from_pdf：将 PDF 的每一页提取为 PIL 图像。

pil_image_to_base64_jpeg：将图像转换为 Base64编码的JPEG格式。

步骤 2：分块与摘要

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

from pydantic import BaseModel, Field
from typing import Literal
import json
import google.generativeai as genai
from langchain_core.documents import Document

class DetectedLayoutItem(BaseModel):
    """
    每个检测到的页面布局元素的模式。
    """
    element_type: Literal["Table", "Figure", "Image", "Text-block"] = Field(
        ..., 
        description="检测到的项目类型。示例：表格、图形、图像、文本块。"
    )
    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 页面中所有相关的布局元素，格式如下： "
            f"{LayoutElements.schema_json()}. "
            f"表格应至少有两列和至少两行。 "
            f"坐标应与每个布局项目重叠。",
            {"mime_type": "image/jpeg", "data": state.base64_jpeg},
        ]

        result = self.model.generate_content(messages)
        data = json.loads(result.text)

        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 模式定义了输出的结构，每种布局项目类型（表格、图形等）及其摘要。

步骤 3：页面的并行处理

为了提高速度，页面以并行方式处理。以下方法创建一个任务列表，以便同时处理所有页面图像，因为处理是 I/O 绑定的：

from langgraph.types import Send

class DocumentParsingAgent:
    @classmethod
    def continue_to_find_layout_items(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 构建的，将图像提取和布局检测步骤链接成一个统一的管道：

from langgraph.graph import StateGraph, START, END

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

        builder.add_node("get_images", self.get_images)
        builder.add_node("find_layout_items", self.find_layout_items)

        builder.add_edge(START, "get_images")
        builder.add_conditional_edges("get_images", self.continue_to_find_layout_items)
        builder.add_edge("find_layout_items", END)

        self.graph = builder.compile()

要在示例 PDF 上运行代理，我们执行：

if __name__ == "__main__":
    _state = DocumentLayoutParsingState(
        document_path="path/to/document.pdf"
    )
    agent = DocumentParsingAgent()

    result_images = agent.get_images(_state)
    _state.pages_as_base64_jpeg_images = result_images["pages_as_base64_jpeg_images"]

    result_layout = agent.find_layout_items(
        FindLayoutItemsInput(
            base64_jpeg=_state.pages_as_base64_jpeg_images[0],
            page_number=0,
            document_path=_state.document_path,
        )
    )

    for item in result_layout["documents"]:
        print(item.page_content)
        print(item.metadata["element_type"])

这将生成一个已解析、分块和摘要的 PDF 表示，这是我们接下来将构建的第二个代理的输入。

RAG Agent

这个第二个代理处理索引和检索部分。它将前一个代理的文档保存到向量数据库中，并使用结果进行检索。这可以分为两个独立的步骤：索引和检索。

步骤 1：索引分块文档

使用生成的摘要，我们将其向量化并保存到 ChromaDB 数据库中：

class DocumentRAGAgent:
    def index_documents(self, state: DocumentRAGState):
        """
        Index the parsed documents into the vector store.
        """
        assert state.documents, "Documents should have at least one element"

        if self.vector_store.get(where={"document_path": state.document_path})["ids"]:
            logger.info(
                "Documents for this file are already indexed, exiting this node"
            )
            return

        self.vector_store.add_documents(state.documents)
        logger.info(f"Indexed {len(state.documents)} documents for {state.document_path}")

index_documents 方法将分块摘要嵌入到向量存储中。我们保留元数据，例如文档路径和页码，以备后用。

步骤 2：处理问题

当用户提问时，代理会在向量存储中搜索最相关的分块。它检索摘要和相应的页面图像以进行上下文理解。

class DocumentRAGAgent:
    def answer_question(self, state: DocumentRAGState):
        """
        Retrieve relevant chunks and generate a response to the user's question.
        """
        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"Responding to question: {state.question}")

        messages = (
            [{"mime_type": "image/jpeg", "data": base64_jpeg} for base64_jpeg in images]
            + [doc.page_content for doc in relevant_documents]
            + [
                f"Answer this question using the context images and text elements only: {state.question}",
            ]
        )

        response = self.model.generate_content(messages)
        return {"response": response.text, "relevant_documents": relevant_documents}

检索器查询向量存储，以找到与用户问题最相关的分块。然后，我们为大语言模型（Gemini）构建上下文，将文本分块和图像结合，以生成响应。

完整的代理工作流程

代理工作流程分为两个阶段：索引阶段和问答阶段：

class DocumentRAGAgent:
    def build_agent(self):
        """
        Build the RAG agent workflow.
        """
        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,
    )

    state1 = DocumentLayoutParsingState(
        document_path=str(Path(__file__).parents[1] / "data" / "docs.pdf")
    )
    agent1 = DocumentParsingAgent()
    result1 = agent1.graph.invoke(state1)

    state2 = DocumentRAGState(
        question="Who was acknowledged in this paper?",
        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="What is the macro average when fine-tuning on PubLayNet using M-RCNN?",
        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）。

第一步：解析和摘要文档（代理 1）

• 执行时间：解析39页文档耗时29秒。
• 结果：代理1生成了一个索引文档，包含分块摘要和每页的Base64编码的JPEG图像。

第2步：询问文档（代理2）

我们提出以下问题：“解释LoRA，给出相关方程”

LoRA，或低秩适应，是一种用于机器学习的技术，特别是在微调大语言模型的背景下。它通过引入权重更新的低秩分解，允许高效地适应预训练模型，从而减少可训练参数的数量。

相关方程

1. 低秩分解： LoRA的核心思想是将权重更新 ( \Delta W ) 分解为两个低秩矩阵 ( A ) 和 ( B )： [ \Delta W = A \cdot B ] 其中 ( A ) 的大小为 ( m \times r )，( B ) 的大小为 ( r \times n )，并且 ( r ) 远小于 ( m ) 和 ( n )（原始权重矩阵的维度）。

2. 参数效率： LoRA引入的参数数量可以显著小于直接微调模型权重的数量。额外参数的总数为： [ \text{总参数} = (m \cdot r) + (r \cdot n) ]

3. 前向传播： 在前向传播过程中，输出 ( y ) 可以计算为： [ y = W \cdot x + A \cdot (B \cdot x) ] 其中 ( W ) 是原始权重矩阵，( x ) 是输入，( y ) 是输出。

4. 损失函数： 在训练过程中使用的损失函数可以定义为： [ \mathcal{L} = \mathcal{L}{original} + \lambda \cdot \mathcal{L}{LoRA} ] 其中 ( \mathcal{L}{original} ) 是原始任务的损失，( \mathcal{L}{LoRA} ) 是与低秩适应相关的损失，( \lambda ) 是控制两者损失之间权衡的超参数。

结果:

检索到的页面:

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

来自大语言模型的响应

作者提供的图片。

大语言模型能够利用视觉上下文，将方程和图形纳入其响应中，从而基于文档生成连贯且正确的响应。

结论

在这个快速教程中，我们看到如何通过利用最近的大语言模型的多模态特性，并使用每个文档中可用的完整视觉上下文，将您的自动化文档处理管道提升到一个新的水平，希望能够提高您从信息提取或RAG管道中获得的输出质量。

我们构建了一个更强大的文档分割步骤，能够检测重要项目，如段落、表格和图形，并对其进行摘要。然后，我们使用第一步的结果查询项目和页面的集合，以使用Gemini提供相关和精确的答案。作为下一步，您可以在您的用例和文档上尝试，尝试使用可扩展的向量数据库，并将这些代理作为您的人工智能驱动的应用的一部分进行部署。

完整代码和示例可在此处获得：https://github.com/CVxTz/document_ai_agents