diff --git a/.github/workflows/publish-models.yml b/.github/workflows/publish-models.yml index 75bc57fd..48800fa5 100644 --- a/.github/workflows/publish-models.yml +++ b/.github/workflows/publish-models.yml @@ -5,10 +5,12 @@ # # The layout model and TableFormer are PyTorch→ONNX *format conversions* of # docling-project's own models (Apache-2.0 / CDLA-Permissive-2.0 — see -# MODELS_NOTICE.md); no weights are retrained or modified here. pdfium and the -# OCR model are re-hosted third-party binaries, unmodified, from their own -# public releases — mirrored here only so `download_dependencies.sh` needs one -# host instead of three. +# MODELS_NOTICE.md); no weights are retrained here. The `*_int8` assets are +# post-training quantizations of those same exports (scripts/quantize_models.py +# — a numeric re-encoding of the weights, no retraining; conformance-validated +# in PDF_PERFORMANCE.md). pdfium and the OCR model are re-hosted third-party +# binaries, unmodified, from their own public releases — mirrored here only so +# `download_dependencies.sh` needs one host instead of three. # # Trigger: manual only (workflow_dispatch) — it installs torch and exports two # models, several GB and 10-20 minutes, not something to run on every push. @@ -65,6 +67,26 @@ jobs: id: tf-export run: python scripts/export_tableformer.py models/tableformer + # --- INT8 quantization (optional — fp32 assets publish regardless) --- + # Conv-only static INT8 of the layout model (calibrated on the repo's PDF + # corpus, checked out above) + dynamic INT8 of the TableFormer decoder. + # Validated conformance-neutral; see PDF_PERFORMANCE.md. Two separate + # steps so a TableFormer export failure doesn't block the layout int8. + - name: Install quantization deps + continue-on-error: true + id: quant-deps + run: pip install onnx onnxruntime sympy pypdfium2 pillow numpy + + - name: Quantize layout model (static INT8, conv-only) + if: steps.quant-deps.outcome == 'success' + continue-on-error: true + run: python scripts/quantize_models.py layout + + - name: Quantize TableFormer decoder (dynamic INT8) + if: steps.quant-deps.outcome == 'success' && steps.tf-export.outcome == 'success' + continue-on-error: true + run: python scripts/quantize_models.py tableformer-decoder + # --- pdfium + OCR (re-hosted third-party binaries, not exports) ----- - name: Fetch pdfium + OCR assets run: | @@ -97,8 +119,10 @@ jobs: stage third-party/ocr_rec.onnx ocr_rec.onnx stage third-party/ppocr_keys_v1.txt ppocr_keys_v1.txt stage models/layout_heron.onnx layout_heron.onnx + stage models/layout_heron_int8.onnx layout_heron_int8.onnx stage models/tableformer/encoder.onnx encoder.onnx stage models/tableformer/decoder.onnx decoder.onnx + stage models/tableformer/decoder_int8.onnx decoder_int8.onnx stage models/tableformer/bbox.onnx bbox.onnx echo "staged:" ls -la release-assets/ diff --git a/PDF_PERFORMANCE.md b/PDF_PERFORMANCE.md new file mode 100644 index 00000000..6af7aa57 --- /dev/null +++ b/PDF_PERFORMANCE.md @@ -0,0 +1,207 @@ +# PDF pipeline — performance review & profiling notes + +Post-migration review of the PDF processing path: where the time actually goes, +what was measured, which optimizations are validated, and a ranked backlog of +further ideas that do **not** trade away output quality. + +Measured on a 4-core AVX-512(+VNNI/AMX) Xeon, release build (`lto = "thin"`), +models from `scripts/download_dependencies.sh`, `FLEISCHWOLF_TIMING=1`. + +## Where the time goes + +Per-stage wall-clock share (summed across workers): + +| Stage | 1913-page text-heavy PDF¹ | 16-page table-heavy paper² | scanned page³ | +|---|---:|---:|---:| +| `layout.predict` (RT-DETR ONNX) | **80.3%** | 55.4% | 64.9% | +| `image.resize` (3×→2× CatmullRom) | 14.9% | 7.9% | 18.5% | +| `tableformer` | 2.8% | 32.1% | — | +| `pdfium.render` | 1.8% | 3.7% | 16.5% | +| `textparse` + assembly | ~0.2% | ~0.3% | ~0.1% | + +¹ `tests/data/pdf/large/dotnet-csharp-language-reference.pdf` — 936 s wall, ~0.49 s/page. +² `tests/data/pdf/sources/2203.01017v2.pdf`. +³ `tests/data/scanned/sources/ocr_test.pdf`. + +Two conclusions drive everything below: + +1. **ONNX inference is ~85–95% of PDF conversion time.** All the Rust-side text + extraction, parsing, and assembly work combined is under 1%. Rust-code + micro-optimizations are irrelevant to PDF throughput until the models get + faster; model-level and preprocessing-level changes are the only levers that + matter. +2. Within TableFormer, the **autoregressive decode loop** dominates + (`tableformer.structure` ≈ 96% of the stage; the per-table page resample + `tableformer.inter_area` is ~1% of a conversion). + +The worker-pool topology heuristic in `lib.rs` (`workers × intra ≈ cores`, +default 2×2 on 4 cores) was re-validated: 2×2 beat both 4×1 and 1×4 on the +16-page document (11.6 s vs 12.2 s vs 15.6 s). + +## Validated win: INT8 quantization (quality-checked) + +`scripts/quantize_models.py` produces two quantized models. Point +`DOCLING_LAYOUT_ONNX` / `DOCLING_TABLEFORMER_DECODER` at them to opt in. + +**These are now the default:** when the `*_int8` files sit next to the fp32 +models at the default paths, the pipeline loads them automatically. +`FLEISCHWOLF_FP32=1` forces full precision, and an explicit +`DOCLING_LAYOUT_ONNX` / `DOCLING_TABLEFORMER_DECODER` always wins (the +conformance/groundtruth scripts pin fp32 explicitly, so snapshots stay +deterministic). + +### Layout: static QDQ INT8, **Conv ops only** (~2.4× faster layout) + +Calibrated on 42 real corpus pages preprocessed exactly like +`layout.rs::predict`. Only the HGNetv2 backbone convolutions are quantized; +the transformer decoder and detection-head MatMuls stay fp32. + +| Configuration | layout.predict (16-page doc) | end-to-end wall | model size | +|---|---:|---:|---:| +| fp32 baseline | 17.2 s | 16.6 s | 172 MB | +| **INT8 conv-only** | **7.2 s (2.4×)** | 11.5 s (1.45×) | 68 MB | +| + INT8 TableFormer decoder | — | **12.3 s → see note** | — | + +On text-dominated documents (layout = 80% of time) the end-to-end gain +approaches ~1.7–2×; on table-heavy ones it is ~1.4×. + +Full-scale run — the 1913-page `dotnet-csharp-language-reference.pdf`, +INT8 layout + INT8 TableFormer decoder vs fp32, same machine and binary, +back-to-back: + +| | fp32 | INT8 | ratio | +|---|---:|---:|---:| +| wall clock | 1406 s (0.74 s/page) | **770 s (0.40 s/page)** | **1.83×** | +| `layout.predict` (summed) | 2667 s | 1350 s | 1.98× | +| output difference | — | 1199 of 52,615 Markdown lines (2.3%) | | + +The 2.3% of differing lines are the same near-threshold classification flips +seen on the corpus (where groundtruth conformance measured *equal or slightly +better* under INT8 — 812 vs 833 summed diff-lines), not a systematic +degradation. With layout halved, `image.resize` becomes the next stage +(24.8% of the INT8 run), which is why backlog item 4 matters more after +quantization. + +**Quality gate.** Markdown diffed across the full PDF+scanned corpus (23 files): + +- Conv-only INT8: 12/23 byte-identical to fp32; remaining diffs are small + region-classification flips. Against the committed groundtruth the summed + diff-line distance is **812 (INT8) vs 833 (fp32)** — i.e. conformance-neutral + (INT8 is marginally better on 3 fixtures, marginally worse on 2). +- Full INT8 (convs + MatMuls) was **rejected**: 3/23 exact, with clear quality + loss (section headers demoted to plain text, page-footer text leaking into + the output) — the RT-DETR head's class scores sit near the 0.3 threshold and + cannot tolerate activation quantization. +- Dynamic (weights-only) INT8 of the whole layout model was also rejected: it + is *slower* than fp32 (3.2 s vs 2.1 s per page-with-table) because inserted + per-activation quantize ops outweigh the MatMul savings while the conv + backbone stays fp32. + +### TableFormer decoder: dynamic INT8 (~10% faster tables, byte-identical) + +The autoregressive tag decoder is MatMul-only; weights-only dynamic INT8 +produced **byte-identical corpus output** and ~10% faster table decode +(784 → 695 ms/table), 78 → 50 MB. Small but free. + +The decoder speed is *not* weight-bound — it is per-step overhead (see backlog +item 2), which is why quantization helps so little there. + +## Ranked backlog of further ideas + +Ordered by expected impact ÷ risk. Items 1–3 attack the 85–95%. + +1. ~~**Ship/document the INT8 layout model as the default CPU + configuration**~~ **Done on this branch:** the pipeline prefers the int8 + models when present (`FLEISCHWOLF_FP32=1` opts out), + `download_dependencies.sh` fetches them by default, and + `publish-models.yml` builds them. Biggest single validated win: ~1.4–2× + end-to-end. +2. **TableFormer decode-loop overhead** (~800 ms/table, ~60–500 steps): + - ~~`decode_step` copies the whole KV cache out (`ocache.to_vec()`) and back + in every step — O(steps²·6·512) float traffic.~~ **Done on this branch:** + the cache and the encoder's cross-K/V + `enc_out` stay owned `ort` values + fed straight back into the next run (~9% faster structure decode, + byte-identical output). + - The exported graph still re-embeds the **full tag sequence** every step + (`tags` grows each iteration) even though attention is KV-cached. Re-export + the decoder to take only the last tag (the cache carries the history) — + this is the remaining per-step cost; see `scripts/export_tableformer.py`. +3. **Layout batching for the parallel path**: the pool currently runs batch-1 + inference per page. RT-DETR's 640×640 input is fixed-shape, so pages can be + batched (e.g. batch-4) per worker with one session — better core utilization + and less framework overhead on wide machines. Output is per-image, so + quality is unaffected. (Needs a re-export with a dynamic batch dim.) +4. **Render → resize pipeline copies** (`pdfium_backend.rs:264-272`, ~15% on + text-heavy docs): pdfium's BGRA bitmap goes through `as_image()` (copy + + swizzle) then `.into_rgb8()` (second copy) before the 3×→2× CatmullRom + downsample (third buffer). A single BGRA→RGB pass into the resize source + removes one full-page copy + traversal per page. Keep the 1.5× supersample + + CatmullRom itself — it is deliberate PIL-BICUBIC parity for model input. + Also: when `layout.predict` is the only image consumer at 640×640, the + 2×-page intermediate is only needed for TableFormer crops and OCR — a page + with no table regions and a text layer never needs it; rendering could be + deferred/skipped per page in a `no_ocr`-like fast path decided *after* + layout runs on a cheaper raster. +5. **textparse font caching** (marginal for PDFs — textparse is ≤1% — but + real for `no_ocr` mode where it becomes the bottleneck): + - fonts are fully re-parsed (ToUnicode CMap decompression + tokenization, + Type1 program scan, width maps) for **every page** and every Form-XObject + invocation (`textparse.rs:794`); cache parsed `Font`s per document keyed + by the font dict's `ObjectId`, and cache decoded Form XObject content. + - ~~`line_cells` + `word_cells` run the identical build+contract twice per + page; one pass can emit both.~~ **Done on this branch** + (`dp_lines::line_and_word_cells`): ~1.25× faster `--no-ocr` conversion, + identical output. + - `decode_code`/`decompose_ligatures` allocate a `String` per glyph + (`textparse.rs:94-145`); decompose once at font-parse time and return + borrowed `&str`. + - RTL merge is O(n²) (string prepend + `Vec::remove(0)`, + `dp_lines.rs:87-155`); accumulate reversed and flip once per line. +6. **OCR line batching** (`ocr.rs::recognize`): lines are recognized one at a + time on one thread (deliberately, for CTC determinism). Batching same-width + buckets keeps determinism per line and would speed scanned documents + several-fold; alternatively run multiple single-thread recognitions across + the existing worker pool. +7. **ort session options**: checked — ONNX Runtime's C-API default is already + `ORT_ENABLE_ALL`, so an explicit optimization level gains nothing. + `with_optimized_model_path` (caching the optimized graph on disk) could + still shave per-worker model-load latency; only worth it if pool spin-up + shows up in a real deployment. + +## Correctness notes found during review (quality, not speed) + +- `textparse.rs` `"` operator: the `aw ac string "` form must set word/char + spacing (`tw`/`tc`) from its first two operands before showing the string; + they are currently ignored (`Tj | ' | "` share one arm), so documents using + `"` get wrong inter-word advances. **Fixed in this branch.** +- `textparse.rs::page_size` ignores a non-zero MediaBox origin; a page with + e.g. `[9 9 621 801]` offsets all parser cells relative to pdfium's raster. + Rare, but cheap to guard: subtract the box origin when emitting glyph boxes. +- OCR recognition ran un-instrumented; `ocr.page` is now a timed stage (this + branch), so scanned-corpus profiles attribute it correctly. + +## Reproducing + +```bash +scripts/download_dependencies.sh +cargo build --release + +# stage timing +FLEISCHWOLF_TIMING=1 ./target/release/fleischwolf input.pdf > /dev/null + +# build the int8 models (used automatically once present) +uv venv .venv-quant && uv pip install --python .venv-quant/bin/python \ + onnx onnxruntime sympy pypdfium2 pillow numpy +.venv-quant/bin/python scripts/quantize_models.py + +# force full precision for a run +FLEISCHWOLF_FP32=1 ./target/release/fleischwolf input.pdf > /dev/null +``` + +Integration points: `scripts/download_dependencies.sh` fetches the +pre-quantized assets by default (`--no-int8` skips; published by +`.github/workflows/publish-models.yml`, which quantizes after export); +`scripts/pdf_setup.sh` quantizes locally unless `FLEISCHWOLF_FP32=1`; +`scripts/performance.sh` benchmarks whatever the pipeline default resolves to +(int8 when present, `FLEISCHWOLF_FP32=1` for fp32); `examples/Dockerfile` +bakes both precisions and defaults to int8 (`--build-arg INT8=0` for fp32). diff --git a/README.md b/README.md index 1a5a442b..ebb0fe3f 100644 --- a/README.md +++ b/README.md @@ -296,6 +296,7 @@ curl -fsSL https://raw.githubusercontent.com/artiz/fleischwolf/master/scripts/do | PP-OCRv3 rec + dictionary | `models/ocr_rec.onnx`, `models/ppocr_keys_v1.txt` | | TableFormer (optional) | `models/tableformer/{encoder,decoder,bbox}.onnx` (+ `.data` sidecars where the export needs them) | | Whisper tiny (audio/ASR; skip with `--no-asr`) | `models/asr/{encoder_model,decoder_model}.onnx`, `models/asr/vocab.json` (+ `added_tokens.json` for language selection) | +| INT8 CPU models (optional; fetch with `--int8`) | `models/layout_heron_int8.onnx`, `models/tableformer/decoder_int8.onnx` | Idempotent — safe to re-run; it skips files already on disk. Pass `--force` to re-fetch everything, or set `$FLEISCHWOLF_MODELS_URL` to fetch from a @@ -305,6 +306,31 @@ come from Hugging Face (`$FLEISCHWOLF_ASR_MODELS_URL` overrides, or point only for now — other platforms, or building the models from source, need [`scripts/pdf_setup.sh`](#testing) instead. +### INT8 models (faster PDF conversion on CPU — the default) + +The `*_int8` assets are post-training quantizations of the same models: +Conv-only static INT8 of the layout detector (calibrated on this repo's PDF +corpus) and dynamic INT8 of the TableFormer decoder. On CPUs with AVX-512 +VNNI they make layout inference — the dominant PDF cost — **~2.4× faster** +(~1.4–1.8× end-to-end) at conformance validated as unchanged against the +corpus groundtruth; the TableFormer output is byte-identical. See +[`PDF_PERFORMANCE.md`](./PDF_PERFORMANCE.md) for the measurements. + +**The pipeline uses them automatically** whenever they sit next to the fp32 +files at the default paths (`download_dependencies.sh` fetches them by +default; `--no-int8` skips, or build them with `python +scripts/quantize_models.py`). To force full precision: + +```bash +FLEISCHWOLF_FP32=1 fleischwolf input.pdf # keep the int8 files, use fp32 +# or pin a model explicitly — an explicit path always wins: +export DOCLING_LAYOUT_ONNX=$PWD/models/layout_heron.onnx +export DOCLING_TABLEFORMER_DECODER=$PWD/models/tableformer/decoder.onnx +``` + +(The [example Dockerfile](./examples/Dockerfile) bakes both precisions and +defaults to INT8; build with `--build-arg INT8=0` for pure fp32.) + Then either: ```bash @@ -446,14 +472,17 @@ The image converts PDFs/images fully offline; the model export (torch + ## Performance -`scripts/performance.sh` runs the **largest fixture of each supported type** through -both engines (published Python `docling` vs the Rust release binary) and reports -peak RSS, CPU utilization, and conversion time. Ratios below are docling ÷ -fleischwolf — bigger means Rust wins by more. +`scripts/performance.sh` runs a representative fixture of each supported type +through both engines (published Python `docling` vs the Rust release binary) and +reports peak RSS, CPU utilization, and conversion time. Ratios below are +docling ÷ fleischwolf — bigger means Rust wins by more. The PDF row is the +**fp32** stack; the optional [INT8 models](#int8-models-faster-pdf-conversion-on-cpu) +roughly double layout-inference speed on top of it (measured 1.83× end-to-end +on a 1913-page document — see [`PDF_PERFORMANCE.md`](./PDF_PERFORMANCE.md)). | File | Size | Peak-memory ratio | CPU ratio | Warm-conversion speedup | |---|---:|---:|---:|---:| -| `2203.01017v2.pdf` (PDF, 47 pp) | 6.9 MB | **2.2× less** | 1.3× | 1.2× | +| `picture_classification.pdf` (PDF) | 208 KB | **2.3× less** | 1.0× | 2.3× | | `docx_rich_tables_01.docx` (DOCX) | 3.1 MB | **41× less** | 2.7× | 21× | | `wiki_duck.html` (HTML) | 240 KB | **57× less** | 3.2× | 46× | | `elife-56337.nxml` (JATS XML) | 180 KB | **61× less** | 2.9× | 10× | @@ -465,10 +494,11 @@ fleischwolf — bigger means Rust wins by more. - **Peak memory** is where Rust wins decisively: a declarative conversion holds a few MB versus docling's ~750 MB (it imports torch even for non-ML formats). The PDF runs the full ML pipeline in both engines (torch vs ONNX), so the gap there - is 2.2× rather than 50×+, but Rust still peaks at 1.4 GB vs docling's 3.1 GB. + is 2.3× rather than 50×+, but Rust still peaks at 0.77 GB vs docling's 1.75 GB — + and the PDF converts **4.1× faster end-to-end** (docling re-pays its torch + import + model load on every invocation). - **CPU**: docling spreads across 2.7–3.2 cores for declarative work that Rust does - on a single core (~100%); on the PDF both go multi-core (Rust 525% vs docling - 674%). + on a single core (~100%); on the PDF both go multi-core (~330% each here). - **Warm-conversion speedup** isolates the parse/convert work — it times docling *in-process* (excluding its ~3 s interpreter + import startup) against the Rust whole-process figure. Rust wins on substantial inputs (HTML 46×, DOCX 21×); the diff --git a/crates/fleischwolf-pdf/src/dp_lines.rs b/crates/fleischwolf-pdf/src/dp_lines.rs index f500594d..3f4f2150 100644 --- a/crates/fleischwolf-pdf/src/dp_lines.rs +++ b/crates/fleischwolf-pdf/src/dp_lines.rs @@ -342,11 +342,47 @@ fn build_cells(glyphs: &[Glyph], euclidean: bool) -> Vec { /// Build line cells from a page's glyph stream via the docling-parse contraction. pub(crate) fn line_cells(glyphs: &[Glyph], page_h: f32, euclidean: bool) -> Vec { + line_and_word_cells(glyphs, page_h, euclidean).0 +} + +/// Build **word** cells from a page's glyph stream via the same contraction as +/// [`line_cells`]: each line splits into its constituent words at exactly the +/// points where the contraction inserted a separator space. This reproduces +/// docling-parse's `word_cells` (the per-word tokens TableFormer matches against +/// table-grid cells), letting the pipeline drop pdfium's text path entirely +/// (roadmap item 6). Empty words (overprint-cleared) are skipped. +pub(crate) fn word_cells(glyphs: &[Glyph], page_h: f32, euclidean: bool) -> Vec { + line_and_word_cells(glyphs, page_h, euclidean).1 +} + +/// Build the line cells **and** the word cells from one shared contraction — the +/// build+contract pass is the expensive step and is identical for both views, so +/// callers that need both (the default text layer) pay it once. Line cells come +/// from each contracted cell's text/box; word cells from its recorded word +/// segments. +pub(crate) fn line_and_word_cells( + glyphs: &[Glyph], + page_h: f32, + euclidean: bool, +) -> (Vec, Vec) { let mut cells = build_cells(glyphs, euclidean); contract(&mut cells, euclidean); - cells + let mut words = Vec::new(); + let lines = cells .into_iter() .map(|c| { + for w in c.words { + if w.text.trim().is_empty() { + continue; + } + words.push(TextCell { + text: w.text, + l: w.l as f32, + t: page_h - w.t as f32, + r: w.r as f32, + b: page_h - w.b as f32, + }); + } let l = c.rx0.min(c.rx1).min(c.rx2).min(c.rx3) as f32; let r = c.rx0.max(c.rx1).max(c.rx2).max(c.rx3) as f32; let top = c.ry0.max(c.ry1).max(c.ry2).max(c.ry3) as f32; @@ -359,34 +395,8 @@ pub(crate) fn line_cells(glyphs: &[Glyph], page_h: f32, euclidean: bool) -> Vec< b: page_h - bot, } }) - .collect() -} - -/// Build **word** cells from a page's glyph stream via the same contraction as -/// [`line_cells`], then split each line into its constituent words at exactly the -/// points where the contraction inserted a separator space. This reproduces -/// docling-parse's `word_cells` (the per-word tokens TableFormer matches against -/// table-grid cells), letting the pipeline drop pdfium's text path entirely -/// (roadmap item 6). Empty words (overprint-cleared) are skipped. -pub(crate) fn word_cells(glyphs: &[Glyph], page_h: f32, euclidean: bool) -> Vec { - let mut cells = build_cells(glyphs, euclidean); - contract(&mut cells, euclidean); - let mut out = Vec::new(); - for c in cells { - for w in c.words { - if w.text.trim().is_empty() { - continue; - } - out.push(TextCell { - text: w.text, - l: w.l as f32, - t: page_h - w.t as f32, - r: w.r as f32, - b: page_h - w.b as f32, - }); - } - } - out + .collect(); + (lines, words) } fn is_rtl_char(c: char) -> bool { diff --git a/crates/fleischwolf-pdf/src/layout.rs b/crates/fleischwolf-pdf/src/layout.rs index bac683d9..26811a2f 100644 --- a/crates/fleischwolf-pdf/src/layout.rs +++ b/crates/fleischwolf-pdf/src/layout.rs @@ -52,7 +52,9 @@ pub struct LayoutModel { } impl LayoutModel { - /// Load the ONNX model from `DOCLING_LAYOUT_ONNX` (or `models/layout_heron.onnx`). + /// Load the ONNX model from `DOCLING_LAYOUT_ONNX`. Without the override, + /// prefers `models/layout_heron_int8.onnx` when present (the quantized + /// default; `FLEISCHWOLF_FP32=1` opts out), else `models/layout_heron.onnx`. pub fn load() -> Result { Self::load_with(crate::intra_threads()) } @@ -61,8 +63,11 @@ impl LayoutModel { /// parallel page-worker pool loads its helper models on a single thread each /// and gets its speed-up from running pages concurrently instead. pub fn load_with(intra: usize) -> Result { - let path = std::env::var("DOCLING_LAYOUT_ONNX") - .unwrap_or_else(|_| "models/layout_heron.onnx".to_string()); + let path = crate::model_path( + "DOCLING_LAYOUT_ONNX", + "models/layout_heron.onnx", + "models/layout_heron_int8.onnx", + ); let session = Session::builder() .map_err(|e| format!("layout: builder: {e}"))? // Let inference use the available cores (ort otherwise defaults low); diff --git a/crates/fleischwolf-pdf/src/lib.rs b/crates/fleischwolf-pdf/src/lib.rs index 85d8bb89..f608f355 100644 --- a/crates/fleischwolf-pdf/src/lib.rs +++ b/crates/fleischwolf-pdf/src/lib.rs @@ -74,6 +74,29 @@ pub(crate) fn intra_threads() -> usize { .unwrap_or(1) } +/// True when `FLEISCHWOLF_FP32` (any value but `0`) forces the full-precision +/// models even where an INT8 variant sits next to the fp32 default. +pub(crate) fn fp32_forced() -> bool { + std::env::var("FLEISCHWOLF_FP32") + .map(|v| v != "0") + .unwrap_or(false) +} + +/// Resolve a model path: an explicit env override always wins; otherwise the +/// INT8 variant of the default path when it exists on disk (the quantized +/// models are conformance-validated — see PDF_PERFORMANCE.md — and load/run +/// markedly faster on CPU), unless `FLEISCHWOLF_FP32` opts back into full +/// precision; else the fp32 default. +pub(crate) fn model_path(env: &str, fp32_default: &str, int8_default: &str) -> String { + if let Ok(p) = std::env::var(env) { + return p; + } + if !fp32_forced() && std::path::Path::new(int8_default).exists() { + return int8_default.to_string(); + } + fp32_default.to_string() +} + /// One page's assembled output: typed nodes plus the page's hyperlinks, kept /// separate so pages processed out of order can be stitched back in page order. type PageOut = (Vec, Vec<(String, String)>); @@ -148,12 +171,13 @@ impl Worker { if self.ocr.is_none() { self.ocr = Some(ocr::OcrModel::load().map_err(PdfError::Ocr)?); } - let cells = self - .ocr - .as_mut() - .unwrap() - .ocr_page(&page.image, ®ions, page.scale) - .map_err(|e| PdfError::Ocr(format!("page {}: {e}", n + 1)))?; + let cells = timing::timed("ocr.page", || { + self.ocr + .as_mut() + .unwrap() + .ocr_page(&page.image, ®ions, page.scale) + }) + .map_err(|e| PdfError::Ocr(format!("page {}: {e}", n + 1)))?; page.cells = cells; } // TableFormer structure per table region (else geometric fallback). diff --git a/crates/fleischwolf-pdf/src/tableformer.rs b/crates/fleischwolf-pdf/src/tableformer.rs index 2965c039..39c91cff 100644 --- a/crates/fleischwolf-pdf/src/tableformer.rs +++ b/crates/fleischwolf-pdf/src/tableformer.rs @@ -7,7 +7,7 @@ use crate::pdfium_backend::TextCell; use image::RgbImage; use ort::session::Session; -use ort::value::{Tensor, TensorRef}; +use ort::value::{DynValue, Tensor}; const SIDE: u32 = 448; // Verbatim from docling's tm_config.json image_normalization (more digits than @@ -58,14 +58,12 @@ pub struct TableFormer { /// 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. Each is a `(shape, flattened data)` pair. +/// `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_shape: Vec, - ck: Vec, - cv_shape: Vec, - cv: Vec, - eo_shape: Vec, - eo: Vec, + ck: DynValue, + cv: DynValue, + eo: DynValue, } impl TableFormer { @@ -82,8 +80,13 @@ impl TableFormer { pub fn load_with(intra: usize) -> Option { let enc = std::env::var("DOCLING_TABLEFORMER_ENCODER") .unwrap_or_else(|_| "models/tableformer/encoder.onnx".to_string()); - let dec = std::env::var("DOCLING_TABLEFORMER_DECODER") - .unwrap_or_else(|_| "models/tableformer/decoder.onnx".to_string()); + // Prefer the INT8 decoder when present (byte-identical output, faster + // decode; FLEISCHWOLF_FP32=1 opts out) unless explicitly overridden. + let dec = crate::model_path( + "DOCLING_TABLEFORMER_DECODER", + "models/tableformer/decoder.onnx", + "models/tableformer/decoder_int8.onnx", + ); let bbx = std::env::var("DOCLING_TABLEFORMER_BBOX") .unwrap_or_else(|_| "models/tableformer/bbox.onnx".to_string()); if [&enc, &dec, &bbx] @@ -129,77 +132,60 @@ impl TableFormer { /// shape `[N_LAYERS,1,H,S,head_dim]`) and `enc_out` for the bbox decoder. fn encode(&mut self, img: &RgbImage) -> Result { let input = preprocess(img)?; - let enc_out = self + let mut enc_out = self .encoder .run(ort::inputs!["image" => input]) .map_err(|e| format!("tableformer: encode: {e}"))?; - let grab = |name: &str| -> Result<(Vec, Vec), String> { - let (sh, data) = enc_out[name] - .try_extract_tensor::() - .map_err(|e| format!("tableformer: {name}: {e}"))?; - Ok((sh.iter().map(|&x| x as usize).collect(), data.to_vec())) + let mut grab = |name: &str| -> Result { + enc_out + .remove(name) + .ok_or_else(|| format!("tableformer: encoder output {name} missing")) }; - let (ck_shape, ck) = grab("cross_k")?; - let (cv_shape, cv) = grab("cross_v")?; - let (eo_shape, eo) = grab("enc_out")?; Ok(EncodeOut { - ck_shape, - ck, - cv_shape, - cv, - eo_shape, - eo, + ck: grab("cross_k")?, + cv: grab("cross_v")?, + eo: grab("enc_out")?, }) } /// One doubly-cached decode step: feed the current `tags`, the constant cross - /// K/V views, and the growing self-attention `cache`; return the raw argmax tag - /// and the last token's hidden state, advancing the cache. `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). + /// 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 Vec, - cache_past: &mut usize, + cache: &mut Option, empty_cache: &Tensor, ) -> Result<(i64, Vec), String> { let tags_t = Tensor::from_array(([tags.len(), 1usize], tags.to_vec())) .map_err(|e| format!("tableformer: tags: {e}"))?; - // Constant per-table cross-attention K/V — zero-copy views each step. - let ck_t = TensorRef::from_array_view((enc.ck_shape.as_slice(), enc.ck.as_slice())) - .map_err(|e| format!("tableformer: cross_k: {e}"))?; - let cv_t = TensorRef::from_array_view((enc.cv_shape.as_slice(), enc.cv.as_slice())) - .map_err(|e| format!("tableformer: cross_v: {e}"))?; - let dout = if *cache_past == 0 { - self.decoder.run(ort::inputs![ - "tags" => tags_t, "cross_k" => ck_t, "cross_v" => cv_t, "cache" => empty_cache]) - } else { - let cache_t = TensorRef::from_array_view(( - [N_LAYERS, *cache_past, 1, EMBED_DIM], - cache.as_slice(), - )) - .map_err(|e| format!("tableformer: cache: {e}"))?; - self.decoder.run(ort::inputs![ - "tags" => tags_t, "cross_k" => ck_t, "cross_v" => cv_t, "cache" => cache_t]) + let mut dout = match cache.as_ref() { + None => self.decoder.run(ort::inputs![ + "tags" => tags_t, "cross_k" => &enc.ck, "cross_v" => &enc.cv, + "cache" => empty_cache]), + 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::() .map_err(|e| format!("tableformer: logits: {e}"))?; let raw = argmax(logits) as i64; - let (oshape, ocache) = dout["out_cache"] - .try_extract_tensor::() - .map_err(|e| format!("tableformer: out_cache: {e}"))?; - let next_cache = ocache.to_vec(); - let next_past = oshape[1] as usize; let (_, hidden) = dout["hidden"] .try_extract_tensor::() .map_err(|e| format!("tableformer: hidden: {e}"))?; let hidden = hidden.to_vec(); - *cache = next_cache; - *cache_past = next_past; + *cache = Some( + dout.remove("out_cache") + .ok_or_else(|| "tableformer: decoder output out_cache missing".to_string())?, + ); Ok((raw, hidden)) } @@ -218,12 +204,10 @@ impl TableFormer { let mut tags: Vec = vec![START]; let mut out: Vec = Vec::new(); let mut prev_ucel = false; - let mut cache: Vec = Vec::new(); - let mut cache_past = 0usize; + let mut cache: Option = None; let empty = self.empty_cache()?; while out.len() < MAX_STEPS { - let (raw, _hidden) = - self.decode_step(&tags, &enc, &mut cache, &mut cache_past, &empty)?; + let (raw, _hidden) = self.decode_step(&tags, &enc, &mut cache, &empty)?; let mut tag = raw; if tag == XCEL { tag = LCEL; @@ -259,12 +243,10 @@ impl TableFormer { let mut bbox_ind = 0usize; let mut cur_bbox_ind = 0usize; let mut merge: std::collections::HashMap = std::collections::HashMap::new(); - let mut cache: Vec = Vec::new(); - let mut cache_past = 0usize; + let mut cache: Option = None; let empty = self.empty_cache()?; while otsl.len() < MAX_STEPS { - let (raw, hidden) = - self.decode_step(&tags, &enc, &mut cache, &mut cache_past, &empty)?; + let (raw, hidden) = self.decode_step(&tags, &enc, &mut cache, &empty)?; let mut tag = raw; if tag == XCEL { tag = LCEL; @@ -304,11 +286,9 @@ impl TableFormer { } let tag_h = Tensor::from_array(([n, 512usize], hiddens)) .map_err(|e| format!("tableformer: tag_h: {e}"))?; - let eo_t = Tensor::from_array((enc.eo_shape.clone(), enc.eo.clone())) - .map_err(|e| format!("tableformer: eo: {e}"))?; let bout = self .bbox - .run(ort::inputs!["enc_out" => eo_t, "tag_h" => tag_h]) + .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::() @@ -337,7 +317,9 @@ impl TableFormer { // page → 1024px height (cv2.INTER_AREA), then crop the table bbox. let sf = 1024.0 / page_image.height() as f32; let pw = (page_image.width() as f32 * sf) as u32; - let page1024 = crate::resample::inter_area(page_image, pw, 1024); + let page1024 = crate::timing::timed("tableformer.inter_area", || { + crate::resample::inter_area(page_image, pw, 1024) + }); let k = 1024.0 / page_h; let x = (region[0] * k).round().max(0.0) as u32; let y = (region[1] * k).round().max(0.0) as u32; @@ -347,7 +329,10 @@ impl TableFormer { return None; } let crop = image::imageops::crop_imm(&page1024, x, y, x2 - x, y2 - y).to_image(); - let cells = self.predict_table_structure(&crop).ok()?; + let cells = crate::timing::timed("tableformer.structure", || { + self.predict_table_structure(&crop) + }) + .ok()?; if cells.is_empty() { return None; } diff --git a/crates/fleischwolf-pdf/src/textparse.rs b/crates/fleischwolf-pdf/src/textparse.rs index c8ec505b..d4349605 100644 --- a/crates/fleischwolf-pdf/src/textparse.rs +++ b/crates/fleischwolf-pdf/src/textparse.rs @@ -702,9 +702,10 @@ pub fn pdf_all_cells(bytes: &[u8]) -> Vec { .map(|(_, pid)| { let (_w, h) = page_size(&doc, pid); let glyphs = page_glyphs(&doc, pid); + let (prose, words) = crate::dp_lines::line_and_word_cells(&glyphs, h, true); PageParserCells { - prose: crate::dp_lines::line_cells(&glyphs, h, true), - words: crate::dp_lines::word_cells(&glyphs, h, true), + prose, + words, code: crate::pdfium_backend::code_cells_from_glyphs(&glyphs, h), } }) @@ -922,6 +923,12 @@ fn run_content( .then(tlm); tm = tlm; } + if op.operator == "\"" { + // `aw ac string "` sets word- and char-spacing before + // showing the string (PDF 32000-1 §9.4.3), persisting after. + tw = op_f(operands, 0); + tc = op_f(operands, 1); + } if let (Some(f), Some(Object::String(s, _))) = (font, operands.last()) { show_text(f, s, fsize, tc, tw, th, trise, &mut tm, ctm, out); } diff --git a/examples/Dockerfile b/examples/Dockerfile index a5ff6c24..21019519 100644 --- a/examples/Dockerfile +++ b/examples/Dockerfile @@ -11,9 +11,17 @@ # Convert a document (Markdown to stdout): # docker run --rm -v "$PWD:/data" fleischwolf /data/input.pdf # +# By default the models stage also INT8-quantizes the layout model (Conv-only, +# calibrated on the repo's PDF corpus) and the TableFormer decoder, and the +# runtime points at the int8 files — ~2.4x faster layout inference on CPU at +# validated-unchanged conformance (see PDF_PERFORMANCE.md). Build with +# `--build-arg INT8=0` for pure fp32, or override per container: +# docker run -e DOCLING_LAYOUT_ONNX=/opt/models/layout_heron.onnx ... +# (both precisions are baked into the image). +# # Stages: -# models — export the layout + TableFormer ONNX (torch, one-time) and fetch -# the OCR model/dict + pdfium +# models — export the layout + TableFormer ONNX (torch, one-time), fetch +# the OCR model/dict + pdfium, INT8-quantize (unless INT8=0) # builder — compile the CLI (the `ort` crate downloads ONNX Runtime at build) # runtime — slim image with just the binary, native libs and models @@ -26,14 +34,19 @@ RUN apt-get update \ && apt-get install -y --no-install-recommends curl ca-certificates tar \ && rm -rf /var/lib/apt/lists/* -# torch (CPU) + the docling model packages used by the export scripts. +# torch (CPU) + the docling model packages used by the export scripts, plus +# the quantization/calibration deps (pypdfium2, pillow, sympy). RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu \ && pip install --no-cache-dir --extra-index-url https://download.pytorch.org/whl/cpu \ "transformers>=4.45,<5" "onnx>=1.16,<2" \ - docling-ibm-models onnxscript onnxruntime huggingface_hub + docling-ibm-models onnxscript onnxruntime huggingface_hub \ + sympy pypdfium2 pillow numpy WORKDIR /build -COPY scripts/export_layout.py scripts/export_tableformer.py ./scripts/ +COPY scripts/export_layout.py scripts/export_tableformer.py scripts/quantize_models.py ./scripts/ +# Real corpus pages for INT8 calibration (build-stage only, not in the runtime). +COPY tests/data/pdf/sources ./tests/data/pdf/sources +COPY tests/data/scanned/sources ./tests/data/scanned/sources RUN mkdir -p /opt/models /opt/pdfium \ # RT-DETR layout detector -> models/layout_heron.onnx @@ -54,6 +67,21 @@ RUN mkdir -p /opt/models /opt/pdfium \ && tar xzf /tmp/pdfium.tgz -C /opt/pdfium \ && rm /tmp/pdfium.tgz +# INT8=1 (default): quantize layout (static, conv-only, calibrated on the +# corpus copied above) + TableFormer decoder (dynamic), and point the +# `*_active` names — which the runtime ENV references — at the int8 files. +# INT8=0 keeps the active names on fp32. Both precisions ship in the image +# either way, so a container can switch back via the DOCLING_* env vars. +ARG INT8=1 +RUN if [ "$INT8" = "1" ]; then \ + FLEISCHWOLF_MODELS_DIR=/opt/models python scripts/quantize_models.py \ + && ln -s layout_heron_int8.onnx /opt/models/layout_heron_active.onnx \ + && ln -s decoder_int8.onnx /opt/models/tableformer/decoder_active.onnx; \ + else \ + ln -s layout_heron.onnx /opt/models/layout_heron_active.onnx \ + && ln -s decoder.onnx /opt/models/tableformer/decoder_active.onnx; \ + fi + # ---------------------------------------------------------------------------- # Stage 2: Rust build # ---------------------------------------------------------------------------- @@ -97,12 +125,16 @@ COPY --from=models /opt/pdfium /opt/pdfium # The pipeline finds its native lib + models via these env vars (the TableFormer # graphs default to models/tableformer/*; set them explicitly so the binary works # from any working directory). +# The `*_active` names are symlinks resolved by the models stage's INT8 arg +# (int8 by default, fp32 with --build-arg INT8=0); the concrete files are all +# present, so a container can pin either precision explicitly, e.g. +# `-e DOCLING_LAYOUT_ONNX=/opt/models/layout_heron.onnx` for fp32. ENV PDFIUM_DYNAMIC_LIB_PATH=/opt/pdfium/lib \ - DOCLING_LAYOUT_ONNX=/opt/models/layout_heron.onnx \ + DOCLING_LAYOUT_ONNX=/opt/models/layout_heron_active.onnx \ DOCLING_OCR_REC_ONNX=/opt/models/ocr_rec.onnx \ DOCLING_OCR_DICT=/opt/models/ppocr_keys_v1.txt \ DOCLING_TABLEFORMER_ENCODER=/opt/models/tableformer/encoder.onnx \ - DOCLING_TABLEFORMER_DECODER=/opt/models/tableformer/decoder.onnx \ + DOCLING_TABLEFORMER_DECODER=/opt/models/tableformer/decoder_active.onnx \ DOCLING_TABLEFORMER_BBOX=/opt/models/tableformer/bbox.onnx \ LD_LIBRARY_PATH=/usr/local/lib:/opt/pdfium/lib diff --git a/scripts/download_dependencies.sh b/scripts/download_dependencies.sh index a2eea874..190f4fc2 100755 --- a/scripts/download_dependencies.sh +++ b/scripts/download_dependencies.sh @@ -29,6 +29,16 @@ # models/asr/{encoder_model,decoder_model}.onnx + vocab.json (Whisper tiny, # from Hugging Face; skip with --no-asr) # +# Also fetches the INT8-quantized CPU models when the release hosts them (see +# PDF_PERFORMANCE.md — ~2.4x faster layout inference at unchanged conformance): +# models/layout_heron_int8.onnx +# models/tableformer/decoder_int8.onnx +# The pipeline picks these up automatically when they sit next to the fp32 +# files (no env vars needed); set FLEISCHWOLF_FP32=1 at runtime to force full +# precision, or skip fetching them entirely with --no-int8. If the release +# doesn't host the int8 assets (older tag), a note explains how to produce +# them locally with scripts/quantize_models.py. +# # pdfium is Linux x64 only for now, matching what's hosted in the release; for # other platforms (or to build the models from source) see scripts/pdf_setup.sh. # @@ -44,12 +54,15 @@ ASR_BASE_URL="${FLEISCHWOLF_ASR_MODELS_URL:-https://huggingface.co/onnx-communit FORCE=false WITH_ASR=true +WITH_INT8=true for arg in "$@"; do case "$arg" in --force) FORCE=true ;; --no-asr) WITH_ASR=false ;; + --int8) WITH_INT8=true ;; # accepted for compatibility; int8 is the default + --no-int8) WITH_INT8=false ;; *) - echo "usage: download_dependencies.sh [--force] [--no-asr]" >&2 + echo "usage: download_dependencies.sh [--force] [--no-asr] [--no-int8]" >&2 exit 2 ;; esac @@ -109,4 +122,20 @@ if [ "$WITH_ASR" = true ]; then fetch_optional "$ASR_BASE_URL/added_tokens.json" models/asr/added_tokens.json fi +if [ "$WITH_INT8" = true ]; then + # INT8-quantized CPU models (optional release assets). The pipeline prefers + # them automatically when they sit at the default paths; FLEISCHWOLF_FP32=1 + # forces the fp32 models at runtime. + fetch_optional "$BASE_URL/layout_heron_int8.onnx" models/layout_heron_int8.onnx + fetch_optional "$BASE_URL/decoder_int8.onnx" models/tableformer/decoder_int8.onnx + if [ -f models/layout_heron_int8.onnx ]; then + echo "int8 models present — used by default (FLEISCHWOLF_FP32=1 forces full precision)" + else + echo "int8 assets not hosted at $BASE_URL — the fp32 models will be used." + echo "To build the int8 models locally (see PDF_PERFORMANCE.md):" + echo " pip install onnx onnxruntime sympy pypdfium2 pillow numpy" + echo " python scripts/quantize_models.py" + fi +fi + echo "done — models/ and .pdfium/lib populated in $(pwd)" diff --git a/scripts/pdf_setup.sh b/scripts/pdf_setup.sh index 82f80e11..c1f523f5 100755 --- a/scripts/pdf_setup.sh +++ b/scripts/pdf_setup.sh @@ -50,9 +50,26 @@ if [ ! -f models/tableformer/decoder.onnx ]; then fi fi +# INT8-quantize the layout model + TableFormer decoder for faster CPU +# inference (validated conformance-neutral — see PDF_PERFORMANCE.md). The +# pipeline prefers the int8 files automatically once they exist; skip building +# them with FLEISCHWOLF_FP32=1. Needs onnx onnxruntime sympy pypdfium2 pillow +# numpy — a missing-deps failure is non-fatal (fp32 keeps working). +if [ "${FLEISCHWOLF_FP32:-0}" != "1" ] && [ ! -f models/layout_heron_int8.onnx ]; then + echo "→ INT8-quantizing layout + TableFormer decoder" + if ! "${PYTHON:-python3}" scripts/quantize_models.py; then + echo " ! quantization failed (missing deps?). The fp32 models still work;" + echo " re-run after: pip install onnx onnxruntime sympy pypdfium2 pillow numpy" + fi +fi + echo "done. export these before running the pipeline:" echo " export PDFIUM_DYNAMIC_LIB_PATH=$(pwd)/.pdfium/lib" -echo " export DOCLING_LAYOUT_ONNX=$(pwd)/models/layout_heron.onnx" +if [ -f models/layout_heron_int8.onnx ]; then + echo " export DOCLING_LAYOUT_ONNX=$(pwd)/models/layout_heron_int8.onnx # int8 default (fp32: layout_heron.onnx, or FLEISCHWOLF_FP32=1)" +else + echo " export DOCLING_LAYOUT_ONNX=$(pwd)/models/layout_heron.onnx" +fi echo " export DOCLING_OCR_REC_ONNX=$(pwd)/models/ocr_rec.onnx" echo " export DOCLING_OCR_DICT=$(pwd)/models/ppocr_keys_v1.txt" if [ -f models/tableformer/decoder.onnx ]; then @@ -61,6 +78,10 @@ if [ -f models/tableformer/decoder.onnx ]; then # still loads (instead of silently falling back to geometric reconstruction) # no matter where the binary/binding is actually invoked from. echo " export DOCLING_TABLEFORMER_ENCODER=$(pwd)/models/tableformer/encoder.onnx" - echo " export DOCLING_TABLEFORMER_DECODER=$(pwd)/models/tableformer/decoder.onnx" + if [ -f models/tableformer/decoder_int8.onnx ]; then + echo " export DOCLING_TABLEFORMER_DECODER=$(pwd)/models/tableformer/decoder_int8.onnx # int8 default (fp32: decoder.onnx, or FLEISCHWOLF_FP32=1)" + else + echo " export DOCLING_TABLEFORMER_DECODER=$(pwd)/models/tableformer/decoder.onnx" + fi echo " export DOCLING_TABLEFORMER_BBOX=$(pwd)/models/tableformer/bbox.onnx" fi diff --git a/scripts/performance.sh b/scripts/performance.sh index 9bf093c4..938996a4 100755 --- a/scripts/performance.sh +++ b/scripts/performance.sh @@ -42,7 +42,14 @@ RUST_BIN="$(build_rust_release)" # Point the Rust PDF pipeline at the fetched libs/models (scripts/pdf_setup.sh) # using absolute paths, so it runs the full pipeline no matter the caller's CWD. # Harmless for non-PDF inputs (the binary only reads these for PDFs/images). +# Mirrors the pipeline's own default: the INT8 layout model + TableFormer +# decoder when present (scripts/quantize_models.py; see PDF_PERFORMANCE.md), +# fp32 with FLEISCHWOLF_FP32=1. [[ -e "$WORKSPACE_DIR/.pdfium/lib/libpdfium.so" ]] && export PDFIUM_DYNAMIC_LIB_PATH="${PDFIUM_DYNAMIC_LIB_PATH:-$WORKSPACE_DIR/.pdfium/lib}" +if [[ "${FLEISCHWOLF_FP32:-0}" != "1" && -e "$WORKSPACE_DIR/models/layout_heron_int8.onnx" ]]; then + export DOCLING_LAYOUT_ONNX="${DOCLING_LAYOUT_ONNX:-$WORKSPACE_DIR/models/layout_heron_int8.onnx}" + [[ -e "$WORKSPACE_DIR/models/tableformer/decoder_int8.onnx" ]] && export DOCLING_TABLEFORMER_DECODER="${DOCLING_TABLEFORMER_DECODER:-$WORKSPACE_DIR/models/tableformer/decoder_int8.onnx}" +fi [[ -e "$WORKSPACE_DIR/models/layout_heron.onnx" ]] && export DOCLING_LAYOUT_ONNX="${DOCLING_LAYOUT_ONNX:-$WORKSPACE_DIR/models/layout_heron.onnx}" [[ -e "$WORKSPACE_DIR/models/ocr_rec.onnx" ]] && export DOCLING_OCR_REC_ONNX="${DOCLING_OCR_REC_ONNX:-$WORKSPACE_DIR/models/ocr_rec.onnx}" [[ -e "$WORKSPACE_DIR/models/ppocr_keys_v1.txt" ]] && export DOCLING_OCR_DICT="${DOCLING_OCR_DICT:-$WORKSPACE_DIR/models/ppocr_keys_v1.txt}" diff --git a/scripts/quantize_models.py b/scripts/quantize_models.py new file mode 100644 index 00000000..69290e56 --- /dev/null +++ b/scripts/quantize_models.py @@ -0,0 +1,154 @@ +#!/usr/bin/env python3 +"""Quantize the PDF-pipeline ONNX models to INT8 for faster CPU inference. + +Two quantizations, both validated against the PDF corpus (see +PDF_PERFORMANCE.md for the measured speed/quality numbers): + +* **layout** — static QDQ INT8 of the RT-DETR layout model, **Conv ops only** + (the HGNetv2 backbone). Calibrated on real corpus pages rendered exactly the + way `layout.rs::predict` preprocesses them. The transformer decoder and + detection heads stay fp32: quantizing their MatMuls shifts class scores near + the 0.3 threshold and visibly degrades output (headers demoted to text, + page-footers leaking in), while conv-only keeps groundtruth conformance at + fp32 level. ~2.4x faster layout inference on AVX-512-VNNI CPUs, 172 -> 68 MB. + +* **tableformer-decoder** — dynamic INT8 (weights-only MatMul) of the + autoregressive tag decoder. Output is byte-identical on the corpus; + ~10% faster table-structure decode, 78 -> 50 MB. + +Usage (from the repo root, models fetched by scripts/download_dependencies.sh): + + uv venv .venv-quant && uv pip install --python .venv-quant/bin/python \ + onnx onnxruntime sympy pypdfium2 pillow numpy + .venv-quant/bin/python scripts/quantize_models.py layout tableformer-decoder + +Then point the pipeline at the quantized files: + + export DOCLING_LAYOUT_ONNX=$PWD/models/layout_heron_int8.onnx + export DOCLING_TABLEFORMER_DECODER=$PWD/models/tableformer/decoder_int8.onnx + +Re-run scripts/pdf_conformance.sh (or diff Markdown against +tests/data/pdf/groundtruth) after re-quantizing to re-verify quality. +""" + +import glob +import os +import sys + +import numpy as np + +REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +# Where the fp32 models live and the int8 outputs go (a checkout's models/ by +# default); FLEISCHWOLF_MODELS_DIR relocates it (e.g. /opt/models in a Docker +# models stage). +MODELS = os.environ.get("FLEISCHWOLF_MODELS_DIR", f"{REPO}/models") +# FLEISCHWOLF_CALIBRATION_DIR: a directory scanned recursively for calibration +# *.pdf files; defaults to the repo's PDF + scanned corpus (the set the +# published quality numbers were measured with). +CALIB = os.environ.get("FLEISCHWOLF_CALIBRATION_DIR") +SIDE = 640 # layout model input side (layout.rs) + + +def calibration_pages(): + """Render up to 3 pages of every calibration PDF the way layout.rs + preprocesses: pdfium at scale 2.0, resize to 640x640 bilinear, /255, CHW + float32.""" + import pypdfium2 as pdfium + from PIL import Image + + if CALIB: + pdfs = sorted(glob.glob(f"{CALIB}/**/*.pdf", recursive=True)) + else: + pdfs = sorted(glob.glob(f"{REPO}/tests/data/pdf/sources/*.pdf")) + sorted( + glob.glob(f"{REPO}/tests/data/scanned/sources/*.pdf") + ) + if not pdfs: + sys.exit("no calibration PDFs found (set FLEISCHWOLF_CALIBRATION_DIR)") + for path in pdfs: + try: + doc = pdfium.PdfDocument(path) + except Exception: + continue + for i in range(min(3, len(doc))): + bmp = doc[i].render(scale=2.0) + img = bmp.to_pil().convert("RGB").resize((SIDE, SIDE), Image.BILINEAR) + arr = np.asarray(img, dtype=np.float32) / 255.0 + yield np.transpose(arr, (2, 0, 1))[None, ...] + doc.close() + + +def quantize_layout(): + from onnxruntime.quantization import ( + CalibrationDataReader, + QuantFormat, + QuantType, + quantize_static, + ) + from onnxruntime.quantization.shape_inference import quant_pre_process + + src = f"{MODELS}/layout_heron.onnx" + pre = f"{MODELS}/layout_heron_pre.onnx" + dst = f"{MODELS}/layout_heron_int8.onnx" + + class Reader(CalibrationDataReader): + def __init__(self): + self.data = [{"pixel_values": x} for x in calibration_pages()] + print(f"layout: {len(self.data)} calibration samples", flush=True) + self.it = iter(self.data) + + def get_next(self): + return next(self.it, None) + + print("layout: pre-processing (shape inference)...", flush=True) + quant_pre_process(src, pre, skip_symbolic_shape=False) + print("layout: static QDQ INT8 quantization (Conv only)...", flush=True) + quantize_static( + pre, + dst, + Reader(), + quant_format=QuantFormat.QDQ, + activation_type=QuantType.QUInt8, + weight_type=QuantType.QInt8, + per_channel=True, + # Conv-only: the RT-DETR decoder/head MatMuls are threshold-sensitive. + op_types_to_quantize=["Conv"], + ) + os.remove(pre) + print(f"layout: done -> {dst} ({os.path.getsize(dst) / 1e6:.1f} MB)") + + +def quantize_tableformer_decoder(): + import onnx + from onnxruntime.quantization import QuantType, quantize_dynamic + + src = f"{MODELS}/tableformer/decoder.onnx" + tmp = f"{MODELS}/tableformer/decoder_clean.onnx" + dst = f"{MODELS}/tableformer/decoder_int8.onnx" + + # The export carries stale value_info shapes that break ORT's shape + # inference; strip them (external weights get folded into the output). + m = onnx.load(src) + del m.graph.value_info[:] + onnx.save(m, tmp, save_as_external_data=True, location="decoder_clean.onnx.data") + print("tableformer-decoder: dynamic INT8 quantization...", flush=True) + quantize_dynamic( + tmp, dst, weight_type=QuantType.QInt8, extra_options={"MatMulConstBOnly": True} + ) + os.remove(tmp) + os.remove(f"{tmp}.data") + print(f"tableformer-decoder: done -> {dst} ({os.path.getsize(dst) / 1e6:.1f} MB)") + + +def main(): + targets = sys.argv[1:] or ["layout", "tableformer-decoder"] + for t in targets: + if t == "layout": + quantize_layout() + elif t == "tableformer-decoder": + quantize_tableformer_decoder() + else: + sys.exit(f"unknown target {t!r} (expected: layout, tableformer-decoder)") + + +if __name__ == "__main__": + main()