+1 (726) 227-3241

Document Extraction with a Vision-Language Model on SageMaker AI

Invoice fields, claim forms, purchase orders, scanned contracts: the pipeline used to be OCR plus a pile of regular expressions plus a human fixing the 15 percent that broke. A vision-language model (VLM) replaces most of that with one call that takes page pixels and returns typed JSON. This tutorial deploys an open-weight VLM to a SageMaker AI endpoint, forces it to emit schema-valid JSON, and then runs the same model over a 200,000-page backfill without paying for a GPU that sits idle.

Everything here uses the SageMaker Python SDK v2 (pip install "sagemaker>=2.230") from a Studio space or any machine with credentials and a SageMaker execution role.

Why a VLM instead of OCR plus rules

The honest trade-off, because a VLM is not always the right answer:

OCR + rules (e.g. Textract queries)VLM on SageMaker AI
New document layoutNew rules, new releaseUsually just a prompt change
Checkboxes, stamps, handwriting notesFragileHandles them as pixels
Tables spanning pagesPainfulGood with page context in prompt
Cost per pageCents or lessHigher, unless batched on your own GPU
DeterminismHighNeeds constrained decoding and evals
Audit trailBounding boxes for freeYou must add grounding yourself

A pattern that works well in production: Textract for high-volume, stable forms, and a VLM for the long tail of layouts plus anything requiring reasoning across the page. If you only have the long tail, go straight to the VLM.

Step 1: deploy the model

Qwen2.5-VL-7B-Instruct is a reasonable default: strong document understanding, Apache-2.0 licensed, and it fits comfortably on a single 48 GB GPU at BF16. Serve it with the Large Model Inference (LMI) container, which bundles vLLM and therefore handles multi-image prompts and structured output.

import sagemaker
from sagemaker.djl_inference import DJLModel

role = sagemaker.get_execution_role()

model = DJLModel(
    model_id="Qwen/Qwen2.5-VL-7B-Instruct",
    role=role,
    env={
        "OPTION_ROLLING_BATCH": "vllm",
        "OPTION_TENSOR_PARALLEL_DEGREE": "1",
        "OPTION_MAX_MODEL_LEN": "16384",
        "OPTION_LIMIT_MM_PER_PROMPT": "image=4",
        "OPTION_DTYPE": "bf16",
    },
)

predictor = model.deploy(
    initial_instance_count=1,
    instance_type="ml.g6e.2xlarge",
    endpoint_name="doc-vlm",
    container_startup_health_check_timeout=900,
)

Three settings earn their keep:

  • OPTION_LIMIT_MM_PER_PROMPT caps images per request. Without it, one caller sending a 60-page PDF as 60 images will blow up the KV cache and take the endpoint down for everyone.
  • OPTION_MAX_MODEL_LEN matters more than usual because image tokens are expensive. A single 1,600 px page can consume well over a thousand tokens.
  • ml.g6e.2xlarge (one L40S, 48 GB) leaves room for the KV cache at 16k context. ml.g5.2xlarge (24 GB A10G) works only if you shrink context and image resolution; ml.g6e.12xlarge is the step up when you want a 32B-class model.

Give the health check at least 900 seconds. The weights are roughly 16 GB and vLLM compiles CUDA graphs on first load.

Step 2: get the pixels right

Most bad extractions are input problems, not model problems. Render PDFs yourself rather than trusting whatever the upstream system sends:

import base64, io
import pypdfium2 as pdfium

def page_images(pdf_bytes: bytes, dpi: int = 200, max_px: int = 1600):
    doc = pdfium.PdfDocument(pdf_bytes)
    out = []
    for page in doc:
        pil = page.render(scale=dpi / 72).to_pil()
        pil.thumbnail((max_px, max_px))
        buf = io.BytesIO()
        pil.save(buf, format="JPEG", quality=90)
        out.append(base64.b64encode(buf.getvalue()).decode())
    return out

Practical rules learned the hard way:

  • 200 DPI is the sweet spot. At 150 DPI small print in footers starts disappearing; above 300 DPI you pay for tokens that add no accuracy.
  • Cap the long edge around 1,600 px. Large images are silently downscaled by the processor anyway, so you are often paying to upload pixels the model never sees.
  • JPEG quality 90 over PNG: a quarter of the bytes, no measurable accuracy loss on text.
  • De-skew and drop blank pages before the call. A blank page still costs a full image's worth of tokens and invites hallucinated fields.

Step 3: force schema-valid JSON

Never parse prose. vLLM in the LMI container supports guided decoding, so hand it a JSON schema and the sampler can only emit tokens that keep the output valid.

import json

SCHEMA = {
    "type": "object",
    "properties": {
        "invoice_number": {"type": "string"},
        "invoice_date": {"type": "string", "description": "ISO 8601 date"},
        "vendor_name": {"type": "string"},
        "currency": {"type": "string", "enum": ["USD", "EUR", "GBP", "CAD"]},
        "total_amount": {"type": "number"},
        "line_items": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "description": {"type": "string"},
                    "quantity": {"type": "number"},
                    "unit_price": {"type": "number"},
                },
                "required": ["description", "quantity", "unit_price"],
            },
        },
        "fields_not_found": {"type": "array", "items": {"type": "string"}},
    },
    "required": [
        "invoice_number", "invoice_date", "vendor_name",
        "currency", "total_amount", "line_items", "fields_not_found",
    ],
}

PROMPT = (
    "Extract the invoice fields from these page images. "
    "Copy values exactly as printed; do not compute or normalize totals. "
    "If a required field is absent, use an empty string or 0 and list the "
    "field name in fields_not_found. Never guess."
)

def extract(images):
    content = [{"type": "image_url",
                "image_url": {"url": f"data:image/jpeg;base64,{b}"}}
               for b in images]
    content.append({"type": "text", "text": PROMPT})

    payload = {
        "messages": [{"role": "user", "content": content}],
        "temperature": 0.0,
        "max_tokens": 2048,
        "response_format": {
            "type": "json_schema",
            "json_schema": {"name": "invoice", "schema": SCHEMA},
        },
    }
    return json.loads(predictor.predict(payload)["choices"][0]["message"]["content"])

The two details that move accuracy most:

  1. fields_not_found as a required array. Given a schema that demands invoice_number, a model will invent one. Giving it a sanctioned way to say "not on the page" converts silent fabrication into an explicit, routable signal.
  2. "Do not compute." VLMs are mediocre at arithmetic. Extract the printed values and verify the sum in Python: abs(sum(li["quantity"] * li["unit_price"] for li in r["line_items"]) - r["total_amount"]) <= 0.02 * r["total_amount"]. A failed check is a review-queue trigger, not a retry.

Set temperature to 0.0. Sampling on an extraction task only buys you non-reproducible bugs.

Step 4: an eval set before a production rollout

Fifty labelled pages, stratified across vendors and layouts, is enough to make decisions. Score per field, not per document, and keep the two error classes separate:

def score(pred, truth, keys):
    rows = {}
    for k in keys:
        p, t = str(pred.get(k, "")).strip(), str(truth.get(k, "")).strip()
        rows[k] = {
            "exact": p == t,
            # wrong value while claiming to have found it: the expensive error
            "silent_error": p != t and k not in pred.get("fields_not_found", []),
        }
    return rows

Track exact-match rate and silent-error rate per field. A field with 92 percent exact match and near-zero silent errors can be automated with a review queue; 97 percent exact match with 3 percent confident-and-wrong cannot, because nobody downstream knows which rows to distrust. Wire this script into the pipeline that promotes the endpoint so a prompt tweak cannot regress a field unnoticed — the same gating idea as in Evaluate Before You Ship.

Step 5: the backfill, without paying for idle GPUs

A real-time endpoint is the wrong tool for 200,000 archived pages. Two better options:

Async inference with scale-to-zero for a queue that drains overnight:

from sagemaker.async_inference import AsyncInferenceConfig

async_predictor = model.deploy(
    initial_instance_count=1,
    instance_type="ml.g6e.2xlarge",
    endpoint_name="doc-vlm-async",
    async_inference_config=AsyncInferenceConfig(
        output_path="s3://my-bucket/vlm/out/",
        failure_path="s3://my-bucket/vlm/err/",
        max_concurrent_invocations_per_instance=4,
    ),
)

Then register the variant with Application Auto Scaling using the ApproximateBacklogSizePerInstance metric and MinCapacity=0. Invocations return immediately with an S3 output location, the endpoint scales up from zero when the queue fills, and scales back to zero when it empties. For a nightly job this is usually the cheapest managed option.

A Processing job when the work is embarrassingly parallel and you would rather not manage an endpoint at all: 8 instances, each loading the model once with vLLM in offline mode, sharding an S3 manifest by instance index. No per-request HTTP overhead, and continuous batching keeps the GPUs near saturation — typically 2 to 4x the throughput per GPU-hour of the same model behind a real-time endpoint. Use managed spot instances and checkpoint completed shard IDs to S3 so an interruption resumes instead of restarting.

Per-page cost falls out of throughput: at roughly 1,800 image-plus-text tokens per page and around 6 pages per second on one L40S with continuous batching, a ml.g6e.2xlarge at a couple of dollars per hour lands near a tenth of a cent per page. Measure it on your own documents before you quote a number to anyone; page density swings this by 3x.

Production checklist

  • PII. Documents are the highest-risk data you will put through a model. Deploy VPC-only with no internet route, KMS-encrypt the S3 input and output paths, and turn off data capture or route it to a locked-down bucket. See Private by Default.
  • Input validation. Reject requests over your image-per-prompt cap at the application layer with a clear error instead of letting them reach the endpoint and evict everyone's KV cache.
  • Grounding. If auditors need to see where a value came from, ask the model for a bounding box per field, or run Textract alongside and match extracted strings to word coordinates. Do not promise grounding a bare VLM cannot give you.
  • Drift. A new vendor template is drift. Alarm on fields_not_found rate and on validation-check failures per document type; both move before anyone files a complaint. Mechanics in Detect Drift in Production.
  • Human review. Route documents that fail arithmetic checks or report missing required fields to a queue. Log the corrections: that is next quarter's fine-tuning set.
  • Teardown. predictor.delete_endpoint(delete_endpoint_config=True) and predictor.delete_model(). A forgotten ml.g6e endpoint is a four-figure monthly surprise.

Where this goes next

Once corrected extractions accumulate, a LoRA fine-tune on a few thousand of your own pages usually beats prompt engineering on a larger model, and it lets you drop to a 7B or even 3B model for the same accuracy — cheaper per page and faster. Multi-adapter serving means one endpoint can host a per-document-type adapter set: see Serve Many Fine-Tunes from One GPU.

Building a document extraction pipeline on SageMaker AI and want a second pair of eyes on the architecture or the eval harness? Talk to us.