El procesamiento automatizado de documentos es uno de los mayores ganadores de la revolución ChatGPT, ya que los LLM pueden abordar una amplia gama de temas y tareas en una configuración de cero disparos, es decir, sin datos de capacitación etiquetados en el dominio. Esto ha facilitado mucho la creación de aplicaciones impulsadas por IA para procesar, analizar y comprender automáticamente documentos arbitrarios. Aunque los enfoques ingenuos que utilizan LLM todavía se ven obstaculizados por contextos no textuales, como figuras, imágenes y tablas, esto es lo que intentaremos abordar en esta publicación de blog, con especial atención en los archivos PDF.
En un nivel básico, los archivos PDF son sólo una colección de caracteres, imágenes y líneas junto con sus coordenadas exactas. No tienen una estructura de “texto” inherente y no fueron creados para ser procesados como texto, sino únicamente para ser vistos tal como están. Esto es lo que dificulta trabajar con ellos, ya que los enfoques de solo texto no logran capturar todo el diseño y los elementos visuales en este tipo de documentos, lo que resulta en una pérdida significativa de contexto e información.
Una forma de evitar esta limitación de “solo texto” es realizar un preprocesamiento intensivo del documento detectando tablas, imágenes y diseños antes de enviarlos al LLM. Las tablas se pueden analizar en Markdown o JSON, las imágenes y figuras se pueden representar mediante sus títulos y el texto se puede alimentar tal cual. Sin embargo, este enfoque requiere modelos personalizados y aun así resultará en cierta pérdida de información, entonces, ¿podemos hacerlo mejor?
Los modelos grandes más recientes son ahora multimodales, lo que significa que pueden procesar múltiples modalidades como texto, código e imágenes. Esto abre el camino a una solución más simple a nuestro problema donde un modelo hace todo a la vez. Entonces, en lugar de subtitular imágenes y analizar tablas, podemos simplemente alimentar la página como una imagen y procesarla tal como está. Nuestra canalización podrá cargar el PDF, extraer cada página como una imagen, dividirla en fragmentos (usando el LLM) e indexar cada fragmento. Si se recupera un fragmento, la página completa se incluye en el contexto LLM para realizar la tarea. A continuación detallaremos cómo se puede implementar esto en la práctica.
El proceso que estamos implementando es un proceso de dos pasos. Primero, segmentamos cada página en partes importantes y resumimos cada una de ellas. En segundo lugar, indexamos fragmentos una vez y luego los buscamos cada vez que recibimos una solicitud e incluimos el contexto completo con cada fragmento recuperado en el contexto LLM.
Paso 1: segmentación y resumen de páginas
Extraemos las páginas como imágenes y pasamos cada una de ellas al LLM multimodal para segmentarlas. Modelos como Gemini pueden comprender y procesar el diseño de la página fácilmente:
- Mesas se identifican como un solo fragmento.
- Figuras formar otro trozo.
- Bloques de texto están segmentados en trozos individuales.
- …
Para cada elemento, el LLM genera un resumen que se puede incrustar e indexar en una base de datos vectorial.
Paso 2: Incrustación y recuperación contextual
En este tutorial usaremos la incrustación de texto solo por simplicidad, pero una mejora sería usar incrustaciones de visión directamente.
Cada entrada en la base de datos incluye:
- El resumen del trozo.
- El número de página donde se encontró.
- Un enlace a la representación de la imagen de la página completa para mayor contexto.
Este esquema permite búsquedas a nivel local (a nivel de fragmento) manteniendo seguimiento del contexto (enlazando a la página completa). Por ejemplo, si una consulta de búsqueda recupera un elemento, el agente puede incluir la imagen de la página completa para proporcionar un diseño completo y contexto adicional al LLM para maximizar la calidad de la respuesta.
Al proporcionar la imagen completa, todas las señales visuales y la información de diseño importante (como imágenes, títulos, viñetas…) y elementos adyacentes (tablas, párrafos,…) están disponibles para el LLM en el momento de generar una respuesta.
Implementaremos cada paso como un agente independiente y reutilizable:
El primer agente es para analizar, fragmentar y resumir. Esto implica la segmentación del documento en partes importantes, seguida de la generación de resúmenes para cada una de ellas. Este agente solo necesita ejecutarse una vez por PDF para preprocesar el documento.
El segundo agente gestiona la indexación, la búsqueda y la recuperación. Esto incluye la inserción de fragmentos en la base de datos de vectores para una búsqueda eficiente. La indexación se realiza una vez por documento, mientras que las búsquedas se pueden repetir tantas veces como sea necesario para diferentes consultas.
Para ambos agentes, utilizamos Géminisun LLM multimodal con sólidas capacidades de comprensión de la visión.
Agente de análisis y fragmentación
El primer agente se encarga de segmentar cada página en fragmentos significativos y resumir cada uno de ellos, siguiendo estos pasos:
Paso 1: extraer páginas PDF como imágenes
Usamos el pdf2image biblioteca. Luego, las imágenes se codifican en formato Base64 para simplificar su adición a la solicitud de LLM.
Aquí está la implementación:
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 Pathclass DocumentParsingAgent:
@classmethod
def get_images(cls, state):
"""
Extract pages of a PDF as Base64-encoded JPEG images.
"""
assert Path(state.document_path).is_file(), "File does not exist"
# Extract images from PDF
images = extract_images_from_pdf(state.document_path)
assert images, "No images extracted"
# Convert images to Base64-encoded 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: Extrae cada página del PDF como una imagen PIL.
pil_image_to_base64_jpeg: convierte la imagen a un formato JPEG codificado en Base64.
Paso 2: fragmentación y resumen
Luego, cada imagen se envía al LLM para su segmentación y resumen. Usamos resultados estructurados para asegurarnos de obtener las predicciones en el formato que esperamos:
from pydantic import BaseModel, Field
from typing import Literal
import json
import google.generativeai as genai
from langchain_core.documents import Documentclass DetectedLayoutItem(BaseModel):
"""
Schema for each detected layout element on a page.
"""
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):
"""
Schema for the list of layout elements on a page.
"""
layout_items: list[DetectedLayoutItem] = []
class FindLayoutItemsInput(BaseModel):
"""
Input schema for processing a single page.
"""
document_path: str
base64_jpeg: str
page_number: int
class DocumentParsingAgent:
def __init__(self, model_name="gemini-1.5-flash-002"):
"""
Initialize the LLM with the appropriate schema.
"""
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},
]
# Send the prompt to the LLM
result = self.model.generate_content(messages)
data = json.loads(result.text)
# Convert the JSON output into documents
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}
El LayoutElements El esquema define la estructura de la salida, con cada tipo de elemento de diseño (Tabla, Figura,…) y su resumen.
Paso 3: procesamiento paralelo de páginas
Las páginas se procesan en paralelo para mayor velocidad. El siguiente método crea una lista de tareas para manejar toda la imagen de la página a la vez, ya que el procesamiento está vinculado a io:
from langgraph.types import Sendclass DocumentParsingAgent:
@classmethod
def continue_to_find_layout_items(cls, state):
"""
Generate tasks to process each page in parallel.
"""
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)
]
Cada página se envía al find_layout_items funcionar como una tarea independiente.
Flujo de trabajo completo
El flujo de trabajo del agente se construye utilizando un StateGraphvinculando los pasos de extracción de imágenes y detección de diseño en una canalización unificada ->
from langgraph.graph import StateGraph, START, ENDclass DocumentParsingAgent:
def build_agent(self):
"""
Build the agent workflow using a state graph.
"""
builder = StateGraph(DocumentLayoutParsingState)
# Add nodes for image extraction and layout item detection
builder.add_node("get_images", self.get_images)
builder.add_node("find_layout_items", self.find_layout_items)
# Define the flow of the graph
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()
Para ejecutar el agente en un PDF de muestra hacemos:
if __name__ == "__main__":
_state = DocumentLayoutParsingState(
document_path="path/to/document.pdf"
)
agent = DocumentParsingAgent()# Step 1: Extract images from PDF
result_images = agent.get_images(_state)
_state.pages_as_base64_jpeg_images = result_images["pages_as_base64_jpeg_images"]
# Step 2: Process the first page (as an example)
result_layout = agent.find_layout_items(
FindLayoutItemsInput(
base64_jpeg=_state.pages_as_base64_jpeg_images[0],
page_number=0,
document_path=_state.document_path,
)
)
# Display the results
for item in result_layout["documents"]:
print(item.page_content)
print(item.metadata["element_type"])
Esto da como resultado una representación analizada, segmentada y resumida del PDF, que es la entrada del segundo agente que crearemos a continuación.
Agente RAG
Este segundo agente se encarga de la parte de indexación y recuperación. Guarda los documentos del agente anterior en una base de datos vectorial y utiliza el resultado para recuperarlos. Esto se puede dividir en dos pasos separados, indexación y recuperación.
Paso 1: indexar el documento dividido
Utilizando los resúmenes generados, los vectorizamos y los guardamos en una base de datos 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"
# Check if the document is already indexed
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 # Skip indexing if already done
# Add parsed documents to the vector store
self.vector_store.add_documents(state.documents)
logger.info(f"Indexed {len(state.documents)} documents for {state.document_path}")
El index_documents El método incorpora los resúmenes de fragmentos en el almacén de vectores. Guardamos metadatos como la ruta del documento y el número de página para su uso posterior.
Paso 2: Manejo de preguntas
Cuando un usuario hace una pregunta, el agente busca los fragmentos más relevantes en el almacén de vectores. Recupera los resúmenes y las imágenes de las páginas correspondientes para una comprensión contextual.
class DocumentRAGAgent:
def answer_question(self, state: DocumentRAGState):
"""
Retrieve relevant chunks and generate a response to the user's question.
"""
# Retrieve the top-k relevant documents based on the query
relevant_documents: list[Document] = self.retriever.invoke(state.question)# Retrieve corresponding page images (avoid duplicates)
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}")
# Construct the prompt: Combine images, relevant summaries, and the 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}",
]
)
# Generate the response using the LLM
response = self.model.generate_content(messages)
return {"response": response.text, "relevant_documents": relevant_documents}
El recuperador consulta el almacén de vectores para encontrar los fragmentos más relevantes para la pregunta del usuario. Luego construimos el contexto para el LLM (Gemini), que combina fragmentos de texto e imágenes para generar una respuesta.
El flujo de trabajo completo del agente
El flujo de trabajo del agente tiene dos etapas, una etapa de indexación y una etapa de respuesta a preguntas:
class DocumentRAGAgent:
def build_agent(self):
"""
Build the RAG agent workflow.
"""
builder = StateGraph(DocumentRAGState)
# Add nodes for indexing and answering questions
builder.add_node("index_documents", self.index_documents)
builder.add_node("answer_question", self.answer_question)
# Define the workflow
builder.add_edge(START, "index_documents")
builder.add_edge("index_documents", "answer_question")
builder.add_edge("answer_question", END)
self.graph = builder.compile()
Ejecución de ejemplo
if __name__ == "__main__":
from pathlib import Path# Import the first agent to parse the document
from document_ai_agents.document_parsing_agent import (
DocumentLayoutParsingState,
DocumentParsingAgent,
)
# Step 1: Parse the document using the first agent
state1 = DocumentLayoutParsingState(
document_path=str(Path(__file__).parents[1] / "data" / "docs.pdf")
)
agent1 = DocumentParsingAgent()
result1 = agent1.graph.invoke(state1)
# Step 2: Set up the second agent for retrieval and answering
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()
# Index the documents
agent2.graph.invoke(state2)
# Answer the first question
result2 = agent2.graph.invoke(state2)
print(result2["response"])
# Answer a second question
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"])
Con esta implementación, se completa el proceso para el procesamiento, la recuperación y la respuesta a preguntas de documentos.
Veamos un ejemplo práctico utilizando el documento. LLM y Adaptación.pdf un conjunto de 39 diapositivas que contienen texto, ecuaciones y figuras (CC BY 4.0).
Paso 1: Analizar y resumir el documento (Agente 1)
- Tiempo de ejecución: El análisis del documento de 39 páginas tomó 29 segundos.
- Resultado: El Agente 1 produce un documento indexado que consta de resúmenes de fragmentos e imágenes JPEG codificadas en base64 de cada página.
Paso 2: Cuestionar el documento (Agente 2)
Nos hacemos la siguiente pregunta:
“Explique LoRA, proporcione las ecuaciones relevantes.“
Resultado:
Páginas recuperadas: