译者 | 朱先忠
审校 | 重楼
引言
近年来，自动化文档处理成为ChatGPT革命的最大赢家之一，因为LLM能够在零样本设置中处理广泛的主题和任务，这意味着无需域内标记的训练数据。这使得构建AI驱动的应用程序来处理、解析和自动理解任意文档变得更加容易。虽然使用LLM的简单方法仍然受到非文本上下文（例如图形、图像和表格）的阻碍，但是这正是我们将在本文中尝试解决的问题，而且我们特别关注PDF文件格式。
从根本上讲，PDF只是字符、图像和线条及其精确坐标的集合。它们没有固有的“文本”结构，也不是为作为文本处理而构建的，而只是按这些内容原样进行查看。这也正是使它们变得困难的原因，因为纯文本方法无法捕获这些类型的文档中的所有布局和视觉元素，从而导致上下文和信息的大量丢失。
绕过这种“纯文本”限制的一种方法是，在将文档输入LLM之前，通过检测表格、图像和布局对文档进行大量预处理。表格可以解析为Markdown或JSON格式，图像和图形可以用其标题表示，文本可以按原样输入。但是，这种方法需要自定义模型，并且仍会导致一些信息丢失。那么，我们能做得更好一些吗？
多模态LLM
现在，大多数最新的大型模型都是多模态的；这意味着，它们可以处理文本、代码和图像等多种模态形式的数据。这为我们的问题提供了一种更简单的解决方案，即一个模型可以同时完成所有工作。因此，我们不必为图像添加标题和解析表格，而是可以将页面作为图像输入并按原样处理。我们在本文中提出的管道方案将能够加载PDF，将每个页面提取为图像，将其拆分为块（使用LLM），并索引每个块。如果检索到块，则将整个页面包含在LLM上下文中以执行任务。
接下来，我们将详细介绍如何在实践中实现这一方案。
管道方案
概括来看，我们正在实施的管道是一个两步的过程。首先，我们将每个页面分割成重要的块并总结每个块。其次，我们对块进行一次索引，然后在每次收到请求时搜索这些块，并在LLM上下文中包含每个检索到的块的完整上下文信息。
第1步：页面分割和摘要
我们将页面提取为图像，并将它们中的每一个传递给多模态LLM进行分割。像Gemini这样的模型可以轻松理解和处理页面布局：
- 表格被识别为一个块。
- 图形形成另一个块。
- 文本块被分割成单独的块。
- …
对于每个元素，LLM都会生成一个摘要，可以将其嵌入并索引到向量数据库中。
第2步：嵌入和上下文检索
在本文中，我们将仅使用文本嵌入以简化操作，但一个改进是直接使用视觉嵌入。
数据库中的每个条目包括：
- 块的摘要。
- 找到它的页码。
- 指向完整页面图像表示的链接，用于添加上下文。
此架构允许在本地级别搜索（在块级别）的同时跟踪上下文（通过链接返回到完整页面）。例如，如果搜索查询检索到某个项目，则代理可以包含整个页面图像，以便向LLM提供完整布局和额外上下文，从而最大限度地提高响应质量。
通过提供完整图像，所有视觉提示和重要布局信息（如图像、标题、项目符号……）和相邻项目（表格、段落……）在生成响应时都可供LLM使用。
代理
我们将把每个步骤实现为单独的可重复使用的代理：
- 第一个代理用于解析、分块和摘要。这涉及将文档分割成重要的块，然后为每个块生成摘要。此代理只需对每个PDF运行一次即可对文档进行预处理。
- 第二个代理管理索引、搜索和检索。这包括将块的嵌入插入到向量数据库中以实现高效搜索。每个文档执行一次索引，而搜索可以根据不同查询的需要重复多次。
对于这两个代理，我们都使用谷歌开发的开源模型Gemini，这是一种具有强大视觉理解能力的多模态LLM。
解析和分块代理
第一个代理负责将每个页面分割成有意义的块并总结每个块，步骤如下：
第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(), "File does not exist"
# 从PDF中提取图像
images = extract_images_from_pdf(state.document_path)
assert images, "No images extracted"
# 转换图像到Base64编码的JPEG
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="Type of detected item. Examples: Table, Figure, Image, Text-block."
)
summary: str = Field(..., description="A detailed description of the layout item.")
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):
"""
Send a page image to the LLM for segmentation and summarization.
"""
messages = [
f"Find and summarize all the relevant layout elements in this PDF page in the following format: "
f"{LayoutElements.schema_json()}. "
f"Tables should have at least two columns and at least two rows. "
f"The coordinates should overlap with each layout item.",
{"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架构定义了输出的结构，包括每个布局项类型（表格、图形等）及其摘要。
第3步：页面的并行处理
为了提高速度，页面是并行处理的。由于处理是io绑定的，因此以下方法会创建一个任务列表来一次性处理所有页面图像：
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()
# 步骤1：从PDF中提取图像
result_images = agent.get_images(_state)
_state.pages_as_base64_jpeg_images = result_images["pages_as_base64_jpeg_images"]
#步骤2：处理第一页（作为一个示例）
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代理
第二个代理负责处理索引和检索部分。它将前一个代理的文档保存到向量数据库中，并使用其结果进行检索。这可以分为两个独立的步骤，即索引和检索。
第1步：索引拆分文档
使用生成的摘要，我们将其向量化并保存在ChromaDB数据库中：
class DocumentRAGAgent:
def index_documents(self, state: DocumentRAGState):
"""
将解析后的文档索引到向量存储区中。
"""
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):
"""
检索相关的数据块，并生成针对用户问题的响应。
"""
# 根据查询检索前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"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}",
]
)
#使用LLM生成响应
response = self.model.generate_content(messages)
return {"response": response.text, "relevant_documents": relevant_documents}在上述代码中，检索器查询向量存储以找到与用户问题最相关的块。然后，我们为LLM（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="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步：解析和摘要文档（代理1）
- 执行时间：解析39页的文档需要29秒。
- 结果：代理1生成一个索引文档，其中包含每个页面的块摘要和Base64编码的JPEG图像。
第2步：询问文档（代理2）
我们提出以下问题：“（Explain LoRA, give the relevant equations）解释LoRA，给出相关方程式”
结果：
检索到的页面如下：
来源：LLM&Adaptation.pdf（CC-BY许可）
LLM的回复
很明显，LLM能够利用视觉上下文根据文档生成连贯且正确的响应，从而将方程式和图形纳入其响应中。
结论
在本文中，我们了解了如何利用最新的LLM多模态性并使用每个文档中可用的完整视觉上下文信息将文档AI处理管道继续推进一步。我非常希望这一思想能够提高你从信息提取或RAG管道中获得的输出质量。
具体地说，我们构建了一个更强大的文档分割步骤，能够检测段落、表格和图形等重要项目并对其进行总结；然后，我们使用第一步的结果查询项目和页面的集合，以使用Gemini模型给出相关且准确的答案。接下来，你可以在自己的具体场景的文档上尝试这一方案，尝试使用可扩展的向量数据库，并将这些代理部署为AI应用程序的一部分。
最后，本文示例工程完整的代码可从链接https://github.com/CVxTz/document_ai_agents处获得。
译者介绍
朱先忠，51CTO社区编辑，51CTO专家博客、讲师，潍坊一所高校计算机教师，自由编程界老兵一枚。
原文标题：Build a Document AI Pipeline for Any Type of PDF with Gemini，作者：Youness Mansar
