baidu/Unlimited-OCR Colab Quickstart: layout-aware OCR, tiled vs global modes

Quick start, three steps to run the Colab demo

  • Open a Google Colab runtime and enable a GPU (if you see “No GPU detected! In Colab: Runtime -> Change runtime type -> GPU.”, switch the runtime).
  • Install the notebook dependencies with a single pip cell:

    pip install transformers==4.57.1 Pillow matplotlib einops addict easydict pymupdf psutil accelerate
  • Run the model‑download cell that loads baidu/Unlimited-OCR, the notebook prints a status line like “>> Downloading model (~6 GB for 3B params in BF16). First run takes a while…”.

What this pipeline does, and when to temper expectations

A single vision‑language model (baidu/Unlimited-OCR, 3B parameters) can often produce useful structured Markdown and JSON that captures headings, paragraphs and tables from high‑resolution scans. That reduces the need for separate detection + recognition + layout parsing stages used in classical OCR stacks.

Results vary by scan quality, font size, language, and document cleanliness. Treat the Colab demo as an engineering starting point: it shows how to load the model, switch precision based on GPU support, choose tiled versus global image strategies, and run long‑context multi‑page decoding. The demo does not provide production SLAs, benchmarked accuracy numbers, or hardened privacy controls, see the Limitations and Production Considerations section before you deploy.

Why a single model?

Vision‑language OCR models combine image feature extraction and text decoding in one model. That makes it easier to output structured artifacts (Markdown/JSON) without stitching multiple tools together. The tradeoffs are memory, decoding stability on long documents, and careful handling of tiled versus global image views for dense layouts.

Environment & dependencies

Target environment: a GPU‑enabled Google Colab runtime. The notebook installs these packages (exact pip strings used):

  • transformers==4.57.1
  • Pillow
  • matplotlib
  • einops
  • addict
  • easydict
  • pymupdf
  • psutil
  • accelerate

After installation the notebook prints helpful runtime diagnostics such as:

“>> GPU: {gpu_name}”

“>> Using dtype: {DTYPE}”

GPU precision: BF16 vs FP16

The demo auto‑selects mixed precision based on hardware:

  • use_bf16 = torch.cuda.is_bf16_supported()
  • DTYPE = torch.bfloat16 if use_bf16 else torch.float16

Why it matters: bfloat16 (BF16) keeps a wider exponent range than float16 (FP16) and can reduce numeric instability on GPUs that support it. If BF16 is available the notebook uses it; otherwise it falls back to FP16. BF16 support also depends on compatible PyTorch/CUDA builds, verify your runtime if you expect BF16 to be available.

Loading the model (trust_remote_code, safetensors), with a security note

The notebook loads the tokenizer and model like this (exact calls used):

  • tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
  • model = AutoModel.from_pretrained(MODEL_NAME, trust_remote_code=True, use_safetensors=True, torch_dtype=DTYPE)
  • model.eval(); model.cuda()

MODEL_NAME is “baidu/Unlimited-OCR” (3B parameters). The notebook prints status lines such as:

“>> Downloading model (~6 GB for 3B params in BF16). First run takes a while…”

“>> Model loaded and moved to GPU.”

Note on size: the printed “~6 GB for 3B params in BF16” refers roughly to the weight file size (3B parameters × 2 bytes ≈ 6 GB). Expect additional memory and disk overhead for optimizer states, activation buffers, caches and the tokenizer.

Security note: trust_remote_code=True executes model code supplied by the model repo. Inspect that code before running in an environment with sensitive data, or run it in an isolated runtime. Treat the repository like any third‑party code you would audit before production use.

Creating synthetic test pages (the demo samples)

The tutorial builds layout‑rich sample pages (image size W = 1240, H = 1754) using PIL to exercise headings, paragraphs, tables and footers. The sample table strings are exact and useful for validating table extraction:

  • Header row: [“Region”, “Q1”, “Q2”, “Q3”]
  • Data rows:
    • [“North”, “12.4”, “13.1”, “15.0”]
    • [“South”, “9.8”, “10.2”, “11.7”]
    • [“East”, “14.3”, “13.9”, “16.2”]
    • [“West”, “11.1”, “12.5”, “12.9”]

These synthetic pages are a safe starting point for early validation before you upload real invoices, reports, or forms.

Single‑page inference: Gundam (tiled + global) vs Base (single view)

The notebook exposes two single‑image inference strategies. Use the tiled approach for dense, small fonts and the global single‑view for clean printed pages.

Brief definitions:

  • Gundam, notebook nickname for tiled+global fusion: crop_mode=True with smaller tiles so the model sees local detail plus a global context.
  • Base, single global view: crop_mode=False and a larger image_size for simpler layouts and faster runtime.

base_size and image_size control how the image is scaled and how tiles are sized. base_size sets the reference scale for the model, image_size controls the crop/tile resolution.

Exact example arguments used in the demo

Gundam‑mode infer() (dense / small text):

  • prompt = “<image>document parsing.”
  • image_file = IMAGE_PATH
  • output_path = “outputs/single_gundam”
  • base_size = 1024
  • image_size = 640
  • crop_mode = True
  • max_length = 32768
  • no_repeat_ngram_size = 35
  • ngram_window = 128
  • save_results = True

Base‑mode infer() (clean print):

  • prompt = “<image>document parsing.”
  • image_file = IMAGE_PATH
  • output_path = “outputs/single_base”
  • base_size = 1024
  • image_size = 1024
  • crop_mode = False
  • max_length = 32768
  • no_repeat_ngram_size = 35
  • ngram_window = 128
  • save_results = True

Generation knobs: the demo sets max_length=32768 to allow long structured outputs and uses repetition controls (no_repeat_ngram_size and ngram_window) to discourage decoder loops on long contexts. Two important caveats:

  • Not every tokenizer/model supports extremely long token windows; verify your model/config can handle the desired max_length.
  • ngram_window is a wrapper parameter in the demo’s infer() helper, it is not a standard Transformers.generate argument. The notebook maps it into generation logic to control repetition across a sliding context window.

Multi‑page / PDF workflow

For PDFs the demo rasterizes pages to PNGs and decodes across pages with a single long horizon infer_multi() call so the model can reason about page structure in sequence.

PDF rasterization (exact function outline):

  • def pdf_to_images(pdf_path, dpi=300)
  • matrix = fitz.Matrix(dpi / 72, dpi / 72)
  • save pages as “page_{i + 1:04d}.png”

Why 300 dpi? It balances OCR fidelity and downstream recognition for small table text. If memory is constrained, test lower DPI (e.g., 150) but expect more recognition errors on tiny fonts.

Multi‑page infer_multi() arguments used in the demo:

  • prompt = “<image>Multi page parsing.”
  • image_files = page_images
  • output_path = “outputs/multi_page”
  • image_size = 1024
  • max_length = 32768
  • no_repeat_ngram_size = 35
  • ngram_window = 1024
  • save_results = True

The demo widens ngram_window to 1024 for multi‑page decoding so repetition controls operate across a longer multi‑page context.

Inspecting outputs

The notebook saves structured files and provides a small preview routine. It scans these text extensions exactly as shown in the code:

TEXT_EXTS = {“.txt”, “.md”, “.mmd”, “.json”}

Previewing is limited to the first 1, 500 characters to give a quick glance without loading huge artifacts.

Quick decision guide, mode selection

  • Single image, dense/small text: use infer() with Gundam‑style tiling, image_size 640, crop_mode=True.
  • Single image, clean print: use infer() with Base (single view), image_size 1024, crop_mode=False.
  • Multi‑page or PDF: use infer_multi() with image_size=1024 and a larger ngram_window (1024).
  • Long documents: keep max_length=32768 and use the no_repeat_ngram settings to reduce generation degeneration.
  • Your own files: upload via the Colab sidebar and point image_file / pdf_to_images() at them.

Limitations and production considerations

Known gaps in the demo

  • No quantitative OCR accuracy metrics or table‑extraction scores are provided for common benchmarks or the synthetic samples. Validate against your own labeled set.
  • No measured GPU memory footprints or latency numbers are supplied, the printed “~6 GB for 3B params in BF16” is an approximate weight size; runtime memory and activation buffers add to that.
  • Robustness to noisy scans, heavy skew, handwriting, non‑Latin scripts, and very low DPI is not demonstrated.
  • The exact schema used for saved JSON/MMD outputs should be inspected; implement a deterministic post‑processor for reliable ingestion into downstream systems.

Production considerations

  • Trust_remote_code=True poses security risk, audit the model repo or run inference behind hardened controls.
  • Privacy and compliance: avoid uploading sensitive PII to public Colab notebooks unless you have explicit controls; consider on‑prem or private cloud deployment for regulated data.
  • Scaling: for thousands of PDFs you’ll need batching, autoscaling inference servers, model quantization, or sharded pipelines, the Colab demo is not a production throughput solution.
  • Error handling: build retries, chunking fallbacks (e.g., reduce image_size or switch from Base to Gundam), and monitoring to catch decoder collapse or partial outputs.

Action checklist, what to run next

  • Run the Colab notebook on a GPU runtime and validate the device diagnostics: confirm torch.cuda.is_bf16_supported().
  • Test both modes on 5-10 representative PDFs: try Gundam for scanned invoices and Base for clean reports.
  • Inspect saved JSON/MMD outputs and write a small deterministic extractor that maps model output fields to your schema.
  • Measure peak GPU memory and per‑page latency on your hardware to decide whether a 3B model fits your cost/throughput targets (estimate a modern GPU with 16-24 GB for comfortable tiled runs; smaller GPUs may run but expect memory pressure).
  • Audit trust_remote_code content, and decide whether to run inference in a private environment for sensitive documents.

Key takeaways, quick Q&A

  • Do I need a GPU to run this pipeline?

    Yes. The notebook asserts torch.cuda.is_available() and prints the guidance “No GPU detected! In Colab: Runtime -> Change runtime type -> GPU.” A GPU is required for practical inference with the 3B model; for comfortable tiled inference plan for a modern GPU (many practitioners target 16-24 GB VRAM).

  • What dtype should I use, BF16 or FP16?

    The notebook automatically selects BF16 if torch.cuda.is_bf16_supported() is true; otherwise it falls back to FP16. BF16 is preferable when supported because it preserves a wider numeric range, but BF16 availability also depends on matching PyTorch/CUDA builds.

  • When should I use Gundam (tiled) mode vs Base (single view)?

    Use Gundam (crop_mode=True, image_size=640) for dense or small text and complex layouts; use Base (crop_mode=False, image_size=1024) for cleaner printed pages. Gundam trades runtime for improved local detail.

  • How do I parse a multi‑page PDF?

    Rasterize the PDF at dpi=300 with PyMuPDF (fitz.Matrix(dpi/72, dpi/72)) into PNGs named “page_{i + 1:04d}.png”, then call infer_multi() with image_size=1024, max_length=32768 and ngram_window=1024 to let the model decode across pages in a longer context.

  • How does the pipeline avoid repetition during long outputs?

    It sets max_length=32768 and uses repetition controls: no_repeat_ngram_size=35 plus an ngram_window (128 for single pages, 1024 for multi‑page) implemented by the notebook’s generation wrapper. These knobs reduce degenerative loops, but you should tune them on your documents.

Final note

The Colab demo is a practical, reproducible starting point for exploring layout‑aware OCR with a single vision‑language model. Use the quick checklist to validate on your documents, audit the model code for security, and measure memory/latency on your target GPUs. If you want help converting the model output to a deterministic table extractor or sketching a scalable, secure inference architecture, I can outline a deployment plan tuned to your SLAs and document mix.

0