Layout-Aware Document Ingestion Adds Image Context to Spark NLP Retrieval Pipelines

PDFs, slide decks, and reports place important information in charts, diagrams, and other images. A text extraction pipeline can miss that information, while sending every page to a vision-language model (VLM) applies image inference even where extracted text is sufficient. Spark NLP’s LayoutAlignerForVision and LayoutAlignerForText let you pair extracted text with nearby images, generate image captions with optional text context, and insert those captions into reconstructed document text before chunking and embedding.

Why document layout matters for retrieval

A chart’s title, surrounding paragraph, and visual trend can each contribute to its meaning. A caption generated from the image alone may miss context that appears elsewhere on the page. If the caption is indexed separately from the document text, retrieval may also lose the relationship between the visual and the discussion around it.

The LayoutAligners workflow connects those steps. ReaderAssembler produces text and image annotation columns. LayoutAlignerForVision matches text annotations with nearby images using layout metadata and produces aligned document, image, and prompt columns. A VLM generates captions for those aligned images. LayoutAlignerForText then places the captions into the text flow, which you can pass to familiar chunking and embedding stages.

This design applies VLM inference to aligned images instead of asking the model to caption every page. It does not, by itself, establish a retrieval-quality or cost improvement; those outcomes depend on your documents, captioning model, and downstream evaluation.

The ingestion workflow

The example below follows the Spark NLP LayoutAligners document-understanding notebook. It uses PDF input, the AutoGGUFVisionModel vision model, character-based chunking, and sentence embeddings. The notebook demonstrates the pipeline with three example PDFs.

Read text and image annotations

ReaderAssembler reads files from contentPath and emits separate annotation columns. With outputCol="data", the text and image columns are named data_text and data_image.

from pyspark.ml import Pipeline

import sparknlp

from sparknlp.annotator import AutoGGUFVisionModel

from sparknlp.reader import LayoutAlignerForText, LayoutAlignerForVision, ReaderAssembler



spark = sparknlp.start()

pdf_directory = "./pdf-files"

empty_df = spark.createDataFrame([], "string").toDF("text")



reader = (

    ReaderAssembler()

    .setContentType("application/pdf")

    .setContentPath(pdf_directory)

    .setOutputAsDocument(False)

    .setOutputCol("data")

    .setUseEncodedImageBytes(True)

)



reader_pipeline = Pipeline(stages=[reader])

reader_df = reader_pipeline.fit(empty_df).transform(empty_df)

The resulting data_text annotations contain extracted document text. data_image contains image annotations for the source files.

Align text and images

The aligner uses layout metadata to associate an image with nearby text. Neighbor text is optional. When enabled, neighborTextCharsWindow controls the number of characters taken from before and after the aligned text when building the prompt context.

aligner_vision = (

    LayoutAlignerForVision()

    .setInputCols(["data_text", "data_image"])

    .setOutputCol("aligned")

    .setImageCaptionBasePrompt(

        "Describe the image with concise business and technical details."

    )

    .setAddNeighborText(True)

    .setNeighborTextCharsWindow(500)

)



vision_pipeline = Pipeline(stages=[aligner_vision])

aligned_df = vision_pipeline.fit(reader_df).transform(reader_df)

With the output column set to aligned_doc, aligned_prompt. The prompt includes the configured base prompt and, when neighbor text is enabled, text context around the aligned paragraph. A window value of .setAddNeighborText(True) and .setNeighborTextCharsWindow(500). The aligner still produces paired document and image annotations with its default prompt.

If you do not need neighboring text, omit .setAddNeighborText(True) and .setNeighborTextCharsWindow(500). The aligner still produces paired document and image annotations with its default prompt.

Caption aligned images

The notebook passes the aligned prompt and image columns to AutoGGUFVisionModel:

vision_model = (

    AutoGGUFVisionModel.pretrained()

    .setInputCols(["aligned_prompt", "aligned_image"])

    .setOutputCol("image_caption")

    .setBatchSize(1)

    .setNGpuLayers(99)

    .setNCtx(4096)

    .setMinKeep(0)

    .setMinP(0.05)

    .setNPredict(40)

    .setPenalizeNl(True)

    .setRepeatPenalty(1.18)

    .setTemperature(0.05)

    .setTopK(40)

    .setTopP(0.95)

)



caption_pipeline = Pipeline(stages=[vision_model])

caption_df = caption_pipeline.fit(aligned_df).transform(aligned_df)

The notebook output shows AutoGGUFVisionModel loading the Qwen2.5_VL_3B_Instruct_Q4_K_M_gguf model. Model availability, hardware, and inference settings affect execution; the example’s GPU-layer setting should be adjusted for the environment where you run it.

Rebuild text in document order
LayoutAlignerForText combines the aligned document annotations with the generated captions. Pass aligned_doc as the document input so the captions stay associated with the matched text elements.

aligner_text = (

    LayoutAlignerForText()

    .setInputCols(["aligned_doc", "image_caption"])

    .setOutputCol("aligned_text")

)



text_pipeline = Pipeline(stages=[aligner_text])

result_df = text_pipeline.fit(caption_df).transform(caption_df)

By default, LayoutAlignerForText emits a merged document annotation per file and drops the two input annotation columns. You can set preserveColumns=True when you need those inputs for inspection.

Chunk and embed the reconstructed text
The notebook uses DocumentCharacterTextSplitter and BertSentenceEmbeddings to prepare text for a vector-indexing step:

from sparknlp.annotator import BertSentenceEmbeddings, DocumentCharacterTextSplitter

from sparknlp.base import EmbeddingsFinisher



text_splitter = (

    DocumentCharacterTextSplitter()

    .setInputCols(["aligned_text"])

    .setOutputCol("chunked_docs")

    .setChunkSize(1200)

    .setChunkOverlap(120)

    .setExplodeSplits(False)

)



sentence_embeddings = (

    BertSentenceEmbeddings.pretrained("sent_small_bert_L2_128", "en")

    .setInputCols(["chunked_docs"])

    .setOutputCol("chunk_embeddings")

)



embeddings_finisher = (

    EmbeddingsFinisher()

    .setInputCols(["chunk_embeddings"])

    .setOutputCols(["finished_embeddings"])

    .setOutputAsVector(True)

    .setCleanAnnotations(False)

)



embedding_pipeline = Pipeline(

    stages=[text_splitter, sentence_embeddings, embeddings_finisher]

)

embedding_df = embedding_pipeline.fit(result_df).transform(result_df)

The notebook flattens chunk text, metadata, and vectors into rows that can be passed to a search index. It stops before writing to Elasticsearch; the index schema, metadata mapping, and retrieval configuration are left to the application.

Evaluate the result in your pipeline

The notebook illustrates the workflow with revenue-report, clinical-trial, and cloud-architecture PDFs. Its outputs show extracted text, aligned image and text annotations, generated captions, and reconstructed document text. Those examples demonstrate pipeline behavior; they are not a benchmark of caption accuracy, retrieval quality, throughput, or cost.

For your own corpus, inspect the aligned pairs and generated captions before indexing. Then evaluate retrieval against questions whose answers depend on figures or diagrams, and compare the results with your existing ingestion path. Keep the source document, page or slide metadata, and chunk metadata available so retrieved text can be traced back to its location.

Use layout alignment before standard retrieval stages

LayoutAlignerForVision and LayoutAlignerForText connect extracted text, image captioning, and document reconstruction in a Spark NLP pipeline. After reconstruction, standard Spark NLP stages can split the text and generate embeddings. The notebook provides a runnable starting point for adapting the workflow to your own document collection and vector store.

For implementation details, see the LayoutAligners document-understanding notebook and the Spark NLP API documentation for LayoutAlignerForVision and LayoutAlignerForText.

Frequently asked questions

What does LayoutAlignerForVision output?

It creates three annotation columns based on the configured output column: document annotations, image annotations, and prompt annotations. For an output column named aligned, those columns are aligned_doc, aligned_image, and aligned_prompt.

What does neighbor text add to image captioning?

When enabled, addNeighborText adds local text context around an aligned paragraph to the prompt passed to the vision model. neighborTextCharsWindow sets the character window on each side. The option provides context to the model; it does not guarantee that a generated caption is correct.

Which columns should I pass to LayoutAlignerForText?

Pass the document column from LayoutAlignerForVision and the caption column from your vision model, such as aligned_doc and image_caption. The transformer uses document layout metadata and caption metadata to rebuild text around the matching elements.

Does the example write vectors to Elasticsearch?

No. The notebook prepares chunk text, metadata, and embedding vectors, then stops before the Elasticsearch write. You need to implement the index mapping and write step for your own deployment.

Does this workflow guarantee better retrieval or lower cost?

No such guarantee follows from the pipeline alone. It limits the VLM stage to images that the aligner pairs with text, but caption quality, inference cost, and retrieval performance depend on your documents, model, configuration, and evaluation.