GLM-OCR: image-to-text with Transformers

This page is the Transformers contract of the GLM-OCR MIC and also a usage example. doc-builder renders the runnable blocks below as documentation, and its pytest plugin runs them as written against zai-org/GLM-OCR. The MIC runner runs the same blocks against the slim fixture and the checkpoint in integration.json. Lines marked # doc-builder: hide only matter there: they swap in the pinned checkpoint or variant and print the result record.

Read text from an image

import torch
from PIL import Image, ImageDraw
from transformers import AutoModelForImageTextToText, AutoProcessor, set_seed
import json  # doc-builder: hide
import os  # doc-builder: hide

model_id = "zai-org/GLM-OCR"
model_id = os.environ.get("MIC_MODEL", model_id)  # doc-builder: hide
variant = json.loads(os.environ.get("MIC_MODEL_KWARGS", "{}"))  # doc-builder: hide
device = "cuda" if torch.cuda.is_available() else "cpu"
set_seed(0)  # seeds torch, NumPy, and Python: loading, dropout, and data augmentation draw random numbers

processor = AutoProcessor.from_pretrained(model_id)
model = AutoModelForImageTextToText.from_pretrained(model_id).to(device)
if variant:  # doc-builder: hide
    del model  # doc-builder: hide
    model = AutoModelForImageTextToText.from_pretrained(model_id, **variant).to(device)  # doc-builder: hide
model.eval()

# Two small synthetic document images so the batch perturbation is meaningful.
def make_doc(text):
    image = Image.new("RGB", (224, 96), "white")
    ImageDraw.Draw(image).text((8, 40), text, fill="black")
    return image

images = [make_doc("INVOICE 42"), make_doc("TOTAL 7.50")]
messages = [[{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "Read the text in this image."}]}] for _ in images]
texts = [processor.apply_chat_template(m, add_generation_prompt=True, tokenize=False) for m in messages]
inputs = processor(text=texts, images=images, padding=True, return_tensors="pt")
inputs = inputs.to(device, dtype=model.dtype)

with torch.no_grad():
    logits = model(**inputs).logits
# With right padding the last real token of each row is at its attention_mask sum - 1.
last = inputs["attention_mask"].sum(dim=-1) - 1
probe_logits = logits[torch.arange(logits.shape[0]), last].float()
predicted = probe_logits.argmax(dim=-1)
print(processor.tokenizer.convert_ids_to_tokens(predicted))

One training step

# Language-model labels: copy the input ids and mask padded positions.
labels = inputs["input_ids"].clone()
labels[inputs["attention_mask"] == 0] = -100

model.train()
optimizer = torch.optim.SGD([p for p in model.parameters() if p.requires_grad], lr=0.01)
loss = model(**inputs, labels=labels).loss
loss.backward()
before = {n: p.detach().cpu().clone() for n, p in model.named_parameters() if p.requires_grad}  # doc-builder: hide
optimizer.step()
print(f"loss: {loss.item():.4f}")

if os.environ.get("MIC_RUNNER"):  # doc-builder: hide
    named = dict(model.named_parameters())
    checked = (
        "model.language_model.layers.0.self_attn.q_proj.weight",
        "model.language_model.layers.0.mlp.down_proj.weight",
        "model.visual.blocks.0.attn.qkv.weight",
    )
    result = {
        "runtime": {
            "device": model.device.type,
            "dtype": str(model.dtype).removeprefix("torch."),
            "attn_implementation": model.config._attn_implementation,
        },
        "predicted_tokens": predicted.tolist(),
        "answer_logits": probe_logits[:, predicted].tolist(),
        "training": {
            "loss": loss.item(),
            "gradients": {n: bool(named[n].grad is not None and torch.count_nonzero(named[n].grad) > 0 and torch.isfinite(named[n].grad).all()) for n in checked},
            "updated": any(not torch.equal(p.detach().cpu(), before[n]) for n, p in model.named_parameters() if n in before),
        },
    }
    print(json.dumps(result))