feat(tableformer): export TableFormer to ONNX (encoder + autoregressive decoder)

First milestone of the TableFormer port: scripts/export_tableformer.py loads
docling's TableModel04_rs (accurate) from the docling-ibm-models safetensors and
exports two ONNX graphs, both verified against PyTorch (max abs diff < 1.5e-5):

  encoder.onnx : image[1,3,448,448]      -> memory[784,1,512]
  decoder.onnx : tags[seq,1] + memory     -> logits[1,13], hidden[1,512]

Key findings that de-risk the rest of the port:
- The OTSL structure vocabulary is just 13 tokens (wordmap.json).
- The decoder runs with cache=None (re-embeds the full prefix each step), so it
  exports cleanly with a dynamic seq axis and is a plain Rust loop — no KV-cache.
- Needed `requires_grad_(False)` on the params for the decoder to trace.

The critical feasibility gate (ONNX-exporting an autoregressive transformer) is
passed. Remaining work is the Rust side — decode loop, OTSL→grid serialization
(port otsl.py), and cell-text matching — scoped in PDF_CONFORMANCE.md. The ONNX
weights are gitignored (exported/downloaded like the layout/OCR models).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: artiz

PDF_CONFORMANCE.md

@@ -120,24 +120,37 @@ rows/columns); docling runs **TableFormer**, an autoregressive transformer that
 predicts the table structure as an OTSL/HTML tag sequence plus per-cell bounding
 boxes, which recovers spanning headers and merged cells we cannot.
 
-**Scope of a port** (large — own PR, likely staged over several):
-
-1. **Weights.** TableFormer ships in `docling-ibm-models` (`TableModel04_rs`,
-   "accurate"/"fast" variants). Export the encoder + the two decoders to ONNX
-   from the published checkpoint; confirm the license permits redistribution of
-   a converted model.
-2. **Inference loop.** Unlike the layout/OCR models (single `Session::run`),
-   TableFormer is **autoregressive**: encode the table-crop image once, then step
-   the structure decoder to emit OTSL tokens until `<end>`, feeding each token
-   back in. The cell-bbox decoder runs per predicted cell. This is a real
-   decoding loop in `fleischwolf-pdf`, not a one-shot call — budget for KV-cache
-   handling and a token vocabulary/OTSL grammar.
-3. **Cell content.** Map predicted cell bounding boxes back onto the PDF text
-   cells (we have these) to fill cell text — the same matching docling does for
-   "PDF" tables (it does not OCR programmatic tables).
-4. **Serialization.** Convert the predicted OTSL grid (with row/col spans) to the
-   `Table` node; the Markdown table serializer already exists but assumes a plain
-   grid, so spans need representing.
+**Status — ONNX export DONE ✅.** `scripts/export_tableformer.py` loads
+`TableModel04_rs` (`accurate`, resnet18 + 6-layer encoder + 6-layer decoder) from
+the `docling-ibm-models` safetensors and exports two graphs, **both verified
+against PyTorch** (max abs diff < 1.5e-5):
+
+- `encoder.onnx` — `image[1,3,448,448] → memory[784,1,512]`
+- `decoder.onnx` — `tags[seq,1] + memory → logits[1,13], hidden[1,512]`
+
+The OTSL vocabulary is only **13 tokens** (`wordmap.json`). The decoder runs with
+`cache=None` (re-embeds the full token prefix each step), so it exports cleanly
+with a dynamic `seq` axis and is driven as a simple loop from Rust — no KV-cache
+machinery needed. The ONNX/weights are gitignored (downloaded/exported, like the
+layout/OCR models).
+
+**Remaining (the Rust inference, staged):**
+
+1. **Decode loop** in `fleischwolf-pdf`: crop+resize the table region to 448²,
+   normalise, `encoder.onnx` once, then loop `decoder.onnx` — `argmax` the logits,
+   apply docling's two structure-correction rules (first-line `xcel→lcel`,
+   `ucel`-then-`lcel → fcel`), append, stop at `<end>`. Yields the OTSL tag run.
+2. **Bbox decoder** (`bbox.onnx`, not yet exported): per-cell hidden → box, for
+   matching. Interim: skip it and match by grid geometry.
+3. **OTSL → grid.** `docling_ibm_models/tableformer/otsl.py` is the reference
+   (`fcel/ched/rhed/srow/ecel/lcel/ucel/xcel/nl`); port it to a `Table` with
+   row/col spans (the Markdown serializer needs span support added).
+4. **Cell content.** Map the PDF text cells we already extract onto the grid
+   cells (docling does not OCR programmatic tables).
+
+A cheaper interim improvement (not docling-exact, but closes some diff): better
+geometric reconstruction — detect header rows, merge obvious spanning cells, and
+handle the multi-line header cells that currently shatter into many columns.
 
 A cheaper interim improvement (not docling-exact, but closes some diff): better
 geometric reconstruction — detect header rows, merge obvious spanning cells, and

scripts/export_tableformer.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+"""Export docling's TableFormer (TableModel04_rs) to ONNX for the Rust pipeline.
+
+TableFormer is autoregressive: an image encoder + a tag-transformer encoder run
+once to produce a memory tensor, then a decoder step is looped to emit OTSL
+structure tokens, and a bbox decoder turns the per-cell hidden states into boxes.
+We export three graphs and drive the loop from Rust:
+
+  encoder.onnx : image[1,3,448,448]            -> memory[784,1,512]
+  decoder.onnx : tags[seq,1] + memory           -> logits[1,V], hidden[1,512]
+  bbox.onnx    : memory + cell_hidden[ncells,512] -> classes, coords  (optional)
+
+Run inside the docling venv:
+  .venv-compare/bin/python scripts/export_tableformer.py <accurate-artifacts-dir> [out_dir]
+"""
+import json
+import os
+import sys
+import warnings
+
+import torch
+import torch.nn as nn
+
+warnings.filterwarnings("ignore")
+
+ART = sys.argv[1]
+OUT = sys.argv[2] if len(sys.argv) > 2 else "models/tableformer"
+os.makedirs(OUT, exist_ok=True)
+
+cfg = json.load(open(f"{ART}/tm_config.json"))
+cfg["model"]["save_dir"] = ART
+cfg["predict"]["profiling"] = False
+
+from docling_ibm_models.tableformer.data_management.tf_predictor import TFPredictor  # noqa: E402
+
+pred = TFPredictor(cfg, device="cpu")
+m = pred._model
+m.eval()
+torch.set_grad_enabled(False)
+for p in m.parameters():
+    p.requires_grad_(False)
+tt = m._tag_transformer
+nh = tt._n_heads
+word_map = pred._init_data["word_map"]["word_map_tag"]
+start = word_map["<start>"]
+
+
+class Encode(nn.Module):
+    def forward(self, img):
+        eo = m._encoder(img)
+        eo = tt._input_filter(eo.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
+        ei = eo.reshape(1, -1, eo.size(-1)).permute(1, 0, 2)
+        pos = ei.shape[0]
+        return tt._encoder(ei, mask=torch.zeros((nh, pos, pos), dtype=torch.bool))
+
+
+class Decode(nn.Module):
+    def forward(self, tags, memory):
+        emb = tt._positional_encoding(tt._embedding(tags))
+        decoded, _ = tt._decoder(emb, memory, None, memory_key_padding_mask=None)
+        last = decoded[-1]
+        return tt._fc(last), last
+
+
+def check(name, a, b):
+    import numpy as np
+
+    d = float(np.abs(a - b).max())
+    print(f"  {name}: shape {tuple(a.shape)} | max|onnx-torch| = {d:.2e}")
+    return d
+
+
+img = torch.randn(1, 3, 448, 448)
+with torch.no_grad():
+    mem = Encode()(img)
+torch.onnx.export(
+    Encode(), (img,), f"{OUT}/encoder.onnx",
+    input_names=["image"], output_names=["memory"], opset_version=17, dynamo=False,
+)
+tags = torch.full((3, 1), start, dtype=torch.long)
+with torch.no_grad():
+    logits, hidden = Decode()(tags, mem)
+torch.onnx.export(
+    Decode(), (tags, mem), f"{OUT}/decoder.onnx",
+    input_names=["tags", "memory"], output_names=["logits", "hidden"],
+    dynamic_axes={"tags": {0: "seq"}}, opset_version=17, dynamo=False,
+)
+
+import onnxruntime as ort  # noqa: E402
+
+print("encoder.onnx:")
+eo = ort.InferenceSession(f"{OUT}/encoder.onnx").run(None, {"image": img.numpy()})[0]
+check("memory", eo, mem.numpy())
+print("decoder.onnx:")
+do = ort.InferenceSession(f"{OUT}/decoder.onnx").run(
+    None, {"tags": tags.numpy(), "memory": mem.numpy()}
+)
+check("logits", do[0], logits.numpy())
+check("hidden", do[1], hidden.numpy())
+
+# word map → tokens file for the Rust decode loop
+json.dump(
+    {"word_map_tag": word_map, "start": start, "end": word_map["<end>"]},
+    open(f"{OUT}/wordmap.json", "w"),
+)
+print("wrote wordmap.json; OTSL vocab size:", len(word_map))