docling_pdf/

tableformer.rs

//! TableFormer: table-structure recovery via docling-ibm-models, exported to
//! ONNX by `scripts/install/export_tableformer.py`. The image encoder + tag-transformer
//! encoder run once to a memory tensor; the decoder is then stepped
//! autoregressively to emit an OTSL structure-token sequence (the same model
//! docling runs). See docs/PDF_CONFORMANCE.md.

use crate::pdfium_backend::TextCell;
// The ONNX-free half (preprocessing, structure corrections, bbox bookkeeping,
// span merge, OTSL→grid) lives in tf_core so the browser build (#157 stage 3)
// runs the same logic; this file owns the three `ort` sessions and the
// owned-value KV-cache fast path.
use crate::tf_core::{
    argmax, build_table_cells, correct, merge_spans, preprocess_input, BboxBook, TableCell, END,
    MAX_STEPS, START, UCEL,
};
use image::RgbImage;
use ort::session::Session;
use ort::value::{DynValue, Tensor};

const SIDE: usize = crate::tf_core::SIDE as usize;
const EMBED_DIM: usize = crate::tf_core::EMBED_DIM;
/// Decoder geometry, fixed by the exported TableModel04_rs graph: the cached
/// decoder threads a `[N_LAYERS, past, 1, EMBED_DIM]` per-layer state cache.
const N_LAYERS: usize = 6;

pub struct TableFormer {
    encoder: Session,
    decoder: Session,
    bbox: Session,
    /// Which decoder graph flavour is loaded, detected from the session's
    /// input names (so an explicit `DOCLING_TABLEFORMER_DECODER` override
    /// works with any of them).
    style: DecoderStyle,
}

/// The three decoder-graph generations the loop supports.
#[derive(Clone, Copy, PartialEq, Eq)]
enum DecoderStyle {
    /// `decoder.onnx`: layer-output cache; feeds the full `tags` prefix and a
    /// single `cache` every step.
    Legacy,
    /// The pre-#97 `decoder_kv.onnx`: one tag per step, `cache_k`/`cache_v`,
    /// with the stacked `cross_k`/`cross_v` re-split inside every step.
    KvStacked,
    /// The #97 `decoder_kv.onnx`: one tag per step, and the constant cross
    /// tensors arrive as 2×`N_LAYERS` per-layer inputs (`cross_kt_i` already
    /// transposed for q·Kᵀ, `cross_v_i`), computed once per table by the
    /// encoder — the step graph does no work proportional to their size.
    KvHoisted,
}

/// KV-cache geometry fixed by the `decoder_kv.onnx` export
/// (`[N_LAYERS, 1, KV_HEADS, past, KV_HEAD_DIM]`, `KV_HEADS × KV_HEAD_DIM = EMBED_DIM`).
const KV_HEADS: usize = 8;
const KV_HEAD_DIM: usize = 64;

/// The autoregressive decode state: `a` is the legacy layer-output cache, or
/// `cache_k` for the KV graph; `b` is `cache_v` (KV graph only). `None` = first
/// step (the zero-`past` empties are allocated per table by [`TableFormer::empty_cache`]).
#[derive(Default)]
struct DecodeCache {
    a: Option<DynValue>,
    b: Option<DynValue>,
}

/// Zero-`past` first-step cache tensors: `(cache, None)` for the legacy graph,
/// `(cache_k, Some(cache_v))` for the KV graph.
type EmptyCache = (Tensor<f32>, Option<Tensor<f32>>);

/// Encoder outputs that drive the cached decode loop: the per-layer cross-attention
/// K/V (projected from the image memory once, constant across decode steps) and
/// `enc_out` for the bbox decoder. Kept as owned `ort` values so each decode step
/// (and the bbox run) borrows them directly — no per-step extract/copy/re-wrap.
struct EncodeOut {
    ck: DynValue,
    cv: DynValue,
    eo: DynValue,
    /// `KvHoisted` only: per-layer `[cross_kt_0..N, cross_v_0..N]`, index-aligned
    /// with the decoder's input names, borrowed by every decode step.
    per_layer: Vec<(String, DynValue)>,
}

impl TableFormer {
    /// Load the exported encoder/decoder/bbox ONNX graphs (env overrides, else
    /// `models/tableformer/{encoder,decoder,bbox}.onnx`). Returns `None` if any is
    /// absent, so the pipeline falls back to geometric reconstruction.
    pub fn load() -> Option<Self> {
        Self::load_with(crate::intra_threads())
    }

    /// Like [`load`](Self::load) but with an explicit intra-op thread count, so a
    /// parallel page-worker pool can run each table model on fewer threads (the
    /// throughput comes from running pages concurrently, not from one fat model).
    pub fn load_with(intra: usize) -> Option<Self> {
        let enc = std::env::var("DOCLING_TABLEFORMER_ENCODER")
            .unwrap_or_else(|_| crate::resolve_asset("models/tableformer/encoder.onnx"));
        // Decoder preference (explicit override wins): INT8 variants first
        // unless DOCLING_RS_FP32 opts out; within a precision the true-KV-cache
        // export (`decoder_kv*`, one token per step, O(past) step cost) ranks
        // ahead of the legacy layer-output-cache graph it matches byte-for-byte
        // (91/91 snapshot corpus exact with either). Re-measured warm on the
        // corpus fixtures: the KV graph is ~13% faster on ordinary tables
        // (2206.01062) and ~17% on the huge-table page (2305.03393v1-pg9),
        // for +36 MB on disk — table-heavy single-page PDFs are exactly where
        // the pipeline is tightest against Python docling, so speed wins the
        // default and the legacy file stays as the smaller fallback.
        let dec = std::env::var("DOCLING_TABLEFORMER_DECODER").unwrap_or_else(|_| {
            let candidates: &[&str] = if crate::prefer_fp32() {
                &[
                    "models/tableformer/decoder_kv.onnx",
                    "models/tableformer/decoder.onnx",
                ]
            } else {
                // decoder_kv ranks ABOVE decoder_int8: the #97 hoisted fp32 KV
                // graph is faster than the quantized legacy graph on every
                // machine measured, and it is byte-exact (its own int8 variant
                // is not produced — see quantize_models.py).
                &[
                    "models/tableformer/decoder_kv_int8.onnx",
                    "models/tableformer/decoder_kv.onnx",
                    "models/tableformer/decoder_int8.onnx",
                    "models/tableformer/decoder.onnx",
                ]
            };
            candidates
                .iter()
                .map(|p| crate::resolve_asset(p))
                .find(|p| std::path::Path::new(p).exists())
                .unwrap_or_else(|| "models/tableformer/decoder.onnx".to_string())
        });
        let bbx = std::env::var("DOCLING_TABLEFORMER_BBOX")
            .unwrap_or_else(|_| crate::resolve_asset("models/tableformer/bbox.onnx"));
        if crate::timing::enabled() {
            eprintln!("docling-pdf: tableformer decoder: {dec}");
        }
        if [&enc, &dec, &bbx]
            .iter()
            .any(|p| !std::path::Path::new(p).exists())
        {
            // The geometric fallback is a supported, intentional configuration
            // (docling has no ML table-structure equivalent baked in either), so
            // this stays a single quiet stderr note rather than an error — but it
            // fires every process (not per-worker) so a CWD-relative default that
            // silently misses its files (a very easy mistake for anything not run
            // from the repo root, e.g. an embedding app) is at least visible once.
            warn_missing_once(&enc, &dec, &bbx);
            return None;
        }
        // The decoder's KV-cache grows by one entry every autoregressive step, so
        // its input shapes differ on every `run()` call. ONNX Runtime's memory
        // pattern optimizer assumes stable shapes to plan buffer reuse; disabling
        // it for this session avoids repeatedly re-validating/re-touching that
        // plan (and the external-weights file) on each step.
        let build = |path: &str, mem_pattern: bool| -> Result<Session, String> {
            let builder = Session::builder()
                .map_err(|e| e.to_string())?
                .with_intra_threads(intra)
                .map_err(|e| e.to_string())?
                .with_memory_pattern(mem_pattern)
                .map_err(|e| e.to_string())?;
            crate::ep::apply(builder)?
                .commit_from_file(path)
                .map_err(|e| format!("tableformer load {path}: {e}"))
        };
        match (build(&enc, true), build(&dec, false), build(&bbx, true)) {
            (Ok(encoder), Ok(decoder), Ok(bbox)) => {
                let has = |n: &str| decoder.inputs().iter().any(|i| i.name() == n);
                let style = if has("cross_kt_0") {
                    DecoderStyle::KvHoisted
                } else if has("cache_k") {
                    DecoderStyle::KvStacked
                } else {
                    DecoderStyle::Legacy
                };
                if style == DecoderStyle::KvHoisted
                    && !encoder.outputs().iter().any(|o| o.name() == "cross_kt_0")
                {
                    eprintln!(
                        "docling-pdf: tableformer decoder needs per-layer cross tensors \
                         (cross_kt_*) the encoder doesn't emit — re-download or re-export \
                         the model set (scripts/install/export_tableformer.py); \
                         falling back to geometric tables"
                    );
                    return None;
                }
                Some(Self {
                    encoder,
                    decoder,
                    bbox,
                    style,
                })
            }
            _ => None,
        }
    }

    /// Run the image encoder and capture what the cached decoder loop needs: each
    /// decoder layer's cross-attention K/V (projected from the image memory once,
    /// shape `[N_LAYERS,1,H,S,head_dim]`) and `enc_out` for the bbox decoder.
    fn encode(&mut self, img: &RgbImage) -> Result<EncodeOut, String> {
        let input = preprocess(img)?;
        let mut enc_out = self
            .encoder
            .run(ort::inputs!["image" => input])
            .map_err(|e| format!("tableformer: encode: {e}"))?;
        let mut per_layer = Vec::new();
        if self.style == DecoderStyle::KvHoisted {
            for prefix in ["cross_kt_", "cross_v_"] {
                for i in 0.. {
                    let name = format!("{prefix}{i}");
                    match enc_out.remove(&name) {
                        Some(v) => per_layer.push((name, v)),
                        None => break,
                    }
                }
            }
            if per_layer.is_empty() {
                return Err("tableformer: encoder emitted no cross_kt_* outputs".into());
            }
        }
        let mut grab = |name: &str| -> Result<DynValue, String> {
            enc_out
                .remove(name)
                .ok_or_else(|| format!("tableformer: encoder output {name} missing"))
        };
        Ok(EncodeOut {
            ck: grab("cross_k")?,
            cv: grab("cross_v")?,
            eo: grab("enc_out")?,
            per_layer,
        })
    }

    /// One doubly-cached decode step: feed the current `tags`, the constant cross
    /// K/V, and the growing self-attention `cache`; return the raw argmax tag and
    /// the last token's hidden state, advancing the cache. The cache stays an owned
    /// `ort` value — the previous step's `out_cache` output is fed back directly,
    /// never extracted or copied (it grows every step, so per-step copies were
    /// O(steps²) float traffic). `empty_cache` is the zero-`past` value used on the
    /// first step (ort's array constructors reject a 0-length dim, so it is
    /// allocated through the session allocator by the caller).
    fn decode_step(
        &mut self,
        tags: &[i64],
        enc: &EncodeOut,
        cache: &mut DecodeCache,
        empty: &EmptyCache,
    ) -> Result<(i64, Vec<f32>), String> {
        crate::timing::timed("tf.decode_step", || {
            self.decode_step_inner(tags, enc, cache, empty)
        })
    }

    fn decode_step_inner(
        &mut self,
        tags: &[i64],
        enc: &EncodeOut,
        cache: &mut DecodeCache,
        empty: &EmptyCache,
    ) -> Result<(i64, Vec<f32>), String> {
        let mut dout = match self.style {
            DecoderStyle::KvHoisted => {
                // #97 graph: one tag; the constant per-layer cross tensors are
                // borrowed views — the step pays nothing proportional to them.
                let last = *tags.last().expect("decode starts from <start>");
                let tag_t = Tensor::from_array(([1usize, 1usize], vec![last]))
                    .map_err(|e| format!("tableformer: tag: {e}"))?;
                let mut inputs: Vec<(
                    std::borrow::Cow<'_, str>,
                    ort::session::SessionInputValue<'_>,
                )> = Vec::with_capacity(3 + enc.per_layer.len());
                inputs.push(("tag".into(), tag_t.into()));
                match (cache.a.as_ref(), cache.b.as_ref()) {
                    (Some(k), Some(v)) => {
                        inputs.push(("cache_k".into(), k.into()));
                        inputs.push(("cache_v".into(), v.into()));
                    }
                    _ => {
                        inputs.push(("cache_k".into(), (&empty.0).into()));
                        inputs.push((
                            "cache_v".into(),
                            empty
                                .1
                                .as_ref()
                                .expect("kv empty cache has both halves")
                                .into(),
                        ));
                    }
                }
                for (name, v) in &enc.per_layer {
                    inputs.push((name.as_str().into(), v.into()));
                }
                self.decoder.run(inputs)
            }
            DecoderStyle::KvStacked => {
                // Pre-#97 KV graph: feed only the newly emitted tag; the projected
                // K/V for the whole prefix live in cache_k/cache_v and are fed
                // back as-is.
                let last = *tags.last().expect("decode starts from <start>");
                let tag_t = Tensor::from_array(([1usize, 1usize], vec![last]))
                    .map_err(|e| format!("tableformer: tag: {e}"))?;
                match (cache.a.as_ref(), cache.b.as_ref()) {
                    (Some(k), Some(v)) => self.decoder.run(ort::inputs![
                        "tag" => tag_t, "cross_k" => &enc.ck, "cross_v" => &enc.cv,
                        "cache_k" => k, "cache_v" => v]),
                    _ => self.decoder.run(ort::inputs![
                        "tag" => tag_t, "cross_k" => &enc.ck, "cross_v" => &enc.cv,
                        "cache_k" => &empty.0,
                        "cache_v" => empty.1.as_ref().expect("kv empty cache has both halves")]),
                }
            }
            DecoderStyle::Legacy => {
                let tags_t = Tensor::from_array(([tags.len(), 1usize], tags.to_vec()))
                    .map_err(|e| format!("tableformer: tags: {e}"))?;
                match cache.a.as_ref() {
                    None => self.decoder.run(ort::inputs![
                        "tags" => tags_t, "cross_k" => &enc.ck, "cross_v" => &enc.cv,
                        "cache" => &empty.0]),
                    Some(c) => self.decoder.run(ort::inputs![
                        "tags" => tags_t, "cross_k" => &enc.ck, "cross_v" => &enc.cv,
                        "cache" => c]),
                }
            }
        }
        .map_err(|e| format!("tableformer: decode: {e}"))?;
        let (_, logits) = dout["logits"]
            .try_extract_tensor::<f32>()
            .map_err(|e| format!("tableformer: logits: {e}"))?;
        let raw = argmax(logits) as i64;
        let (_, hidden) = dout["hidden"]
            .try_extract_tensor::<f32>()
            .map_err(|e| format!("tableformer: hidden: {e}"))?;
        let hidden = hidden.to_vec();
        if self.style != DecoderStyle::Legacy {
            cache.a = Some(
                dout.remove("out_cache_k")
                    .ok_or_else(|| "tableformer: out_cache_k missing".to_string())?,
            );
            cache.b = Some(
                dout.remove("out_cache_v")
                    .ok_or_else(|| "tableformer: out_cache_v missing".to_string())?,
            );
        } else {
            cache.a = Some(
                dout.remove("out_cache")
                    .ok_or_else(|| "tableformer: decoder output out_cache missing".to_string())?,
            );
        }
        Ok((raw, hidden))
    }

    /// The zero-`past` first-step cache(s), allocated through the session
    /// allocator (ort's array constructors reject a 0-length dim; the C API does
    /// allow it).
    fn empty_cache(&self) -> Result<EmptyCache, String> {
        let alloc = self.decoder.allocator();
        if self.style != DecoderStyle::Legacy {
            let mk = || {
                Tensor::<f32>::new(alloc, [N_LAYERS, 1, KV_HEADS, 0usize, KV_HEAD_DIM])
                    .map_err(|e| format!("tableformer: empty kv cache: {e}"))
            };
            Ok((mk()?, Some(mk()?)))
        } else {
            let c = Tensor::<f32>::new(alloc, [N_LAYERS, 0usize, 1, EMBED_DIM])
                .map_err(|e| format!("tableformer: empty cache: {e}"))?;
            Ok((c, None))
        }
    }

    /// Predict the OTSL structure-token sequence for a table-region image.
    pub fn predict_otsl(&mut self, img: &RgbImage) -> Result<Vec<i64>, String> {
        let enc = self.encode(img)?;
        // Structure corrections live in tf_core::correct (shared with the wasm
        // path); docling's line_num is never incremented, so xcel→lcel fires on
        // every row.
        let mut tags: Vec<i64> = vec![START];
        let mut out: Vec<i64> = Vec::new();
        let mut prev_ucel = false;
        let mut cache = DecodeCache::default();
        let empty = self.empty_cache()?;
        while out.len() < MAX_STEPS {
            let (raw, _hidden) = self.decode_step(&tags, &enc, &mut cache, &empty)?;
            let tag = correct(raw, prev_ucel);
            if tag == END {
                break;
            }
            out.push(tag);
            tags.push(tag);
            prev_ucel = tag == UCEL;
        }
        Ok(out)
    }

    /// Full structure prediction: OTSL grid cells with per-cell boxes (in the 448
    /// image, normalized cxcywh). Collects per-cell decoder hidden states using
    /// docling's exact bbox bookkeeping (skip-after-row-break, first-lcel of a
    /// horizontal span), runs the bbox decoder, merges span boxes, then lays the
    /// cells onto the OTSL grid with row/col spans.
    pub fn predict_table_structure(&mut self, img: &RgbImage) -> Result<Vec<TableCell>, String> {
        let enc = self.encode(img)?;

        // The autoregressive loop's bbox bookkeeping lives in tf_core::BboxBook
        // (shared with the wasm path); this loop only steps the decoder.
        let mut book = BboxBook::new();
        let mut cache = DecodeCache::default();
        let empty = self.empty_cache()?;
        while book.otsl.len() < MAX_STEPS {
            let (raw, hidden) = self.decode_step(&book.tags, &enc, &mut cache, &empty)?;
            if !book.step(raw, &hidden) {
                break;
            }
        }
        if book.n == 0 {
            return Ok(Vec::new());
        }
        let tag_h = Tensor::from_array(([book.n, EMBED_DIM], std::mem::take(&mut book.hiddens)))
            .map_err(|e| format!("tableformer: tag_h: {e}"))?;
        let bout = self
            .bbox
            .run(ort::inputs!["enc_out" => &enc.eo, "tag_h" => tag_h])
            .map_err(|e| format!("tableformer: bbox: {e}"))?;
        let (_, raw) = bout["boxes"]
            .try_extract_tensor::<f32>()
            .map_err(|e| format!("tableformer: boxes: {e}"))?;
        let boxes: Vec<[f32; 4]> = raw
            .chunks_exact(4)
            .map(|c| [c[0], c[1], c[2], c[3]])
            .collect();
        // Per-cell class logits [n, 3] → argmax (docling's `outputs_class`).
        let (_, craw) = bout["classes"]
            .try_extract_tensor::<f32>()
            .map_err(|e| format!("tableformer: classes: {e}"))?;
        let classes: Vec<i64> = craw.chunks_exact(3).map(|c| argmax(c) as i64).collect();
        let (merged, merged_classes) = merge_spans(&boxes, &classes, &book.merge);
        Ok(build_table_cells(&book.otsl, &merged, &merged_classes))
    }

    /// Predict a table region's Markdown grid: crop the region (docling's
    /// page→1024px box-average then bbox crop), run the structure model, then
    /// match the page's word cells into the predicted cells with docling's
    /// matching post-processor ([`crate::tf_match`]) and expand spans into a
    /// dense `rows × cols` grid. `region` is `(l, t, r, b)` in page points
    /// (top-left). Returns `None` if no structure is predicted.
    pub fn predict_table_rows(
        &mut self,
        page_image: &RgbImage,
        region: [f32; 4],
        words: &[TextCell],
    ) -> Option<Vec<Vec<String>>> {
        // page → 1024px height (cv2.INTER_AREA), then crop the table bbox.
        // docling's coordinate chain, rounding included: the cluster bbox is
        // rounded to integer page points *first* (`round(cluster.bbox.l) *
        // scale`, banker's rounding), scaled by 2 (its table-structure page
        // scale), then by `1024 / <2x page-image height>`, and the crop indices
        // round again. Rounding after scaling instead shifts some crops by a
        // pixel — enough to change TableFormer's cell boxes on tall tables
        // (redp5110's TOC).
        let sf = 1024.0 / page_image.height() as f32;
        let pw = (page_image.width() as f32 * sf) as u32;
        let page1024 = crate::timing::timed("tableformer.inter_area", || {
            crate::resample::inter_area(page_image, pw, 1024)
        });
        let k = 2.0 * 1024.0 / page_image.height() as f64;
        let px = |v: f32| (v as f64).round_ties_even() * k;
        let x = (px(region[0]).round_ties_even()).max(0.0) as u32;
        let y = (px(region[1]).round_ties_even()).max(0.0) as u32;
        let x2 = (px(region[2]).round_ties_even() as u32).min(page1024.width());
        let y2 = (px(region[3]).round_ties_even() as u32).min(page1024.height());
        if x2 <= x || y2 <= y {
            return None;
        }
        let crop = image::imageops::crop_imm(&page1024, x, y, x2 - x, y2 - y).to_image();
        let cells = crate::timing::timed("tableformer.structure", || {
            self.predict_table_structure(&crop)
        })
        .ok()?;
        if cells.is_empty() {
            return None;
        }
        // The ort-free tail (word matching + grid assembly) is shared with the
        // browser path in tf_core.
        crate::tf_core::table_rows(&cells, region, words)
    }
}

/// Note once per process that TableFormer's ONNX graphs weren't found, so tables
/// fall back to geometric reconstruction. The default paths are relative
/// (`models/tableformer/*.onnx`), which only resolves when the process's current
/// directory happens to be the repo root — a very easy miss for anything else
/// (an embedding app, a binding invoked from a different working directory, …),
/// and previously failed with no signal at all.
fn warn_missing_once(enc: &str, dec: &str, bbx: &str) {
    static WARNED: std::sync::Once = std::sync::Once::new();
    WARNED.call_once(|| {
        eprintln!(
            "docling.rs: TableFormer models not found (checked {enc}, {dec}, {bbx}); \
             tables will use geometric reconstruction instead of ML table-structure \
             recognition. Set DOCLING_TABLEFORMER_ENCODER / DOCLING_TABLEFORMER_DECODER \
             / DOCLING_TABLEFORMER_BBOX to enable it (see README.md)."
        );
    });
}

/// docling's preprocessing: bilinear (cv2.INTER_LINEAR) resize the crop to 448²,
/// normalize `(x/255 − mean)/std`, laid out as (C, W, H) — docling transposes
/// (2,1,0), so width is the major spatial axis. The page→1024px box-average
/// (cv2.INTER_AREA) is the caller's job.
fn preprocess(img: &RgbImage) -> Result<Tensor<f32>, String> {
    Tensor::from_array(([1usize, 3, SIDE, SIDE], preprocess_input(img)))
        .map_err(|e| format!("tableformer: input: {e}"))
}