#157 stage 3 (phase 4): TableFormer demo wiring

Add the browser TableFormer profile to the scan demo. A "TableFormer
tables" checkbox opts in; when checked, pipeline.js lazily loads the three
graphs same-origin from ./models/tableformer/ (encoder.onnx self-contained
~225 MB; decoder_kv.onnx + bbox.onnx with their .onnx.data sidecars via
ort-web's externalData option) and builds JsTfSession — the stateful
session the wasm TfSession interop expects: encode stashes the constant
cross K/V + enc_out and resets the KV-cache, step runs one decoder round
(int64 tag, empty first-step cache [6,1,8,0,64]) and returns logits+hidden,
bbox runs the box/class head. Table regions then go through
ScannedConverter.addPageTf; the lite geometric path stays the default so
the big models aren't fetched unless asked. Models load lazily on the
first table conversion with a progress line.

Refs #157

Signed-off-by: artiz <artem.kustikov@gmail.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EY5KAiquN4YpVf2PXEQkVT

Author: artiz

crates/docling-wasm/www/pipeline.js

@@ -38,6 +38,60 @@ const REC_MODELS = {
   },
 };
 
+// TableFormer graphs, loaded same-origin from ./models/tableformer/ (the same
+// files download_dependencies.sh fetches). The encoder is self-contained; the
+// decoder/bbox use ONNX external data, so their .onnx.data sidecar rides along
+// via ort-web's externalData option (path = the location stored in the .onnx).
+const TF_DIR = "./models/tableformer/";
+
+/// The stateful session the wasm TfSession interop expects (see
+/// src/tableformer.rs): `encode` runs the image encoder once and stashes the
+/// constant cross-attention K/V + enc_out and resets the KV-cache; `step` runs
+/// one decoder step feeding the stored cross + growing cache; `bbox` runs the
+/// bbox decoder. The heavy tensors stay here — only tags and logits/hidden
+/// cross the wasm boundary. Geometry: 6 layers, 8 KV heads, head_dim 64.
+class JsTfSession {
+  constructor(enc, dec, bbox) {
+    this.enc = enc;
+    this.dec = dec;
+    this.bboxSess = bbox;
+    this.cross = null;
+    this.encOut = null;
+    this.cacheK = null;
+    this.cacheV = null;
+  }
+  async encode(image) {
+    const out = await this.enc.run({ image: new ort.Tensor("float32", image, [1, 3, 448, 448]) });
+    this.cross = {};
+    for (let i = 0; i < 6; i++) {
+      this.cross["cross_kt_" + i] = out["cross_kt_" + i];
+      this.cross["cross_v_" + i] = out["cross_v_" + i];
+    }
+    this.encOut = out["enc_out"];
+    // Empty first-step KV-cache: [N_LAYERS, 1, KV_HEADS, past=0, head_dim].
+    this.cacheK = new ort.Tensor("float32", new Float32Array(0), [6, 1, 8, 0, 64]);
+    this.cacheV = new ort.Tensor("float32", new Float32Array(0), [6, 1, 8, 0, 64]);
+  }
+  async step(tag) {
+    const out = await this.dec.run({
+      tag: new ort.Tensor("int64", BigInt64Array.from([BigInt(tag)]), [1, 1]),
+      cache_k: this.cacheK,
+      cache_v: this.cacheV,
+      ...this.cross,
+    });
+    this.cacheK = out["out_cache_k"];
+    this.cacheV = out["out_cache_v"];
+    return { logits: out["logits"].data, hidden: out["hidden"].data };
+  }
+  async bbox(tagH, n) {
+    const out = await this.bboxSess.run({
+      enc_out: this.encOut,
+      tag_h: new ort.Tensor("float32", tagH, [n, 512]),
+    });
+    return { boxes: out["boxes"].data, classes: out["classes"].data };
+  }
+}
+
 export function createOcr({ onStatus }) {
   const status = (msg, spinning = true) => onStatus && onStatus(msg, spinning);
 
@@ -141,15 +195,42 @@ export function createOcr({ onStatus }) {
     }
   }
 
+  // TableFormer sessions, loaded lazily on first use (the encoder alone is
+  // ~225 MB, so this only downloads when the table profile is chosen).
+  let tf = null;
+  async function ensureTf() {
+    if (tf) return tf;
+    const load = async (name, external) => {
+      const model = await fetchProgress(TF_DIR + name + ".onnx", `tableformer ${name}`);
+      const opts = { executionProviders: ["wasm"], logSeverityLevel: 3 };
+      if (external) {
+        const data = await fetchProgress(TF_DIR + name + ".onnx.data", `tableformer ${name} data`);
+        opts.externalData = [{ path: name + ".onnx.data", data: new Uint8Array(data) }];
+      }
+      return ort.InferenceSession.create(model, opts);
+    };
+    const enc = await load("encoder", false);
+    const dec = await load("decoder_kv", true);
+    const bbox = await load("bbox", true);
+    status("starting tableformer sessions …", true);
+    tf = new JsTfSession(enc, dec, bbox);
+    return tf;
+  }
+
   // Multi-page document lifecycle: startDoc → addPage* → finishDoc. Pages come
   // in as RGBA (already rasterized on the main thread), one document at a time.
   let cur = null;
-  async function startDoc(lang) {
+  async function startDoc(lang, useTf) {
     const { dict, rec } = await recFor(lang);
-    cur = { conv: new ScannedConverter(dict), rec };
+    const tfSess = useTf ? await ensureTf() : null;
+    cur = { conv: new ScannedConverter(dict), rec, tf: tfSess };
   }
   async function addPage(rgba, w, h, scale) {
-    await cur.conv.add_page(rgba, w, h, scale, layout, cur.rec);
+    if (cur.tf) {
+      await cur.conv.addPageTf(rgba, w, h, scale, layout, cur.rec, cur.tf);
+    } else {
+      await cur.conv.add_page(rgba, w, h, scale, layout, cur.rec);
+    }
   }
   function finishDoc(name) {
     const md = cur.conv.finish(name, "md");

crates/docling-wasm/www/scan.html

@@ -84,6 +84,10 @@ <h1>docling.rs → wasm: scanned PDF / image, full lite profile</h1>
       <option value="cyrillic">Cyrillic + Latin (cyrillic_PP-OCRv5)</option>
       <option value="ch">Chinese + Latin (ch_PP-OCRv3)</option>
     </select>
+    <label style="margin-left:1rem" title="Stage 3: real table structure via the TableFormer ONNX graphs, instead of the geometric reconstruction. Fetches ~380 MB of models (encoder/decoder/bbox) on first use, same-origin from ./models/tableformer/.">
+      <input type="checkbox" id="tf" disabled> TableFormer tables
+      <span style="color:#888">(+~380 MB models, PDFs)</span>
+    </label>
     <span id="meta"></span>
   </p>
 
@@ -118,6 +122,7 @@ <h1>docling.rs → wasm: scanned PDF / image, full lite profile</h1>
     const meta = document.getElementById("meta");
     const out = document.getElementById("out");
     const lang = document.getElementById("lang");
+    const tf = document.getElementById("tf");
     const drop = document.getElementById("drop");
     const pick = document.getElementById("pick");
 
@@ -128,12 +133,12 @@ <h1>docling.rs → wasm: scanned PDF / image, full lite profile</h1>
     }
     function busy(msg) {
       document.body.classList.add("busy");
-      pick.disabled = lang.disabled = true;
+      pick.disabled = lang.disabled = tf.disabled = true;
       status(msg, true);
     }
     function ready(msg) {
       document.body.classList.remove("busy");
-      pick.disabled = lang.disabled = false;
+      pick.disabled = lang.disabled = tf.disabled = false;
       status(msg, false);
     }
     const onStatus = (msg, spinning) => {
@@ -172,7 +177,7 @@ <h1>docling.rs → wasm: scanned PDF / image, full lite profile</h1>
         worker,
         boot: (l) => rpc("boot", { lang: l }),
         rec: (l) => rpc("rec", { lang: l }),
-        docStart: (l) => rpc("doc-start", { lang: l }),
+        docStart: (l, useTf) => rpc("doc-start", { lang: l, useTf }),
         addPage: (rgba, w, h, scale) => rpc("doc-page", { rgba, w, h, scale }, [rgba]),
         docFinish: (name) => rpc("doc-finish", { name }).then((r) => r.md),
         convertImage: (bytes, name, l) => rpc("convert-image", { bytes, name, lang: l }, [bytes]).then((r) => r.md),
@@ -190,7 +195,7 @@ <h1>docling.rs → wasm: scanned PDF / image, full lite profile</h1>
           return { kind, threads: THREADS, main: true };
         },
         rec: async (l) => { await ocr.recFor(l); await ocr.warmup(l); return {}; },
-        docStart: (l) => ocr.startDoc(l),
+        docStart: (l, useTf) => ocr.startDoc(l, useTf),
         addPage: (rgba, w, h, scale) => ocr.addPage(new Uint8Array(rgba), w, h, scale),
         docFinish: (name) => Promise.resolve(ocr.finishDoc(name)),
         convertImage: (bytes, name, l) => ocr.convertImage(bytes, name, l),
@@ -205,7 +210,7 @@ <h1>docling.rs → wasm: scanned PDF / image, full lite profile</h1>
     async function handlePdf(file) {
       const buf = await file.arrayBuffer();
       const pdf = await pdfjs.getDocument({ data: buf, ...PDFJS_ASSETS }).promise;
-      await drv.docStart(lang.value);
+      await drv.docStart(lang.value, tf.checked);
       const canvas = document.createElement("canvas");
       const ctx = canvas.getContext("2d", { willReadFrequently: true });
       async function render(p) {

crates/docling-wasm/www/worker.js

@@ -32,7 +32,7 @@ async function handle(m) {
       await ocr.warmup(m.lang);
       return {};
     case "doc-start":
-      await ocr.startDoc(m.lang);
+      await ocr.startDoc(m.lang, m.useTf);
       return {};
     case "doc-page":
       await ocr.addPage(new Uint8Array(m.rgba), m.w, m.h, m.scale);