基于谷歌Gemini多模態模型實現PDF文檔自動化處理 原創
本文將提出一種新的PDF文檔自動化處理方案，從而成功處理其中的表格、圖像、圖形或方程式等對象。
引言
近年來，自動化文檔處理成為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
