Most of the vision work we are hired for used to end at a label: is this a defect, which of these twelve categories, what is the bounding box. That framing is now the wrong tool for a large share of client problems. When the real question is "pull the invoice number, the vendor, and the line items out of this scanned PDF" or "describe what is wrong in this inspection photo in a sentence a technician can act on", you want a vision-language model (VLM): one model that takes an image plus a text prompt and returns text.
On the TensorFlow and Keras stack, KerasHub ships open-weight VLMs you can fine-tune and host yourself - PaliGemma and PaliGemma 2 being the workhorses. This tutorial builds a document-extraction task end to end: prompt the base model, fine-tune it with LoRA on a few hundred labelled pages, evaluate with a metric that reflects the business requirement, and export something servable. It assumes you have read our Gemma fine-tuning tutorial; the mechanics are similar and we do not repeat the basics here.
Why a VLM instead of OCR plus a classifier
The classic pipeline is OCR to text, then regex or a trained extractor. It still wins when your documents are uniform, high-resolution, and machine-printed - it is cheaper, deterministic, and easier to audit. A VLM earns its keep when:
- Layout varies wildly between vendors or sites, so template rules multiply forever.
- The signal is partly visual - a stamp, a checkbox, a handwritten annotation, a photo of a broken part.
- You need free-text output (a summary, a rationale) rather than a fixed label set.
- You have few labels. A 200-example LoRA fine-tune of a pretrained VLM often beats a from-scratch extractor trained on the same 200 examples by a wide margin.
We usually recommend a hybrid in production: OCR text and the page image into the prompt, so the model gets crisp character evidence plus layout.
Setup
pip install -U keras keras-hub pillow
# plus a backend; we train on JAX
import os
os.environ["KERAS_BACKEND"] = "jax"
import keras, keras_hub, numpy as np
from PIL import Image
keras.mixed_precision.set_global_policy("mixed_bfloat16")
PaliGemma checkpoints are gated the same way Gemma is: accept the licence on Kaggle and set KAGGLE_USERNAME / KAGGLE_KEY so KerasHub can download the preset. Check the KerasHub model list for the presets available to you and their input resolutions - the suffix on a preset name (224, 448, 896) is the image resolution it was trained at, and it matters more than the parameter count for document work.
Baseline: prompt before you train
vlm = keras_hub.models.PaliGemmaCausalLM.from_preset("pali_gemma2_mix_3b_448")
image = np.array(Image.open("invoice_0001.png").convert("RGB").resize((448, 448)))
out = vlm.generate(
{
"images": image,
"prompts": "answer en What is the invoice total?",
},
max_length=64,
)
print(out)
Two things to internalise here. First, the mix presets are instruction-tuned and respond to task prefixes such as answer en, caption en, or detect; the plain pretrained presets mostly do not, and will look broken if you prompt them conversationally. Second, resolution is the single biggest lever for dense documents. A 448 preset reads a clean A4 invoice; small print, multi-column forms, or engineering drawings usually need 896 or a tiling strategy where you cut the page into overlapping crops and query each one.
Always measure this baseline on your evaluation set before fine-tuning. On a meaningful minority of client projects the zero-shot mix model is already good enough for the field values that matter, and the project turns into an evaluation-and-serving job rather than a training job.
Data format
Each training example is an image, a prompt, and the target response. Keep the prompt template byte-identical between training and inference - a VLM will happily learn to depend on a trailing space.
PROMPT = "answer en Extract invoice_number, vendor, total as JSON."
def to_example(record):
img = Image.open(record["path"]).convert("RGB").resize((448, 448))
return {
"images": np.array(img),
"prompts": PROMPT,
"responses": record["target_json"], # e.g. '{"invoice_number": "A-1182", ...}'
}
Build it into a tf.data pipeline so loading and resizing overlap with the accelerator step - the same discipline as in our tf.data tutorial:
import tensorflow as tf
def gen(records):
for r in records:
yield to_example(r)
sig = {
"images": tf.TensorSpec((448, 448, 3), tf.uint8),
"prompts": tf.TensorSpec((), tf.string),
"responses": tf.TensorSpec((), tf.string),
}
train_ds = (
tf.data.Dataset.from_generator(lambda: gen(train_records), output_signature=sig)
.shuffle(1024)
.batch(2)
.prefetch(tf.data.AUTOTUNE)
)
Practical notes on labelling, which is where these projects actually succeed or fail:
- Serialise targets deterministically. Fixed key order, no pretty-printing, one canonical date and currency format. Inconsistent targets teach the model to be inconsistent.
- Label the absent case. Include pages where a field genuinely is not present and make the target explicit (
"total": null). Otherwise the model hallucinates a plausible number for every page, which is the failure mode clients notice first. - Stratify the split by vendor / template / capture device, not randomly. A random split leaks near-duplicate pages across train and test and will tell you the model is far better than it is.
- A few hundred pages is a real dataset. We have shipped useful extractors on 300 labelled pages. Spend the effort on coverage of layout variety, not volume.
LoRA fine-tune
Freeze the weights, train small adapters, keep the sequence length tight:
vlm.preprocessor.sequence_length = 256
vlm.backbone.enable_lora(rank=8)
vlm.summary() # trainable params should be a few million, not billions
vlm.compile(
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=keras.optimizers.AdamW(learning_rate=5e-5, weight_decay=0.01),
weighted_metrics=[keras.metrics.SparseCategoricalAccuracy()],
)
vlm.fit(
train_ds,
validation_data=val_ds,
epochs=3,
callbacks=[
keras.callbacks.EarlyStopping(patience=1, restore_best_weights=True),
keras.callbacks.TensorBoard(log_dir="logs/vlm"),
],
)
Tuning guidance that has held up across engagements:
- Freeze the vision tower first. Adapters on the language decoder plus a frozen image encoder train faster, need less memory, and are usually enough when your images resemble natural photos or scanned documents. Unfreeze the tower only if your imagery is genuinely out of distribution (thermal, radiographic, microscopy) and the frozen-tower run has plateaued short of requirement.
- Batch size is dominated by image tokens. A 448 image expands to hundreds of tokens before your prompt starts. If batch 2 is all that fits at 24 GB, use
gradient_accumulation_stepson the optimiser rather than dropping resolution - resolution is the thing you cannot get back. - Two to three epochs. These models memorise a small set fast; watch validation loss and stop.
- bfloat16 throughout. Float16 on a VLM decoder is a reliable way to spend an afternoon debugging NaNs.
Evaluate on the business metric
Token accuracy and validation loss will not tell you whether the extractor is deployable. Score the thing the client cares about: per-field exact match after normalisation, plus a parse rate.
import json
def normalise(v):
if v is None: return None
return str(v).strip().lower().replace(",", "")
def score(model, records, fields=("invoice_number", "vendor", "total")):
hits = {f: 0 for f in fields}
parsed = 0
for r in records:
img = np.array(Image.open(r["path"]).convert("RGB").resize((448, 448)))
raw = model.generate({"images": img, "prompts": PROMPT}, max_length=128)
try:
pred = json.loads(raw[len(PROMPT):].strip())
parsed += 1
except json.JSONDecodeError:
continue
truth = json.loads(r["target_json"])
for f in fields:
if normalise(pred.get(f)) == normalise(truth.get(f)):
hits[f] += 1
n = len(records)
return {"parse_rate": parsed / n, **{f: hits[f] / n for f in fields}}
print(score(vlm, eval_records))
Report it per field and per vendor, and report the base-model numbers next to the fine-tuned ones. Two rules we hold to:
- Never validate JSON by trusting the model. Parse it, then validate against a schema (Pydantic or
jsonschema) and reject non-conforming output at the API boundary. A retry with a stricter prompt fixes most of what fails. - Route low confidence to a human. For extraction, sequence log-probability plus a business rule (does the line-item sum match the total?) is a cheap, effective abstention signal. A pipeline that extracts 85% of pages automatically and escalates 15% is worth far more than one that silently guesses on all of them.
Save, export, serve
vlm.save_to_preset("./paligemma2_invoices")
For serving, the honest answer is that a VLM decoder is not a good fit for the SavedModel-plus-TensorFlow-Serving pattern we describe in Serving TensorFlow Models in Production. Autoregressive generation with variable-length output wants a Python service that owns the model in memory. What has worked well for us:
- Reload the preset in a FastAPI (or equivalent) worker on the JAX or TensorFlow backend, one model per GPU process, and queue requests.
- Batch by page, not by document, and cap
max_lengthhard - a runaway generation is your latency outlier. - Push preprocessing (PDF rasterisation, deskew, resize) into a separate CPU service. It is embarrassingly parallel and should not occupy GPU workers.
- For high volume, quantise the decoder before you buy more GPUs; see Shrinking TensorFlow Models for the trade-offs and how to measure the quality loss on your own eval set.
- For on-device or in-browser inference, smaller multimodal Gemma variants converted through the LiteRT path are the route - covered in Migrating from TensorFlow Lite to LiteRT.
Governance you will be asked about
Document extraction touches invoices, claims, and medical forms, so plan for the questions before legal asks them. Keep training images in a controlled bucket with retention rules; log prompts and outputs but redact extracted PII from those logs, or hash it; and record the preset name, LoRA rank, dataset hash, and eval scores next to each shipped checkpoint. If the data cannot leave the client's environment at all, the whole point of an open-weight VLM is that this pipeline runs inside their VPC - and if the training data is sensitive enough to need formal guarantees, see Training on Sensitive Data.
Where this goes wrong
The three failure modes we get called in to fix, in order of frequency: resolution too low for the print size, so the model confidently misreads digits; an evaluation split that leaked near-duplicate pages, so the pilot numbers never survive contact with a new vendor; and no abstention path, so a 92%-accurate extractor quietly writes 8% wrong values into a finance system. All three are cheap to avoid at design time and expensive to fix after go-live.
If you are scoping a multimodal extraction or inspection project on the TensorFlow stack, our Generative AI on the TensorFlow Stack page describes how we run a feasibility sprint - typically two weeks to a measured baseline on your own documents. Get in touch with a description of the pages and the fields and we can tell you quickly whether this is a prompt problem, a fine-tune problem, or still an OCR problem.