Skip to main content

docling_pdf/
lib.rs

1//! PDF backend for docling.rs.
2//!
3//! A port of docling's standard PDF pipeline: pdfium extracts the text layer
4//! (cells with bounding boxes) and renders page images; a discriminative ONNX
5//! stack (layout detection, table structure, OCR) classifies regions; the cells
6//! are assembled in reading order into a [`DoclingDocument`].
7//!
8//! Current stages: pdfium text-cell extraction + page rendering ([`pdfium_backend`])
9//! and the deterministic text/reading-order assembly ([`assemble`]). The layout,
10//! table-structure and OCR ONNX stages land behind [`Pipeline`] next.
11
12// Without `ml` only the text-layer path runs; the shared assembly/label
13// helpers it doesn't exercise stay compiled for API stability (the full
14// build still flags genuinely dead code).
15#![cfg_attr(not(feature = "ml"), allow(dead_code))]
16
17// Reading-order assembly. Public under `ocr-prep` so the browser pipeline can
18// reuse the geometric table reconstruction and its reliability gate (#157).
19#[cfg(feature = "ocr-prep")]
20pub mod assemble;
21#[cfg(not(feature = "ocr-prep"))]
22mod assemble;
23mod dp_lines;
24#[cfg(feature = "ml")]
25pub mod enrich;
26// Public so sibling crates (e.g. docling-rag's ONNX embedder) can route their
27// own `ort` sessions through the same `DOCLING_RS_EP` selection.
28#[cfg(feature = "ml")]
29pub mod ep;
30pub mod layout;
31#[cfg(feature = "ml")]
32mod mets;
33#[cfg(feature = "ml")]
34mod ocr;
35#[cfg(feature = "ocr-prep")]
36pub mod ocr_prep;
37#[cfg(feature = "ml")]
38mod orient;
39pub mod pdfium_backend;
40#[cfg(feature = "ml")]
41pub mod quality;
42mod reading_order;
43// Pure-Rust region resampling (page→1024px box-average, crop→448 bilinear) —
44// available to the browser TableFormer path (#157 stage 3), not just `ml`.
45#[cfg(feature = "ocr-prep")]
46pub mod resample;
47#[cfg(feature = "ocr-prep")]
48pub mod scanned;
49// Built-in standard-14 font metrics for the pure-Rust text parser (#187) —
50// no feature gate: the wasm/pdf-text path needs them like the native one.
51mod std14;
52#[cfg(feature = "ml")]
53pub mod tableformer;
54pub mod textparse;
55#[cfg(feature = "ocr-prep")]
56pub mod tf_core;
57// docling's TableFormer cell matcher — pure Rust, shared with the browser
58// TableFormer path (#157 stage 3).
59#[cfg(feature = "ocr-prep")]
60pub mod tf_match;
61pub mod timing;
62
63#[cfg(feature = "ml")]
64use std::collections::BTreeMap;
65use std::fmt;
66#[cfg(feature = "ml")]
67use std::sync::mpsc::{sync_channel, Receiver};
68#[cfg(feature = "ml")]
69use std::sync::{Arc, Mutex};
70
71use docling_core::DoclingDocument;
72// The env-knob helpers only gate ML-pipeline diagnostics and tuning; the
73// pure text-layer (wasm) build has no call sites.
74#[cfg(feature = "ml")]
75use docling_core::Node;
76#[cfg(feature = "ml")]
77use docling_core::{debug_log, env};
78
79#[cfg(feature = "ml")]
80pub use mets::{convert_mets_gbs, convert_mets_gbs_with_options};
81#[cfg(feature = "ml")]
82pub use ocr::OcrLang;
83#[cfg(feature = "ml")]
84pub use pdfium_backend::PdfDocument;
85pub use pdfium_backend::{PdfPage, TextCell};
86
87/// Errors from the PDF backend. Detailed and surfaced (never silently skipped).
88#[derive(Debug)]
89pub enum PdfError {
90    /// pdfium failed to bind, open, or read the document.
91    Pdfium(String),
92    /// The layout ONNX model failed to load or run.
93    Layout(String),
94    /// The OCR ONNX model failed to load or run.
95    Ocr(String),
96}
97
98impl fmt::Display for PdfError {
99    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
100        match self {
101            PdfError::Pdfium(m) => write!(f, "pdf: pdfium error: {m}"),
102            PdfError::Layout(m) => write!(f, "pdf: {m}"),
103            PdfError::Ocr(m) => write!(f, "pdf: {m}"),
104        }
105    }
106}
107
108impl std::error::Error for PdfError {}
109
110#[cfg(feature = "ml")]
111impl From<pdfium_render::prelude::PdfiumError> for PdfError {
112    fn from(e: pdfium_render::prelude::PdfiumError) -> Self {
113        // A failed dlopen means pdfium was never installed — the #1 first-run
114        // failure after a bare `cargo install` (which ships no runtime
115        // assets). Say what to do instead of leaking the raw loader error.
116        if matches!(e, pdfium_render::prelude::PdfiumError::LoadLibraryError(_)) {
117            // The loader error pretty-prints over several lines; compact it.
118            let detail = e
119                .to_string()
120                .split_whitespace()
121                .collect::<Vec<_>>()
122                .join(" ");
123            return PdfError::Pdfium(format!(
124                "the pdfium library is not installed. PDF/image conversion needs \
125                 pdfium + the ONNX models: fetch both with \
126                 scripts/install/download_dependencies.sh from a docling.rs \
127                 checkout (https://github.com/docling-project/docling.rs), or \
128                 point PDFIUM_DYNAMIC_LIB_PATH at a directory containing the \
129                 pdfium library. A digital PDF's embedded text layer converts \
130                 without either in no-OCR mode (CLI: --no-ocr). Declarative \
131                 formats (DOCX, HTML, Markdown, …) never need them. [{detail}]"
132            ));
133        }
134        PdfError::Pdfium(e.to_string())
135    }
136}
137
138/// Convert a PDF's **embedded text layer only** — no pdfium, no ONNX, no
139/// threads: the pure-Rust content-stream parser ([`textparse`]) feeds the same
140/// orphan-region assembly the `no_ocr` pipeline flag uses, so text-layer PDFs
141/// come out identical to `--no-ocr` (flat, line-grouped paragraphs in reading
142/// order; no headings/lists/tables/pictures, and no hyperlink recovery).
143///
144/// This is the only conversion entry compiled without the `ml` feature (it is
145/// what a wasm32 build runs). A scanned/image-only PDF (no embedded text
146/// layer) yields an empty document rather than an error, same as `no_ocr` —
147/// callers can detect that and fall back to an OCR-capable build.
148pub fn convert_text_layer(bytes: &[u8], name: &str) -> Result<DoclingDocument, PdfError> {
149    convert_text_layer_pages(bytes, name, None)
150}
151
152/// [`convert_text_layer`] restricted to a **1-based inclusive** page window
153/// (issue #80's `--pages`); `None` converts everything. The window is
154/// validated the same way as [`Pipeline::pages`]: `first <= last`, 1-based,
155/// and it must select at least one existing page.
156pub fn convert_text_layer_pages(
157    bytes: &[u8],
158    name: &str,
159    pages: Option<(usize, usize)>,
160) -> Result<DoclingDocument, PdfError> {
161    if let Some((first, last)) = pages {
162        if first == 0 || last < first {
163            return Err(PdfError::Pdfium(format!(
164                "invalid page range {first}-{last} (pages are 1-based, first <= last)"
165            )));
166        }
167    }
168    let mut doc = DoclingDocument::new(name);
169    let mut total = 0usize;
170    let parsed = textparse::pdf_text_pages(bytes);
171    // A vestigial layer (a few typed-in form fields over scanned pages) is not
172    // the document's text: return the empty document, which callers already
173    // report as "no text layer" — so an OCR-capable caller falls back to OCR
174    // instead of proudly extracting thirteen characters.
175    if textparse::text_layer_is_vestigial(&parsed) {
176        return Ok(doc);
177    }
178    for (i, page) in parsed.into_iter().enumerate() {
179        total += 1;
180        if let Some((first, last)) = pages {
181            if i + 1 < first || i + 1 > last {
182                continue;
183            }
184        }
185        let mut regions = Vec::new();
186        assemble::add_orphan_regions(&mut regions, &page.cells);
187        let table_rows = vec![None; regions.len()];
188        let enrich_out = vec![None; regions.len()];
189        let (mut nodes, links) = assemble::assemble_page(&page, regions, &table_rows, &enrich_out);
190        assemble::stamp_page_no(&mut nodes, i + 1);
191        doc.nodes.extend(nodes);
192        doc.links.extend(links);
193    }
194    if let Some((first, last)) = pages {
195        if first > total {
196            return Err(PdfError::Pdfium(format!(
197                "page range {first}-{last} is outside the document ({total} page(s))"
198            )));
199        }
200    }
201    assemble::merge_continuations(&mut doc.nodes);
202    Ok(doc)
203}
204
205/// Threads ONNX inference may use, capped by `DOCLING_RS_PDF_THREADS` if set.
206/// Defaults to the available parallelism (ort otherwise picks a low number).
207#[cfg(feature = "ml")]
208pub(crate) fn intra_threads() -> usize {
209    if let Some(n) = env::parse::<usize>("DOCLING_RS_PDF_THREADS").filter(|&n| n > 0) {
210        return n;
211    }
212    std::thread::available_parallelism()
213        .map(|n| n.get())
214        .unwrap_or(1)
215}
216
217#[cfg(feature = "ml")]
218/// True when `DOCLING_RS_FP32` forces the full-precision models even where
219/// an INT8 variant sits next to the fp32 default.
220pub(crate) fn fp32_forced() -> bool {
221    env::flag("DOCLING_RS_FP32")
222}
223
224#[cfg(feature = "ml")]
225/// Should the int8 model defaults be skipped in favor of fp32? Either the
226/// user said so (`DOCLING_RS_FP32`), or a GPU execution provider is selected
227/// (#74) — the int8 exports are QDQ graphs calibrated for CPU kernels and
228/// only conformance-validated there. An explicit `DOCLING_*_ONNX` path
229/// override still wins over this at every call site.
230pub(crate) fn prefer_fp32() -> bool {
231    fp32_forced() || ep::prefers_fp32()
232}
233
234#[cfg(feature = "ml")]
235/// Resolve a default (CWD-relative) asset path — the shared chain in
236/// [`docling_core::assets`]: CWD, then next to the executable and one level
237/// above it (the `scripts/install/install.sh` layout).
238pub(crate) fn resolve_asset(rel: &str) -> String {
239    docling_core::assets::resolve(rel)
240}
241
242/// One resolved runtime asset — which file a stage would load right now,
243/// given the CWD, the env overrides and the int8/fp32 preference.
244#[cfg(feature = "ml")]
245#[derive(Debug, Clone)]
246pub struct ModelEntry {
247    /// Pipeline stage, e.g. `layout`, `tableformer.decoder`, `ocr.rec`.
248    pub stage: &'static str,
249    /// The resolved path (absolute or CWD-relative, as it will be opened).
250    pub path: String,
251    /// Whether the file exists right now.
252    pub found: bool,
253    /// File size in bytes (0 when missing) — enough to tell an int8 quant
254    /// from an fp32 graph, or a stale model from a re-published one, at a
255    /// glance without hashing gigabytes per request.
256    pub bytes: u64,
257}
258
259/// Resolve the whole runtime model set **without loading anything** — the
260/// exact selection each stage performs at load time (layout honors the
261/// int8/fp32 preference, TableFormer its decoder ranking, OCR the language
262/// pair), plus the pdfium library. docling-serve exposes this at
263/// `/v1/config` and logs it at startup, so "the server picked up different
264/// models" is one `curl` away instead of a mystery of dissolved tables.
265/// Resolution is CWD-relative with an exe-dir fallback, so the answer can
266/// legitimately differ between two working directories.
267#[cfg(feature = "ml")]
268pub fn model_inventory() -> Vec<ModelEntry> {
269    fn entry(stage: &'static str, path: String) -> ModelEntry {
270        let meta = std::fs::metadata(&path).ok();
271        ModelEntry {
272            stage,
273            found: meta.is_some(),
274            bytes: meta.map(|m| m.len()).unwrap_or(0),
275            path,
276        }
277    }
278    let (enc, dec, bbx) = tableformer::resolved_paths();
279    let (rec, dict) = ocr::resolve_rec_pair(ocr::OcrLang::from_env());
280    let pdfium =
281        env::nonempty("PDFIUM_DYNAMIC_LIB_PATH").unwrap_or_else(|| resolve_asset(".pdfium/lib"));
282    vec![
283        entry(
284            "layout",
285            model_path(
286                "DOCLING_LAYOUT_ONNX",
287                ".models/layout_heron.onnx",
288                ".models/layout_heron_int8.onnx",
289            ),
290        ),
291        entry("tableformer.encoder", enc),
292        entry("tableformer.decoder", dec),
293        entry("tableformer.bbox", bbx),
294        entry("ocr.rec", rec),
295        entry("ocr.dict", dict),
296        entry("pdfium", pdfium),
297    ]
298}
299
300/// Resolve a model path: an explicit env override always wins; otherwise the
301/// INT8 variant of the default path when it exists on disk (the quantized
302/// models are conformance-validated — see docs/PDF_CONFORMANCE.md — and load/run
303/// markedly faster on CPU), unless `DOCLING_RS_FP32` opts back into full
304/// precision; else the fp32 default.
305#[cfg(feature = "ml")]
306pub(crate) fn model_path(key: &str, fp32_default: &str, int8_default: &str) -> String {
307    if let Some(p) = env::nonempty(key) {
308        return p;
309    }
310    if !prefer_fp32() {
311        let p = resolve_asset(int8_default);
312        if std::path::Path::new(&p).exists() {
313            return p;
314        }
315    }
316    resolve_asset(fp32_default)
317}
318
319/// Decode a standalone image with hard resource limits. A crafted image can
320/// declare enormous dimensions in a few-KB file; `image::load_from_memory`
321/// then tries to allocate the full pixel buffer (e.g. 60000×60000 → ~10 GB),
322/// and allocation failure aborts the whole process, bypassing the per-request
323/// panic catch. The 256 MiB alloc / 30000-px caps below turn that into a
324/// recoverable decode error instead. `DOCLING_RS_MAX_IMAGE_PIXELS` overrides
325/// the per-side pixel cap for the rare legitimately-huge scan.
326///
327/// Gated on `ml`: the only callers (`convert_image`, the METS backend) are
328/// ML-only, and the `image` crate is an `ml`-feature dependency — the
329/// text-layer wasm build has neither.
330#[cfg(feature = "ml")]
331pub(crate) fn decode_image_limited(bytes: &[u8]) -> Result<image::RgbImage, PdfError> {
332    let max_side: u32 = env::parse("DOCLING_RS_MAX_IMAGE_PIXELS").unwrap_or(30_000);
333    decode_image_with_max_side(bytes, max_side)
334}
335
336#[cfg(feature = "ml")]
337fn decode_image_with_max_side(bytes: &[u8], max_side: u32) -> Result<image::RgbImage, PdfError> {
338    use image::ImageReader;
339    use std::io::Cursor;
340
341    let mut limits = image::Limits::default();
342    limits.max_image_width = Some(max_side);
343    limits.max_image_height = Some(max_side);
344    limits.max_alloc = Some(256 * 1024 * 1024);
345
346    let mut reader = ImageReader::new(Cursor::new(bytes))
347        .with_guessed_format()
348        .map_err(|e| PdfError::Pdfium(format!("image: {e}")))?;
349    reader.limits(limits);
350    Ok(reader
351        .decode()
352        .map_err(|e| PdfError::Pdfium(format!("image: {e}")))?
353        .into_rgb8())
354}
355
356#[cfg(feature = "ml")]
357/// One page's assembled output: typed nodes plus the page's hyperlinks (kept
358/// separate so pages processed out of order can be stitched back in page
359/// order) and its confidence scores (#183).
360type PageOut = (
361    Vec<Node>,
362    Vec<(String, String)>,
363    docling_core::confidence::PageConfidence,
364);
365
366#[cfg(feature = "ml")]
367/// The pool-wide TableFormer slot: one instance shared by every worker, loaded
368/// lazily on the first table region any worker sees. Tables appear on a
369/// minority of pages, so per-worker copies mostly multiplied ~0.4 GB of
370/// weights+arenas by the pool size for nothing; a single shared instance keeps
371/// the peak flat regardless of pool width, and a table's structure prediction
372/// is independent of which worker runs it, so output is byte-identical. The
373/// mutex serialises concurrent tables — the shared instance is loaded with the
374/// full intra-op thread budget to compensate (one wide TableFormer instead of
375/// several narrow ones).
376enum TfSlot {
377    /// Not attempted yet (no table seen so far).
378    Unloaded,
379    /// Load attempted, graphs absent — geometric fallback (warned once).
380    Missing,
381    Ready(tableformer::TableFormer),
382}
383
384#[cfg(feature = "ml")]
385type SharedTables = Arc<Mutex<TfSlot>>;
386
387#[cfg(feature = "ml")]
388/// The same lazy shared-slot pattern for the (rarer still) enrichment models:
389/// one instance per pipeline, loaded on the first region that needs it.
390enum EnrichSlot<T> {
391    Unloaded,
392    /// Load attempted, model files absent — enrichment skipped (warned once).
393    Missing,
394    Ready(T),
395}
396
397#[cfg(feature = "ml")]
398type SharedClassifier = Arc<Mutex<EnrichSlot<enrich::PictureClassifier>>>;
399#[cfg(feature = "ml")]
400type SharedCodeFormula = Arc<Mutex<EnrichSlot<enrich::CodeFormula>>>;
401
402#[cfg(feature = "ml")]
403/// The opt-in enrichment passes, mirroring docling's `PdfPipelineOptions`
404/// flags (`do_picture_classification`, `do_code_enrichment`,
405/// `do_formula_enrichment`). All off by default.
406#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
407pub struct EnrichmentOptions {
408    /// Classify each picture with DocumentFigureClassifier (26 classes).
409    pub picture_classification: bool,
410    /// Rewrite code blocks (and detect their language) with CodeFormulaV2.
411    pub code: bool,
412    /// Decode display formulas to LaTeX with CodeFormulaV2.
413    pub formula: bool,
414}
415
416#[cfg(feature = "ml")]
417impl EnrichmentOptions {
418    fn any(&self) -> bool {
419        self.picture_classification || self.code || self.formula
420    }
421}
422
423#[cfg(feature = "ml")]
424/// The layout model's input for a page: the docling-exact scale-1.0 page
425/// image when the renderer produced one, else the legacy stretch of the 2×
426/// bitmap (browser / METS paths) — see [`layout::LayoutSrc`]. Public so the
427/// diagnostic examples feed [`layout::LayoutModel::predict`] the same input
428/// the pipeline does.
429pub fn layout_src(page: &PdfPage) -> layout::LayoutSrc<'_> {
430    match &page.image_layout {
431        Some(img) => layout::LayoutSrc::PageImage(img),
432        None => layout::LayoutSrc::Raw(&page.image),
433    }
434}
435
436#[cfg(feature = "ml")]
437/// A self-contained set of the per-page models (layout, OCR). Each parallel
438/// page-worker owns its own `Worker` so inference runs concurrently without
439/// sharing an ONNX session (`ort`'s `Session::run` is `&mut self`); only the
440/// rarely-hit TableFormer is shared (see [`TfSlot`]).
441struct Worker {
442    /// `None` when `no_ocr` skips layout entirely — no model load, no inference.
443    layout: Option<layout::LayoutModel>,
444    ocr: Option<ocr::OcrModel>,
445    /// Shared TableFormer slot; `None` when `no_table_former`/`no_ocr` skip it.
446    tables: Option<SharedTables>,
447    /// Shared enrichment slots; `None` unless the corresponding flag is on.
448    classifier: Option<SharedClassifier>,
449    code_formula: Option<SharedCodeFormula>,
450    enrich: EnrichmentOptions,
451    /// Skip layout, OCR, and TableFormer; reconstruct text purely from the PDF's
452    /// embedded text layer. See [`Pipeline::no_ocr`].
453    no_ocr: bool,
454    /// Discard the embedded text layer and OCR every page. See
455    /// [`Pipeline::force_full_page_ocr`].
456    force_full_page_ocr: bool,
457    /// Keep text-panel pictures as pictures instead of demoting them to
458    /// paragraphs. See [`Pipeline::no_text_panels`].
459    no_text_panels: bool,
460    /// Which recognition model [`Self::ocr`] loads. See [`Pipeline::ocr_lang`].
461    ocr_lang: ocr::OcrLang,
462}
463
464#[cfg(feature = "ml")]
465impl Worker {
466    #[allow(clippy::too_many_arguments)] // mirrors the Pipeline's option set
467    fn load(
468        intra: usize,
469        tables: Option<SharedTables>,
470        enrich_slots: (Option<SharedClassifier>, Option<SharedCodeFormula>),
471        enrich: EnrichmentOptions,
472        no_ocr: bool,
473        force_full_page_ocr: bool,
474        no_text_panels: bool,
475        ocr_lang: ocr::OcrLang,
476    ) -> Result<Self, PdfError> {
477        Ok(Self {
478            layout: if no_ocr {
479                None
480            } else {
481                Some(layout::LayoutModel::load_with(intra).map_err(PdfError::Layout)?)
482            },
483            ocr: None,
484            tables,
485            classifier: enrich_slots.0,
486            code_formula: enrich_slots.1,
487            enrich,
488            no_ocr,
489            force_full_page_ocr,
490            no_text_panels,
491            ocr_lang,
492        })
493    }
494
495    /// Run layout (+ OCR for cell-less pages) + TableFormer and assemble page `n`
496    /// into its nodes and links. Pure given the page (mutates only the worker's
497    /// lazily-loaded OCR model), so it is safe to run concurrently across pages.
498    fn process(&mut self, n: usize, page: &mut PdfPage) -> Result<PageOut, PdfError> {
499        if self.no_ocr {
500            // Fastest path: no layout/OCR/TableFormer inference at all. The PDF's
501            // embedded text cells (if any) become flat, line-grouped paragraphs in
502            // reading order via the same orphan-region machinery that normally
503            // rescues text the detector missed — here it rescues *all* of it.
504            // Pages with no embedded text layer (scanned/image-only) yield nothing;
505            // convert those without `no_ocr`.
506            let parse = quality::parse_score(&page.cells);
507            let mut regions = Vec::new();
508            assemble::add_orphan_regions(&mut regions, &page.cells);
509            let table_rows = vec![None; regions.len()];
510            let enrich_out = vec![None; regions.len()];
511            let conf = quality::page_confidence(parse, &regions, &[]);
512            let (nodes, links) = timing::timed("assemble_page", || {
513                assemble::assemble_page(page, regions, &table_rows, &enrich_out)
514            });
515            return Ok((nodes, links, conf));
516        }
517        self.normalize_orientation(n, page)?;
518        let regions = timing::timed("layout.predict", || {
519            self.layout
520                .as_mut()
521                .expect("layout model loaded unless no_ocr")
522                .predict(layout_src(page), page.width, page.height)
523        })
524        .map_err(|e| PdfError::Layout(format!("page {}: {e}", n + 1)))?;
525        self.finish_page(n, page, regions)
526    }
527
528    /// Content-based orientation normalization (#225), before any inference:
529    /// a physically rotated scan (sideways phone photo, landscape-fed sheet)
530    /// has `/Rotate 0`, so the metadata pass in `extract_page` never fires and
531    /// layout+OCR would read a sideways raster. Only pages with no text layer
532    /// at all are probed (a digital page's raster is upright by construction,
533    /// and its cells — not its pixels — carry the text); the detected angle
534    /// composes with any `/Rotate` normalization through the same
535    /// [`PdfPage::unrotate`] + display-space assembly mapping. Detection is
536    /// evidence-gated and degrades to a no-op — see [`orient`].
537    fn normalize_orientation(&mut self, n: usize, page: &mut PdfPage) -> Result<(), PdfError> {
538        let scanned =
539            page.cells.is_empty() && page.word_cells.is_empty() && page.code_cells.is_empty();
540        if self.no_ocr || !scanned || page.image.width() <= 1 || !orient::enabled() {
541            return Ok(());
542        }
543        if self.ocr.is_none() {
544            self.ocr = Some(ocr::OcrModel::load(self.ocr_lang).map_err(PdfError::Ocr)?);
545        }
546        let deg = timing::timed("orient.detect", || {
547            orient::detect(&page.image, self.ocr.as_mut().unwrap())
548        });
549        if deg != 0 {
550            debug_log!(
551                "docling-pdf: page {}: content rotated {deg}° in the raster; \
552                 un-rotating before layout/OCR",
553                n + 1
554            );
555            page.unrotate(deg);
556        }
557        Ok(())
558    }
559
560    /// Layout-detect a whole batch of pages with one inference call (issue #73),
561    /// then run each page's remaining stages (OCR / TableFormer / enrichment /
562    /// assembly) per page. Index-aligned with `items`; a layout failure fails
563    /// every page in the batch (they shared the one inference call).
564    fn process_batch(&mut self, items: &mut [(usize, PdfPage)]) -> Vec<Result<PageOut, PdfError>> {
565        if self.no_ocr {
566            // No layout model to batch — the text-layer-only path is per page.
567            return items
568                .iter_mut()
569                .map(|(n, page)| {
570                    let n = *n;
571                    self.process(n, page)
572                })
573                .collect();
574        }
575        // Orientation-normalize every scanned page before the shared layout
576        // call — the batched inference must see upright bitmaps too (#225).
577        for (n, page) in items.iter_mut() {
578            let n = *n;
579            if let Err(e) = self.normalize_orientation(n, page) {
580                // Model-load failure — every page in the batch needs the same
581                // model, so they all fail alike (mirrors the layout-error arm).
582                let msg = e.to_string();
583                return items
584                    .iter()
585                    .map(|_| Err(PdfError::Ocr(msg.clone())))
586                    .collect();
587            }
588        }
589        let inputs: Vec<(layout::LayoutSrc<'_>, f32, f32)> = items
590            .iter()
591            .map(|(_, page)| (layout_src(page), page.width, page.height))
592            .collect();
593        let batched = timing::timed("layout.predict", || {
594            self.layout
595                .as_mut()
596                .expect("layout model loaded unless no_ocr")
597                .predict_batch(&inputs)
598        });
599        match batched {
600            Ok(all) => items
601                .iter_mut()
602                .zip(all)
603                .map(|((n, page), regions)| self.finish_page(*n, page, regions))
604                .collect(),
605            Err(e) => items
606                .iter()
607                .map(|(n, _)| Err(PdfError::Layout(format!("page {}: {e}", n + 1))))
608                .collect(),
609        }
610    }
611
612    /// Everything after layout detection: per-label confidence thresholds,
613    /// overlap resolution, orphan-text recovery, OCR for cell-less pages,
614    /// TableFormer, enrichment, and page assembly.
615    fn finish_page(
616        &mut self,
617        n: usize,
618        page: &mut PdfPage,
619        regions: Vec<layout::Region>,
620    ) -> Result<PageOut, PdfError> {
621        // Force-OCR is exactly "pretend the text layer is not there": clear
622        // every cell kind the extractors produced before anything reads them,
623        // and the ordinary no-text-layer machinery below — full-page OCR,
624        // OCR-fed TableFormer matching — takes over unchanged. (`no_ocr` wins
625        // when both are set, mirroring docling, where `force_full_page_ocr`
626        // is a sub-option of `do_ocr`; the no-ocr path never reaches here.)
627        // Done here rather than in `process` so the batched layout path
628        // (`process_batch` → `finish_page`) honors the flag too.
629        // Parse quality is scored on the extracted text layer before force-OCR
630        // discards it (docling's page-preprocessing stage runs before OCR too,
631        // so its parse_score also reflects the original text layer).
632        let parse = quality::parse_score(&page.cells);
633        // Recognition confidences of every OCR'd cell on this page → ocr_score.
634        let mut ocr_confs: Vec<f32> = Vec::new();
635        if self.force_full_page_ocr {
636            page.cells.clear();
637            page.code_cells.clear();
638            page.word_cells.clear();
639        }
640        // Quant-robustness guard: the default int8 layout graph keeps its
641        // confidences near the 0.5 label thresholds, and a different CPU's
642        // quantized kernels can flip a whole page's detections under them —
643        // tables and paragraphs then dissolve into orphan one-liners while the
644        // same build converts the page perfectly elsewhere. When a dense
645        // digital page ends up with detections covering almost none of its
646        // text cells, re-run that one page on the fp32 graph (lazy-loaded,
647        // auto-int8 selection only) and keep whichever detections cover more.
648        let mut regions = regions;
649        if !page.cells.is_empty() {
650            let thresholded = |rs: &[layout::Region]| -> Vec<layout::Region> {
651                rs.iter()
652                    .filter(|r| r.score >= layout::label_threshold(r.label))
653                    .cloned()
654                    .collect()
655            };
656            let text_cells = page
657                .cells
658                .iter()
659                .filter(|c| !c.text.trim().is_empty())
660                .count();
661            let cov = assemble::layout_cell_coverage(&thresholded(&regions), &page.cells);
662            if text_cells >= 15 && cov < 0.5 {
663                let retry = self
664                    .layout
665                    .as_mut()
666                    .expect("layout model loaded unless no_ocr")
667                    .predict_fp32_fallback(layout_src(page), page.width, page.height)
668                    .map_err(|e| PdfError::Layout(format!("page {}: {e}", n + 1)))?;
669                if let Some(retry) = retry {
670                    let cov2 = assemble::layout_cell_coverage(&thresholded(&retry), &page.cells);
671                    if cov2 > cov {
672                        debug_log!(
673                            "docling-pdf: page {}: int8 layout covered {:.0}% of the text \
674                             cells; the fp32 retry covers {:.0}% — using it",
675                            n + 1,
676                            cov * 100.0,
677                            cov2 * 100.0
678                        );
679                        regions = retry;
680                    }
681                }
682            }
683        }
684        // docling's LayoutPostprocessor drops each detection below its label's
685        // confidence threshold (stricter than the 0.3 base the predictor keeps),
686        // before any overlap resolution. This removes the low-confidence tables /
687        // pictures / list-items that otherwise double-emit or mis-classify.
688        if env::flag("DOCLING_RS_DEBUG_REGIONS") {
689            for r in &regions {
690                eprintln!(
691                    "DBG raw {} {:.2} [{:.0},{:.0},{:.0},{:.0}]",
692                    r.label, r.score, r.l, r.t, r.r, r.b
693                );
694            }
695        }
696        regions.retain(|r| r.score >= layout::label_threshold(r.label));
697        // docling's same-label picture dedup runs on the thresholded
698        // detections, before overlap resolution: a figure proposed both whole
699        // and as sub-panels collapses to one box (see `dedup_pictures`).
700        assemble::dedup_pictures(&mut regions);
701        // Resolve overlapping detections once, before OCR.
702        let mut regions = assemble::resolve(regions);
703        // Emit text the detector missed as orphan text regions (docling parity).
704        assemble::add_orphan_regions(&mut regions, &page.cells);
705        // Drop phantom empty low-confidence picture boxes (docling parity).
706        assemble::drop_false_pictures(&mut regions, &page.cells, page.width, page.height);
707        // A regular region fully inside a surviving table/index/picture is that
708        // special's child (a cell / in-figure label), not a separate block —
709        // remove it so it isn't emitted twice (docling parity).
710        assemble::drop_contained_regulars(&mut regions);
711        // No text layer → recognise text from the page image via OCR.
712        let ocred = page.cells.is_empty();
713        if ocred {
714            if self.ocr.is_none() {
715                self.ocr = Some(ocr::OcrModel::load(self.ocr_lang).map_err(PdfError::Ocr)?);
716            }
717            let cells = timing::timed("ocr.page", || {
718                self.ocr
719                    .as_mut()
720                    .unwrap()
721                    .ocr_page(&page.image, &regions, page.scale)
722            })
723            .map_err(|e| PdfError::Ocr(format!("page {}: {e}", n + 1)))?;
724            ocr_confs.extend(cells.iter().map(|(_, conf)| conf));
725            page.cells = cells.into_iter().map(|(cell, _)| cell).collect();
726            // Table interiors carry no words yet: region-scoped OCR skips
727            // table labels, and a scanned page has no pdfium text layer — so
728            // TableFormer's cell matcher got an empty word list and the table
729            // dissolved (#173). Recognize the table regions' word crops
730            // (mirroring the browser scanned path): `word_cells` feeds the
731            // matcher, and the same cells join `cells` so the geometric
732            // fallback and the table's region text see them too.
733            if regions.iter().any(|r| assemble::is_table_like(r.label)) {
734                let words = timing::timed("ocr.table_words", || {
735                    self.ocr
736                        .as_mut()
737                        .unwrap()
738                        .ocr_table_words(&page.image, &regions, page.scale)
739                })
740                .map_err(|e| PdfError::Ocr(format!("page {}: {e}", n + 1)))?;
741                ocr_confs.extend(words.iter().map(|(_, conf)| conf));
742                let words: Vec<_> = words.into_iter().map(|(cell, _)| cell).collect();
743                page.cells.extend(words.iter().cloned());
744                page.word_cells = words;
745            }
746        }
747        // Region-scoped OCR skips `picture` interiors, and a digital page's
748        // text layer cannot see into an embedded raster either — so a figure
749        // that is really a text box (terms-and-conditions exported as an
750        // image) lost its words on every page kind. Python docling OCRs the
751        // bitmap-covered areas of *every* page — even digital ones — once they
752        // exceed `bitmap_area_threshold` (5 % of the page); the browser paths
753        // already do. Recognize the big text-less crops here too; the panel
754        // demotion / orphan recovery below place the lines.
755        let mut pic_cells: Vec<pdfium_backend::TextCell> = Vec::new();
756        {
757            let page_area = (page.width * page.height).max(1.0);
758            let has_text = |r: &layout::Region| {
759                page.cells.iter().any(|c| {
760                    let ca = ((c.r - c.l) * (c.b - c.t)).max(1.0);
761                    let ix = (r.r.min(c.r) - r.l.max(c.l)).max(0.0);
762                    let iy = (r.b.min(c.b) - r.t.max(c.t)).max(0.0);
763                    !c.text.trim().is_empty() && ix * iy / ca > 0.5
764                })
765            };
766            // A captioned picture can never demote to a text panel (see
767            // recover_text_panels), and on digital pages its speculative OCR
768            // would be discarded anyway — don't pay for it.
769            let captioned = |r: &layout::Region| {
770                regions.iter().any(|c| {
771                    c.label == "caption"
772                        && c.r.min(r.r) - c.l.max(r.l) > 0.0
773                        && ((c.t >= r.b && c.t - r.b <= 25.0) || (r.t >= c.b && r.t - c.b <= 25.0))
774                })
775            };
776            let bare: Vec<layout::Region> = regions
777                .iter()
778                .filter(|r| {
779                    r.label == "picture"
780                        && (r.r - r.l) * (r.b - r.t) / page_area >= 0.05
781                        && !has_text(r)
782                        && (ocred || !captioned(r))
783                })
784                .map(|r| layout::Region {
785                    label: "text",
786                    ..r.clone()
787                })
788                .collect();
789            if !bare.is_empty() {
790                if self.ocr.is_none() {
791                    self.ocr = Some(ocr::OcrModel::load(self.ocr_lang).map_err(PdfError::Ocr)?);
792                }
793                let scored = timing::timed("ocr.pictures", || {
794                    self.ocr
795                        .as_mut()
796                        .unwrap()
797                        .ocr_page(&page.image, &bare, page.scale)
798                })
799                .map_err(|e| PdfError::Ocr(format!("page {}: {e}", n + 1)))?;
800                // Speculative in-picture OCR counts toward ocr_score only on
801                // OCR'd pages, where the recognized lines actually join the
802                // output; on a digital page they may be discarded below.
803                if ocred {
804                    ocr_confs.extend(scored.iter().map(|(_, conf)| conf));
805                }
806                pic_cells = scored.into_iter().map(|(cell, _)| cell).collect();
807                page.cells.extend(pic_cells.iter().cloned());
808            }
809        }
810        let cells_before_pic_ocr = page.cells.len() - pic_cells.len();
811        // A "picture" that is really a colored text panel — dense, wide,
812        // multi-line — reads out as paragraphs instead of shipping as pixels;
813        // sparse in-picture text (a chart's labels) keeps the crop and stays
814        // inside it as the picture's silent children (docling parity, #200).
815        // `no_text_panels` (#173) opts out entirely for image-extraction
816        // workflows.
817        if !self.no_text_panels {
818            assemble::recover_text_panels(&mut regions, &page.cells);
819        }
820        // On an OCR'd page, in-picture text that did NOT demote its picture
821        // mostly stays silent, exactly as in docling: its postprocess step
822        // "Remove regular clusters that are included in wrappers" walks
823        // SPECIAL_TYPES — which includes PICTURE — so an orphan text cluster
824        // >80 % contained in a kept picture becomes that picture's child and
825        // never reaches the serializer. Only border-straddlers (≤80 %
826        // containment) survive as text. Emitting *everything* here used to
827        // splice a chart's OCR'd axis ticks into the body text right next to
828        // the image chunk (#200) — so the orphan pass places the recognized
829        // lines, then the same containment drop that handled the first wave
830        // re-runs to swallow the in-picture ones.
831        if ocred && !pic_cells.is_empty() {
832            // Pictures (and wrappers) no longer count as claimers (#165), so
833            // the plain orphan pass places the recognized lines directly.
834            assemble::add_orphan_regions(&mut regions, &pic_cells);
835            assemble::drop_contained_regulars(&mut regions);
836        } else if !ocred && !pic_cells.is_empty() {
837            // Digital page, picture kept: its speculative OCR cells must not
838            // linger in the text-cell set (they were appended at the tail).
839            let kept: Vec<layout::Region> = regions
840                .iter()
841                .filter(|r| r.label == "picture")
842                .cloned()
843                .collect();
844            let tail = page.cells.split_off(cells_before_pic_ocr);
845            page.cells.extend(tail.into_iter().filter(|c| {
846                !kept.iter().any(|r| {
847                    let ca = ((c.r - c.l) * (c.b - c.t)).max(1.0);
848                    let ix = (r.r.min(c.r) - r.l.max(c.l)).max(0.0);
849                    let iy = (r.b.min(c.b) - r.t.max(c.t)).max(0.0);
850                    ix * iy / ca > 0.5
851                })
852            }));
853        }
854        // A text-less *table* detected inside a picture on a digital page — a
855        // screenshot of a table (2203's Figure 10) — has no text layer and no
856        // scanned-path OCR to feed it, so its grid used to serialize empty and
857        // the whole element vanished. docling OCRs bitmap-covered areas on
858        // every page kind and its table cluster collects those cells; mirror
859        // the scanned path for exactly these tables: recognize word crops and
860        // feed them to the TableFormer matcher and the cell set.
861        if !ocred {
862            let has_text = |t: &layout::Region| {
863                page.cells.iter().any(|c| {
864                    let ca = ((c.r - c.l) * (c.b - c.t)).max(1.0);
865                    let ix = (t.r.min(c.r) - t.l.max(c.l)).max(0.0);
866                    let iy = (t.b.min(c.b) - t.t.max(c.t)).max(0.0);
867                    !c.text.trim().is_empty() && ix * iy / ca > 0.5
868                })
869            };
870            let in_picture = |t: &layout::Region| {
871                regions.iter().any(|r| {
872                    r.label == "picture" && {
873                        let ta = ((t.r - t.l) * (t.b - t.t)).max(1.0);
874                        let ix = (r.r.min(t.r) - r.l.max(t.l)).max(0.0);
875                        let iy = (r.b.min(t.b) - r.t.max(t.t)).max(0.0);
876                        ix * iy / ta > 0.5
877                    }
878                })
879            };
880            let pic_tables: Vec<layout::Region> = regions
881                .iter()
882                .filter(|t| assemble::is_table_like(t.label) && !has_text(t) && in_picture(t))
883                .cloned()
884                .collect();
885            if !pic_tables.is_empty() {
886                if self.ocr.is_none() {
887                    self.ocr = Some(ocr::OcrModel::load(self.ocr_lang).map_err(PdfError::Ocr)?);
888                }
889                let words = timing::timed("ocr.table_words", || {
890                    self.ocr
891                        .as_mut()
892                        .unwrap()
893                        .ocr_table_words(&page.image, &pic_tables, page.scale)
894                })
895                .map_err(|e| PdfError::Ocr(format!("page {}: {e}", n + 1)))?;
896                ocr_confs.extend(words.iter().map(|(_, conf)| conf));
897                let words: Vec<_> = words.into_iter().map(|(cell, _)| cell).collect();
898                page.cells.extend(words.iter().cloned());
899                page.word_cells.extend(words);
900            }
901        }
902        // TableFormer structure per table region (else geometric fallback). The
903        // shared slot is only locked (and lazily loaded) when the page actually
904        // has a table, so table-free documents never pay for TableFormer at all.
905        let mut table_rows: Vec<Option<Vec<Vec<String>>>> = vec![None; regions.len()];
906        if let Some(slot) = self.tables.as_ref() {
907            if regions.iter().any(|r| assemble::is_table_like(r.label)) {
908                timing::timed("tableformer", || {
909                    let mut guard = slot.lock().unwrap();
910                    if matches!(*guard, TfSlot::Unloaded) {
911                        // Full intra-op width: tables serialise on this mutex, so
912                        // the one instance gets the whole thread budget.
913                        *guard = match tableformer::TableFormer::load_with(intra_threads()) {
914                            Some(tf) => TfSlot::Ready(tf),
915                            None => TfSlot::Missing,
916                        };
917                    }
918                    if let TfSlot::Ready(tf) = &mut *guard {
919                        for (i, r) in regions.iter().enumerate() {
920                            if assemble::is_table_like(r.label) {
921                                table_rows[i] = tf.predict_table_rows(
922                                    &page.image,
923                                    [r.l, r.t, r.r, r.b],
924                                    &page.word_cells,
925                                );
926                            }
927                        }
928                    }
929                });
930            }
931        }
932        if env::flag("DOCLING_RS_DEBUG_REGIONS") {
933            for (i, r) in regions.iter().enumerate() {
934                eprintln!(
935                    "DBG final {} {:.2} [{:.0},{:.0},{:.0},{:.0}] rows={:?}",
936                    r.label,
937                    r.score,
938                    r.l,
939                    r.t,
940                    r.r,
941                    r.b,
942                    table_rows[i]
943                        .as_ref()
944                        .map(|t| (t.len(), t.first().map(|r| r.len())))
945                );
946            }
947            eprintln!(
948                "DBG cells={} words={}",
949                page.cells.len(),
950                page.word_cells.len()
951            );
952        }
953        // Enrichment passes (opt-in): DocumentPictureClassifier over picture
954        // regions, CodeFormulaV2 over code/formula regions. Same shared-slot
955        // shape as TableFormer — one lazily-loaded instance per pipeline, only
956        // ever locked when a page actually has a matching region.
957        let mut enrich_out: Vec<Option<assemble::Enrichment>> = vec![None; regions.len()];
958        if let Some(slot) = self.classifier.as_ref() {
959            if regions.iter().any(|r| r.label == "picture") {
960                timing::timed("picture_classifier", || {
961                    let mut guard = slot.lock().unwrap();
962                    if matches!(*guard, EnrichSlot::Unloaded) {
963                        *guard = match enrich::PictureClassifier::load_with(intra_threads()) {
964                            Some(m) => EnrichSlot::Ready(m),
965                            None => EnrichSlot::Missing,
966                        };
967                    }
968                    if let EnrichSlot::Ready(model) = &mut *guard {
969                        for (i, r) in regions.iter().enumerate() {
970                            if r.label != "picture" {
971                                continue;
972                            }
973                            let Some(crop) = assemble::crop_region_scaled(
974                                page,
975                                [r.l, r.t, r.r, r.b],
976                                enrich::CLASSIFIER_SCALE,
977                            ) else {
978                                continue;
979                            };
980                            match model.classify(&crop) {
981                                Ok(classes) => {
982                                    enrich_out[i] =
983                                        Some(assemble::Enrichment::PictureClasses(classes));
984                                }
985                                Err(e) => eprintln!("docling-pdf: page {}: {e}", n + 1),
986                            }
987                        }
988                    }
989                });
990            }
991        }
992        if let Some(slot) = self.code_formula.as_ref() {
993            let wants = |label: &str| {
994                (label == "code" && self.enrich.code) || (label == "formula" && self.enrich.formula)
995            };
996            if regions.iter().any(|r| wants(r.label)) {
997                timing::timed("code_formula", || {
998                    let mut guard = slot.lock().unwrap();
999                    if matches!(*guard, EnrichSlot::Unloaded) {
1000                        *guard = match enrich::CodeFormula::load_with(intra_threads()) {
1001                            Some(m) => EnrichSlot::Ready(m),
1002                            None => EnrichSlot::Missing,
1003                        };
1004                    }
1005                    if let EnrichSlot::Ready(model) = &mut *guard {
1006                        for (i, r) in regions.iter().enumerate() {
1007                            if !wants(r.label) {
1008                                continue;
1009                            }
1010                            // docling crops the postprocessed cluster box — the
1011                            // union of the region's text cells, not the raw
1012                            // detector box — expanded by 18% per side, at
1013                            // ~120 dpi.
1014                            let [bl, bt, br, bb] = assemble::region_cell_bbox(r, &page.cells)
1015                                .unwrap_or([r.l, r.t, r.r, r.b]);
1016                            let (w, h) = (br - bl, bb - bt);
1017                            let ex = enrich::CODE_FORMULA_EXPANSION;
1018                            let bbox = [bl - w * ex, bt - h * ex, br + w * ex, bb + h * ex];
1019                            let Some(crop) = assemble::crop_region_scaled(
1020                                page,
1021                                bbox,
1022                                enrich::CODE_FORMULA_SCALE,
1023                            ) else {
1024                                continue;
1025                            };
1026                            let kind = if r.label == "code" {
1027                                enrich::CodeFormulaKind::Code
1028                            } else {
1029                                enrich::CodeFormulaKind::Formula
1030                            };
1031                            match model.predict(&crop, kind) {
1032                                Ok(text) => {
1033                                    enrich_out[i] = Some(match kind {
1034                                        enrich::CodeFormulaKind::Code => {
1035                                            let (code, language) =
1036                                                enrich::extract_code_language(&text);
1037                                            assemble::Enrichment::Code {
1038                                                language,
1039                                                text: code,
1040                                            }
1041                                        }
1042                                        enrich::CodeFormulaKind::Formula => {
1043                                            assemble::Enrichment::Formula { latex: text }
1044                                        }
1045                                    });
1046                                }
1047                                Err(e) => eprintln!("docling-pdf: page {}: {e}", n + 1),
1048                            }
1049                        }
1050                    }
1051                });
1052            }
1053        }
1054        // Score the final region set (docling assigns layout_score over the
1055        // postprocessed clusters — the same set assemble_page consumes).
1056        let conf = quality::page_confidence(parse, &regions, &ocr_confs);
1057        let (nodes, links) = timing::timed("assemble_page", || {
1058            assemble::assemble_page(page, regions, &table_rows, &enrich_out)
1059        });
1060        Ok((nodes, links, conf))
1061    }
1062}
1063
1064#[cfg(feature = "ml")]
1065/// Per-worker ONNX intra-op threads. The layout model is memory-bandwidth bound,
1066/// so on a typical machine two threads per worker (sharing one in-cache copy of
1067/// the weights) extracts more throughput than one fat model or many single-thread
1068/// workers. `DOCLING_RS_PDF_INTRA` overrides for per-machine tuning.
1069fn pdf_intra() -> usize {
1070    if let Some(n) = env::parse::<usize>("DOCLING_RS_PDF_INTRA").filter(|&n| n > 0) {
1071        return n;
1072    }
1073    if intra_threads() >= 2 {
1074        2
1075    } else {
1076        1
1077    }
1078}
1079
1080#[cfg(feature = "ml")]
1081/// How many page-workers to spin up for a multi-page PDF. `DOCLING_RS_PDF_WORKERS`
1082/// overrides; otherwise size the pool so `workers × intra ≈ cores`, capped at 4 so
1083/// a worst-case pool holds a bounded amount of model memory (~0.4 GB per worker)
1084/// and does not oversaturate the memory bus with model-weight traffic.
1085fn pdf_worker_count() -> usize {
1086    if let Some(n) = env::parse::<usize>("DOCLING_RS_PDF_WORKERS").filter(|&n| n > 0) {
1087        return n;
1088    }
1089    (intra_threads() / pdf_intra()).clamp(1, 4)
1090}
1091
1092#[cfg(feature = "ml")]
1093/// Max pages a worker layout-detects with one batched inference call (issue
1094/// #73). Workers drain the work channel opportunistically up to this size —
1095/// whatever is already rendered gets batched, so batching never *waits* for
1096/// pages and adds no latency when rendering is the bottleneck.
1097///
1098/// Default: 4 on 8+ cores, 1 (per-page) below. Measured on a 4-core box the
1099/// batch only adds cache pressure and costs pipeline overlap (2 workers × 2
1100/// threads: 8.1 s/conv at batch=1 vs 9.3 s at batch=4 on the 9-page
1101/// 2206.01062 fixture); the single-session amortization it buys needs the
1102/// wider thread budget of a many-core machine. Output is bit-identical at
1103/// every batch size, so this is purely a throughput knob.
1104/// `DOCLING_RS_PDF_LAYOUT_BATCH` overrides; `1` restores per-page inference.
1105fn pdf_layout_batch() -> usize {
1106    env::parse::<usize>("DOCLING_RS_PDF_LAYOUT_BATCH")
1107        .filter(|&n| n > 0)
1108        .unwrap_or_else(|| if intra_threads() >= 8 { 4 } else { 1 })
1109}
1110
1111#[cfg(feature = "ml")]
1112/// Minimum page count before a PDF is worth the parallel worker pool. Below this,
1113/// the serial primary (running its model on every core) is faster than fanning out
1114/// — the helper pool's one-time model-load cost only pays off once enough pages
1115/// share it. `DOCLING_RS_PDF_PARALLEL_MIN` overrides.
1116fn pdf_parallel_min() -> usize {
1117    env::parse::<usize>("DOCLING_RS_PDF_PARALLEL_MIN")
1118        .filter(|&n| n > 0)
1119        .unwrap_or(6)
1120}
1121
1122#[cfg(feature = "ml")]
1123/// A reusable PDF pipeline. The **primary** worker runs its models on every core,
1124/// so a single-page / small / image / METS input is converted at full intra-op
1125/// speed with no pool to load. A document with enough pages instead fans out
1126/// across a **pool** of narrower workers processed concurrently. Both load lazily
1127/// and are cached for reuse, so a one-shot conversion only pays for what it uses.
1128pub struct Pipeline {
1129    /// Full-intra worker for the serial path; loaded on first serial use.
1130    primary: Option<Worker>,
1131    /// Narrower workers (≈cores/`target_workers` threads each) for the parallel
1132    /// path; loaded on first multi-page use and cached.
1133    pool: Vec<Worker>,
1134    /// The single TableFormer instance every worker shares (see [`TfSlot`]).
1135    tables: SharedTables,
1136    /// The shared enrichment-model slots (same pattern as [`TfSlot`]).
1137    classifier: SharedClassifier,
1138    code_formula: SharedCodeFormula,
1139    /// Desired pool size for multi-page documents.
1140    target_workers: usize,
1141    /// Page count at/above which the parallel pool is worth its load cost.
1142    parallel_min: usize,
1143    /// Skip loading/running TableFormer; table regions fall back to geometric
1144    /// reconstruction. See [`Pipeline::no_table_former`].
1145    no_table_former: bool,
1146    /// Skip layout, OCR, and TableFormer entirely. See [`Pipeline::no_ocr`].
1147    no_ocr: bool,
1148    /// OCR every page even when it carries a text layer. See
1149    /// [`Pipeline::force_full_page_ocr`].
1150    force_full_page_ocr: bool,
1151    /// Never demote text-panel pictures. See [`Pipeline::no_text_panels`].
1152    no_text_panels: bool,
1153    /// Opt-in enrichment passes. See [`Pipeline::enrichments`].
1154    enrich: EnrichmentOptions,
1155    /// 1-based inclusive page window to convert. See [`Pipeline::pages`].
1156    page_range: Option<(usize, usize)>,
1157    /// OCR recognition language. See [`Pipeline::ocr_lang`].
1158    ocr_lang: ocr::OcrLang,
1159    /// Optional per-page progress hook `(done, selected_total)`, invoked after
1160    /// each page finishes on both the serial and parallel buffered paths. Set
1161    /// by the CLI batch mode for dot-progress; `None` costs nothing.
1162    progress: Option<Arc<dyn Fn(usize, usize) + Send + Sync>>,
1163}
1164
1165#[cfg(feature = "ml")]
1166impl Pipeline {
1167    /// Construct the pipeline. Models load lazily on first use (full-intra primary
1168    /// for serial inputs, the helper pool for multi-page PDFs), so nothing is
1169    /// loaded that a given document doesn't need.
1170    pub fn new() -> Result<Self, PdfError> {
1171        Ok(Self {
1172            primary: None,
1173            pool: Vec::new(),
1174            tables: Arc::new(Mutex::new(TfSlot::Unloaded)),
1175            classifier: Arc::new(Mutex::new(EnrichSlot::Unloaded)),
1176            code_formula: Arc::new(Mutex::new(EnrichSlot::Unloaded)),
1177            target_workers: pdf_worker_count(),
1178            parallel_min: pdf_parallel_min(),
1179            no_table_former: false,
1180            no_ocr: false,
1181            force_full_page_ocr: false,
1182            no_text_panels: false,
1183            enrich: EnrichmentOptions::default(),
1184            page_range: None,
1185            ocr_lang: ocr::OcrLang::from_env(),
1186            progress: None,
1187        })
1188    }
1189
1190    /// Install (or clear) the per-page progress hook: called with
1191    /// `(pages_done, pages_selected)` after each page completes during
1192    /// [`convert`](Self::convert). Shared with the parallel workers, so the
1193    /// callback must be cheap and thread-safe.
1194    pub fn set_progress(&mut self, cb: Option<Arc<dyn Fn(usize, usize) + Send + Sync>>) {
1195        self.progress = cb;
1196    }
1197
1198    /// Convert only pages `first..=last` (**1-based**, like the page numbers a
1199    /// PDF viewer shows — issue #80's `--pages A-B`). Out-of-range pages are
1200    /// skipped before rasterization, so the cost is proportional to the window,
1201    /// not the document. `last` past the end of the document clamps; a window
1202    /// that selects no pages at all (`first` beyond the last page) is an error
1203    /// at convert time. `None` (the default) converts everything.
1204    pub fn pages(mut self, range: Option<(usize, usize)>) -> Self {
1205        self.page_range = range;
1206        self
1207    }
1208
1209    /// In-place variant of [`pages`](Self::pages) for a long-lived pipeline
1210    /// (e.g. docling-serve's warm instance) that applies a per-request window
1211    /// without rebuilding — unlike the model switches, the window is pure
1212    /// configuration. Set it before every conversion; it stays until changed.
1213    pub fn set_pages(&mut self, range: Option<(usize, usize)>) {
1214        self.page_range = range;
1215    }
1216
1217    /// OCR recognition language (see [`OcrLang`]): English by default, `ch`
1218    /// for the multilingual docling-conformance model. `None` keeps the
1219    /// process default (`DOCLING_RS_OCR_LANG`, else English). Set before the
1220    /// first conversion; for a warm pipeline use
1221    /// [`set_ocr_lang`](Self::set_ocr_lang).
1222    pub fn ocr_lang(mut self, lang: Option<ocr::OcrLang>) -> Self {
1223        self.set_ocr_lang(lang);
1224        self
1225    }
1226
1227    /// In-place variant of [`ocr_lang`](Self::ocr_lang) for a long-lived
1228    /// pipeline (docling-serve's warm instance). Unlike the page window this
1229    /// is a *model* switch: any worker whose cached recognition model was
1230    /// loaded for a different language drops it, to be lazily reloaded on the
1231    /// next OCR-needing page (cheap — the rec models are ~10 MB).
1232    pub fn set_ocr_lang(&mut self, lang: Option<ocr::OcrLang>) {
1233        let lang = lang.unwrap_or_else(ocr::OcrLang::from_env);
1234        self.ocr_lang = lang;
1235        for worker in self.primary.iter_mut().chain(self.pool.iter_mut()) {
1236            if worker.ocr_lang != lang {
1237                worker.ocr_lang = lang;
1238                worker.ocr = None;
1239            }
1240        }
1241    }
1242
1243    /// Resolve the configured 1-based window against a page count into the
1244    /// 0-based inclusive form the backend walks, validating it selects at
1245    /// least one existing page.
1246    fn resolve_range(&self, total: usize) -> Result<Option<(usize, usize)>, PdfError> {
1247        let Some((first, last)) = self.page_range else {
1248            return Ok(None);
1249        };
1250        if first == 0 || last < first {
1251            return Err(PdfError::Pdfium(format!(
1252                "invalid page range {first}-{last} (pages are 1-based, first <= last)"
1253            )));
1254        }
1255        if first > total {
1256            return Err(PdfError::Pdfium(format!(
1257                "page range {first}-{last} is outside the document ({total} page(s))"
1258            )));
1259        }
1260        Ok(Some((first - 1, last.min(total) - 1)))
1261    }
1262
1263    /// Enable the opt-in enrichment passes (docling's
1264    /// `do_picture_classification` / `do_code_enrichment` /
1265    /// `do_formula_enrichment`). Each enabled pass lazily loads its model on
1266    /// the first matching region; a missing model warns once and is skipped.
1267    /// Set before the first conversion (no effect on already-loaded workers).
1268    pub fn enrichments(mut self, opts: EnrichmentOptions) -> Self {
1269        self.enrich = opts;
1270        self
1271    }
1272
1273    /// Skip loading and running the TableFormer table-structure model. Table
1274    /// regions still get emitted, but reconstructed geometrically from cell
1275    /// positions instead of via the ONNX model's predicted structure — faster
1276    /// (no model load, no per-table inference) at the cost of table fidelity.
1277    /// No effect if a worker is already loaded; set this before the first
1278    /// conversion.
1279    pub fn no_table_former(mut self, disable: bool) -> Self {
1280        self.no_table_former = disable;
1281        self
1282    }
1283
1284    /// Keep every detected `picture` region as a picture. By default an
1285    /// *uncaptioned* picture that reads like a dense, uniform text panel (a
1286    /// terms-and-conditions box exported as an image) is demoted into
1287    /// paragraphs (#157); a chart the layout mislabels can still trip that
1288    /// heuristic on scanned pages, and image-extraction workflows may simply
1289    /// want every crop — this flag disables the demotion entirely (#173).
1290    /// No effect on already-loaded workers; set before the first conversion.
1291    pub fn no_text_panels(mut self, disable: bool) -> Self {
1292        self.no_text_panels = disable;
1293        self
1294    }
1295
1296    /// Skip layout detection, OCR, and TableFormer entirely — no model load, no
1297    /// inference of any kind. The PDF's embedded text cells are grouped by line
1298    /// and emitted as plain paragraphs in reading order: no headings, lists,
1299    /// tables, code blocks, or pictures, since that structure comes from the
1300    /// layout model. The fastest possible PDF path, but pages with no embedded
1301    /// text layer (scanned/image-only PDFs) yield no text at all — convert those
1302    /// without this flag. Implies `no_table_former`. No effect if a worker is
1303    /// already loaded; set this before the first conversion.
1304    pub fn no_ocr(mut self, disable: bool) -> Self {
1305        self.no_ocr = disable;
1306        self
1307    }
1308
1309    /// OCR every page from its rendered image even when the page carries an
1310    /// embedded text layer — docling's `force_full_page_ocr`. The escape hatch
1311    /// for text layers that exist but lie: broken encodings, subset fonts with
1312    /// garbage mappings, a scanned form with a few typed-in fields. Ignored
1313    /// when [`no_ocr`](Self::no_ocr) is set, mirroring docling (there
1314    /// `force_full_page_ocr` is a sub-option of `do_ocr`).
1315    pub fn force_full_page_ocr(mut self, force: bool) -> Self {
1316        self.force_full_page_ocr = force;
1317        self
1318    }
1319
1320    /// The shared TableFormer slot handed to each worker, or `None` when the
1321    /// pipeline options skip TableFormer entirely.
1322    fn tables_slot(&self) -> Option<SharedTables> {
1323        if self.no_table_former || self.no_ocr {
1324            None
1325        } else {
1326            Some(Arc::clone(&self.tables))
1327        }
1328    }
1329
1330    /// The shared enrichment slots for a worker (`None` per model unless its
1331    /// flag is on; `no_ocr` skips layout, so there are no regions to enrich).
1332    fn enrich_slots(&self) -> (Option<SharedClassifier>, Option<SharedCodeFormula>) {
1333        if self.no_ocr || !self.enrich.any() {
1334            return (None, None);
1335        }
1336        (
1337            self.enrich
1338                .picture_classification
1339                .then(|| Arc::clone(&self.classifier)),
1340            (self.enrich.code || self.enrich.formula).then(|| Arc::clone(&self.code_formula)),
1341        )
1342    }
1343
1344    /// Eagerly load the models (the full-intra serial worker: layout + OCR, and
1345    /// the shared TableFormer unless disabled) so the first conversion doesn't pay
1346    /// the load cost. Idempotent; respects `no_ocr` / `no_table_former` (with
1347    /// `no_ocr` there is nothing to load). The docling.rs analogue of docling's
1348    /// `DocumentConverter.initialize_pipeline`.
1349    pub fn warm_up(&mut self) -> Result<(), PdfError> {
1350        self.primary()?;
1351        Ok(())
1352    }
1353
1354    /// The full-intra serial worker, loaded on first use.
1355    fn primary(&mut self) -> Result<&mut Worker, PdfError> {
1356        if self.primary.is_none() {
1357            self.primary = Some(Worker::load(
1358                intra_threads(),
1359                self.tables_slot(),
1360                self.enrich_slots(),
1361                self.enrich,
1362                self.no_ocr,
1363                self.force_full_page_ocr,
1364                self.no_text_panels,
1365                self.ocr_lang,
1366            )?);
1367        }
1368        Ok(self.primary.as_mut().unwrap())
1369    }
1370
1371    /// Convert a PDF (bytes) to a [`DoclingDocument`]. A document with fewer than
1372    /// `parallel_min` pages (or a pool size of 1) streams through the full-intra
1373    /// primary; a larger one renders on this thread (pdfium is not thread-safe) and
1374    /// fans the pages out across the worker pool, reassembled in page order so the
1375    /// output is byte-identical to the serial path.
1376    pub fn convert(
1377        &mut self,
1378        bytes: &[u8],
1379        password: Option<&str>,
1380        name: &str,
1381    ) -> Result<DoclingDocument, PdfError> {
1382        let pages = pdfium_backend::page_count(bytes, password)?;
1383        let range = self.resolve_range(pages)?;
1384        // Serial vs parallel is decided by the pages actually converted: a
1385        // 3-page window over a 500-page PDF should not pay the pool load.
1386        let selected = range.map_or(pages, |(a, b)| b - a + 1);
1387        let doc = if self.target_workers >= 2 && selected >= self.parallel_min {
1388            self.convert_parallel(bytes, password, name, range, selected)?
1389        } else {
1390            self.convert_serial(bytes, password, name, range, selected)?
1391        };
1392        timing::report();
1393        Ok(doc)
1394    }
1395
1396    /// Stream pages one at a time through the primary worker — render → process →
1397    /// drop — so the document holds ~one page bitmap (~5 MB) at a time.
1398    fn convert_serial(
1399        &mut self,
1400        bytes: &[u8],
1401        password: Option<&str>,
1402        name: &str,
1403        range: Option<(usize, usize)>,
1404        selected: usize,
1405    ) -> Result<DoclingDocument, PdfError> {
1406        let mut doc = DoclingDocument::new(name);
1407        let mut confs = std::collections::BTreeMap::new();
1408        let render_image = !self.no_ocr;
1409        let progress = self.progress.clone();
1410        let mut done = 0usize;
1411        let worker = self.primary()?;
1412        pdfium_backend::for_each_page(
1413            bytes,
1414            password,
1415            render_image,
1416            range,
1417            |n, _total, mut page| {
1418                let (mut nodes, links, conf) = worker.process(n, &mut page)?;
1419                assemble::stamp_page_no(&mut nodes, n + 1);
1420                doc.nodes.extend(nodes);
1421                doc.links.extend(links);
1422                confs.insert(n + 1, conf);
1423                if let Some(cb) = &progress {
1424                    done += 1;
1425                    cb(done, selected);
1426                }
1427                Ok::<(), PdfError>(())
1428            },
1429        )?;
1430        assemble::merge_continuations(&mut doc.nodes);
1431        doc.confidence = Some(docling_core::ConfidenceReport::from_pages(confs));
1432        Ok(doc)
1433    }
1434
1435    /// Render pages serially on this thread (pdfium) and process them in parallel
1436    /// across the worker pool. A bounded channel applies backpressure so only a
1437    /// handful of page bitmaps are resident at once; results carry their page
1438    /// index and are reassembled in order, so the output is byte-identical to the
1439    /// serial path.
1440    fn convert_parallel(
1441        &mut self,
1442        bytes: &[u8],
1443        password: Option<&str>,
1444        name: &str,
1445        range: Option<(usize, usize)>,
1446        selected: usize,
1447    ) -> Result<DoclingDocument, PdfError> {
1448        self.ensure_pool()?;
1449        let progress = self.progress.clone();
1450        let pages_done = std::sync::atomic::AtomicUsize::new(0);
1451        let n_workers = self.pool.len();
1452        let render_image = !self.no_ocr;
1453        let layout_batch = pdf_layout_batch();
1454        // Bound sized so every worker can accumulate a full layout batch while
1455        // rendering stays ahead (and never below the pre-#73 render-ahead of
1456        // two pages per worker); still a hard cap on resident page bitmaps.
1457        let (work_tx, work_rx) = sync_channel::<(usize, PdfPage)>(n_workers * layout_batch.max(2));
1458        let work_rx: Arc<Mutex<Receiver<(usize, PdfPage)>>> = Arc::new(Mutex::new(work_rx));
1459        let results: Arc<Mutex<Vec<(usize, PageOut)>>> = Arc::new(Mutex::new(Vec::new()));
1460        let first_err: Arc<Mutex<Option<PdfError>>> = Arc::new(Mutex::new(None));
1461
1462        // Move the pool into the scope so each worker gets an exclusive `&mut`.
1463        let mut workers = std::mem::take(&mut self.pool);
1464        std::thread::scope(|s| {
1465            for worker in workers.iter_mut() {
1466                let work_rx = Arc::clone(&work_rx);
1467                let results = Arc::clone(&results);
1468                let first_err = Arc::clone(&first_err);
1469                let progress = progress.clone();
1470                let pages_done = &pages_done;
1471                s.spawn(move || loop {
1472                    // Hold the receiver lock only for the recv (plus a non-blocking
1473                    // drain up to the layout batch size); release before the (long)
1474                    // per-page work so other workers can pull concurrently.
1475                    let mut batch = Vec::new();
1476                    {
1477                        let rx = work_rx.lock().unwrap();
1478                        match rx.recv() {
1479                            Ok(item) => {
1480                                batch.push(item);
1481                                while batch.len() < layout_batch {
1482                                    match rx.try_recv() {
1483                                        Ok(item) => batch.push(item),
1484                                        Err(_) => break,
1485                                    }
1486                                }
1487                            }
1488                            Err(_) => break,
1489                        }
1490                    }
1491                    let outs = worker.process_batch(&mut batch);
1492                    for ((idx, _), out) in batch.iter().zip(outs) {
1493                        match out {
1494                            Ok(out) => {
1495                                results.lock().unwrap().push((*idx, out));
1496                                if let Some(cb) = &progress {
1497                                    let d = pages_done
1498                                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
1499                                        + 1;
1500                                    cb(d, selected);
1501                                }
1502                            }
1503                            Err(e) => {
1504                                let mut slot = first_err.lock().unwrap();
1505                                if slot.is_none() {
1506                                    *slot = Some(e);
1507                                }
1508                            }
1509                        }
1510                    }
1511                });
1512            }
1513            // Render on this thread and feed the workers; backpressure blocks here
1514            // when the channel is full. Dropping `work_tx` afterwards signals the
1515            // workers (recv → Err) to finish.
1516            let render = pdfium_backend::for_each_page(
1517                bytes,
1518                password,
1519                render_image,
1520                range,
1521                |i, _total, page| {
1522                    work_tx
1523                        .send((i, page))
1524                        .map_err(|_| PdfError::Pdfium("page-worker channel closed".into()))
1525                },
1526            );
1527            drop(work_tx);
1528            if let Err(e) = render {
1529                let mut slot = first_err.lock().unwrap();
1530                if slot.is_none() {
1531                    *slot = Some(e);
1532                }
1533            }
1534        });
1535        // Threads have joined; restore the pool for the next conversion.
1536        self.pool = workers;
1537
1538        if let Some(e) = first_err.lock().unwrap().take() {
1539            return Err(e);
1540        }
1541        let mut results = Arc::try_unwrap(results)
1542            .unwrap_or_else(|arc| Mutex::new(arc.lock().unwrap().clone()))
1543            .into_inner()
1544            .unwrap();
1545        results.sort_by_key(|(idx, _)| *idx);
1546        let mut doc = DoclingDocument::new(name);
1547        let mut confs = std::collections::BTreeMap::new();
1548        for (idx, (mut nodes, links, conf)) in results {
1549            assemble::stamp_page_no(&mut nodes, idx + 1);
1550            doc.nodes.extend(nodes);
1551            doc.links.extend(links);
1552            confs.insert(idx + 1, conf);
1553        }
1554        assemble::merge_continuations(&mut doc.nodes);
1555        doc.confidence = Some(docling_core::ConfidenceReport::from_pages(confs));
1556        Ok(doc)
1557    }
1558
1559    /// Convert a PDF in **streaming** mode: `emit` is called with each finalized,
1560    /// in-document-order batch of nodes (and that span's recovered links) as pages
1561    /// complete, so a caller can serialize Markdown page by page instead of waiting
1562    /// for the whole document. The batches are exactly the buffered [`convert`]'s
1563    /// nodes, split at safe block boundaries by [`assemble::StreamAssembler`] — the
1564    /// parallel path reorders pages back into document order before emitting, so
1565    /// the output is identical regardless of worker scheduling.
1566    ///
1567    /// `emit` runs on the calling thread (never a worker), so it needn't be `Send`
1568    /// and its backpressure throttles the whole pipeline. Returning `Err` from
1569    /// `emit` aborts the conversion with that error.
1570    pub fn convert_streaming<F>(
1571        &mut self,
1572        bytes: &[u8],
1573        password: Option<&str>,
1574        name: &str,
1575        emit: F,
1576    ) -> Result<(), PdfError>
1577    where
1578        F: FnMut(Vec<Node>, Vec<(String, String)>) -> Result<(), PdfError>,
1579    {
1580        let _ = name; // page nodes carry no name; the caller owns the document name.
1581        let pages = pdfium_backend::page_count(bytes, password)?;
1582        let range = self.resolve_range(pages)?;
1583        let selected = range.map_or(pages, |(a, b)| b - a + 1);
1584        let r = if self.target_workers >= 2 && selected >= self.parallel_min {
1585            self.convert_streaming_parallel(bytes, password, range, emit)
1586        } else {
1587            self.convert_streaming_serial(bytes, password, range, emit)
1588        };
1589        timing::report();
1590        r
1591    }
1592
1593    /// Serial streaming: render → process → emit, one page at a time, holding back
1594    /// only the tail that might still merge into the next page.
1595    fn convert_streaming_serial<F>(
1596        &mut self,
1597        bytes: &[u8],
1598        password: Option<&str>,
1599        range: Option<(usize, usize)>,
1600        mut emit: F,
1601    ) -> Result<(), PdfError>
1602    where
1603        F: FnMut(Vec<Node>, Vec<(String, String)>) -> Result<(), PdfError>,
1604    {
1605        let mut asm = assemble::StreamAssembler::new();
1606        let render_image = !self.no_ocr;
1607        let worker = self.primary()?;
1608        pdfium_backend::for_each_page(
1609            bytes,
1610            password,
1611            render_image,
1612            range,
1613            |n, _total, mut page| {
1614                // Confidence is dropped on the streaming path: the report is
1615                // only complete once every page has run, which defeats
1616                // page-by-page emission — buffered `convert` carries it.
1617                let (nodes, links, _conf) = worker.process(n, &mut page)?;
1618                emit(asm.push(nodes), links)
1619            },
1620        )?;
1621        emit(asm.finish(), Vec::new())
1622    }
1623
1624    /// Parallel streaming: pages render serially on a dedicated thread (pdfium is
1625    /// not thread-safe) and process across the worker pool; results carry their
1626    /// page index and are reordered on the calling thread into a
1627    /// [`assemble::StreamAssembler`], which emits each page in document order as
1628    /// soon as its predecessors have arrived. Bounded channels keep only a handful
1629    /// of pages resident and let `emit`'s backpressure reach the renderer.
1630    fn convert_streaming_parallel<F>(
1631        &mut self,
1632        bytes: &[u8],
1633        password: Option<&str>,
1634        range: Option<(usize, usize)>,
1635        mut emit: F,
1636    ) -> Result<(), PdfError>
1637    where
1638        F: FnMut(Vec<Node>, Vec<(String, String)>) -> Result<(), PdfError>,
1639    {
1640        self.ensure_pool()?;
1641        let n_workers = self.pool.len();
1642        let render_image = !self.no_ocr;
1643        let layout_batch = pdf_layout_batch();
1644        // Bound sized so every worker can accumulate a full layout batch while
1645        // rendering stays ahead (and never below the pre-#73 render-ahead of
1646        // two pages per worker); still a hard cap on resident page bitmaps.
1647        let (work_tx, work_rx) = sync_channel::<(usize, PdfPage)>(n_workers * layout_batch.max(2));
1648        let work_rx: Arc<Mutex<Receiver<(usize, PdfPage)>>> = Arc::new(Mutex::new(work_rx));
1649        // Workers and the renderer report here; the calling thread drains it in
1650        // page order. Bounded so workers block (bounding resident bitmaps) when the
1651        // consumer falls behind.
1652        let (res_tx, res_rx) = sync_channel::<Result<(usize, PageOut), PdfError>>(n_workers * 2);
1653
1654        let mut workers = std::mem::take(&mut self.pool);
1655        let mut asm = assemble::StreamAssembler::new();
1656        let mut first_err: Option<PdfError> = None;
1657
1658        std::thread::scope(|s| {
1659            // Workers: pull a batch of pages (whatever is already rendered, up
1660            // to the layout batch size), process it, report (index-tagged)
1661            // results.
1662            for worker in workers.iter_mut() {
1663                let work_rx = Arc::clone(&work_rx);
1664                let res_tx = res_tx.clone();
1665                s.spawn(move || 'outer: loop {
1666                    let mut batch = Vec::new();
1667                    {
1668                        let rx = work_rx.lock().unwrap();
1669                        match rx.recv() {
1670                            Ok(item) => {
1671                                batch.push(item);
1672                                while batch.len() < layout_batch {
1673                                    match rx.try_recv() {
1674                                        Ok(item) => batch.push(item),
1675                                        Err(_) => break,
1676                                    }
1677                                }
1678                            }
1679                            Err(_) => break,
1680                        }
1681                    }
1682                    let outs = worker.process_batch(&mut batch);
1683                    for ((idx, _), out) in batch.iter().zip(outs) {
1684                        if res_tx.send(out.map(|o| (*idx, o))).is_err() {
1685                            break 'outer; // consumer gone
1686                        }
1687                    }
1688                });
1689            }
1690            // Renderer: feed pages to the pool on its own thread (pdfium stays on a
1691            // single thread); report a render error through the same channel.
1692            {
1693                let res_tx = res_tx.clone();
1694                s.spawn(move || {
1695                    let render = pdfium_backend::for_each_page(
1696                        bytes,
1697                        password,
1698                        render_image,
1699                        range,
1700                        |i, _total, page| {
1701                            work_tx
1702                                .send((i, page))
1703                                .map_err(|_| PdfError::Pdfium("page-worker channel closed".into()))
1704                        },
1705                    );
1706                    drop(work_tx); // signal workers to finish
1707                    if let Err(e) = render {
1708                        let _ = res_tx.send(Err(e));
1709                    }
1710                });
1711            }
1712            // Drop our own sender so the channel closes once the threads finish.
1713            drop(res_tx);
1714
1715            // Collector (this thread): reorder into document order and emit.
1716            // With a page window, indices start at the window's first page.
1717            let mut buffer: BTreeMap<usize, PageOut> = BTreeMap::new();
1718            let mut next = range.map_or(0, |(first, _)| first);
1719            for msg in res_rx.iter() {
1720                match msg {
1721                    Err(e) => {
1722                        if first_err.is_none() {
1723                            first_err = Some(e);
1724                        }
1725                    }
1726                    Ok((idx, out)) => {
1727                        buffer.insert(idx, out);
1728                        if first_err.is_some() {
1729                            continue; // keep draining so the threads can exit
1730                        }
1731                        while let Some((nodes, links, _conf)) = buffer.remove(&next) {
1732                            if let Err(e) = emit(asm.push(nodes), links) {
1733                                first_err = Some(e);
1734                                break;
1735                            }
1736                            next += 1;
1737                        }
1738                    }
1739                }
1740            }
1741        });
1742        // Threads have joined; restore the pool for the next conversion.
1743        self.pool = workers;
1744
1745        if let Some(e) = first_err {
1746            return Err(e);
1747        }
1748        emit(asm.finish(), Vec::new())
1749    }
1750
1751    /// Lazily grow the pool to `target_workers`, loading the new workers
1752    /// concurrently (model load is mostly I/O + mmap, so N loads overlap to roughly
1753    /// one load's wall-time). Cached for reuse across documents.
1754    fn ensure_pool(&mut self) -> Result<(), PdfError> {
1755        let need = self.target_workers.saturating_sub(self.pool.len());
1756        if need == 0 {
1757            return Ok(());
1758        }
1759        let intra = pdf_intra();
1760        let no_ocr = self.no_ocr;
1761        let force = self.force_full_page_ocr;
1762        let ntp = self.no_text_panels;
1763        let ocr_lang = self.ocr_lang;
1764        let enrich = self.enrich;
1765        let tables = self.tables_slot();
1766        let enrich_slots = self.enrich_slots();
1767        let loaded: Vec<Result<Worker, PdfError>> = std::thread::scope(|s| {
1768            let handles: Vec<_> = (0..need)
1769                .map(|_| {
1770                    let tables = tables.clone();
1771                    let enrich_slots = enrich_slots.clone();
1772                    s.spawn(move || {
1773                        Worker::load(
1774                            intra,
1775                            tables,
1776                            enrich_slots,
1777                            enrich,
1778                            no_ocr,
1779                            force,
1780                            ntp,
1781                            ocr_lang,
1782                        )
1783                    })
1784                })
1785                .collect();
1786            handles.into_iter().map(|h| h.join().unwrap()).collect()
1787        });
1788        for w in loaded {
1789            self.pool.push(w?);
1790        }
1791        Ok(())
1792    }
1793
1794    /// Convert a standalone image (PNG/JPEG/TIFF/WebP/…) as a single page —
1795    /// docling routes images through the same layout+OCR pipeline as a PDF page.
1796    pub fn convert_image(&mut self, bytes: &[u8], name: &str) -> Result<DoclingDocument, PdfError> {
1797        let image = decode_image_limited(bytes)?;
1798        let (w, h) = image.dimensions();
1799        // The image is its own page rendered at 1 px per "point" (scale 1.0); a
1800        // standalone image has no text layer, so OCR supplies the cells.
1801        let page = PdfPage {
1802            width: w as f32,
1803            height: h as f32,
1804            scale: 1.0,
1805            cells: Vec::new(),
1806            code_cells: Vec::new(),
1807            word_cells: Vec::new(),
1808            // A standalone image *is* its own scale-1.0 page image, so the
1809            // layout model sees it through the docling-exact PIL kernel.
1810            image_layout: Some(image.clone()),
1811            image,
1812            links: Vec::new(),
1813            rotation: 0,
1814        };
1815        self.process_pages(vec![page], name)
1816    }
1817
1818    /// Run layout (+ OCR for cell-less pages) and assemble each already-rendered
1819    /// page (image / METS inputs, which are small and already materialised).
1820    fn process_pages(
1821        &mut self,
1822        mut pages: Vec<PdfPage>,
1823        name: &str,
1824    ) -> Result<DoclingDocument, PdfError> {
1825        let mut doc = DoclingDocument::new(name);
1826        let mut confs = std::collections::BTreeMap::new();
1827        let worker = self.primary()?;
1828        for (n, page) in pages.iter_mut().enumerate() {
1829            let (mut nodes, links, conf) = worker.process(n, page)?;
1830            assemble::stamp_page_no(&mut nodes, n + 1);
1831            doc.nodes.extend(nodes);
1832            doc.links.extend(links);
1833            confs.insert(n + 1, conf);
1834        }
1835        assemble::merge_continuations(&mut doc.nodes);
1836        doc.confidence = Some(docling_core::ConfidenceReport::from_pages(confs));
1837        Ok(doc)
1838    }
1839}
1840
1841/// Number of pages in a PDF, without converting anything — what the CLI batch
1842/// mode prints in its per-document start line.
1843#[cfg(feature = "ml")]
1844pub fn page_count(bytes: &[u8], password: Option<&str>) -> Result<usize, PdfError> {
1845    Ok(pdfium_backend::page_count(bytes, password)?)
1846}
1847
1848#[cfg(feature = "ml")]
1849/// Convenience one-shot conversion (loads the pipeline per call). Errors are
1850/// detailed and surfaced (never silently skipped).
1851pub fn convert(
1852    bytes: &[u8],
1853    password: Option<&str>,
1854    name: &str,
1855) -> Result<DoclingDocument, PdfError> {
1856    convert_with_options(
1857        bytes,
1858        password,
1859        name,
1860        false,
1861        false,
1862        false,
1863        false,
1864        EnrichmentOptions::default(),
1865        None,
1866        None,
1867    )
1868}
1869
1870#[cfg(feature = "ml")]
1871/// Like [`convert`], but optionally skips loading/running TableFormer (see
1872/// [`Pipeline::no_table_former`]) and/or layout+OCR+TableFormer entirely (see
1873/// [`Pipeline::no_ocr`]), and/or enables the enrichment passes (see
1874/// [`Pipeline::enrichments`]).
1875// One positional per pipeline switch mirrors the Pipeline builder; growing
1876// past clippy's arity cap is the price of keeping this one-shot signature
1877// stable-ish instead of churning callers into an options struct mid-series.
1878#[allow(clippy::too_many_arguments)]
1879pub fn convert_with_options(
1880    bytes: &[u8],
1881    password: Option<&str>,
1882    name: &str,
1883    no_table_former: bool,
1884    no_ocr: bool,
1885    force_full_page_ocr: bool,
1886    no_text_panels: bool,
1887    enrich: EnrichmentOptions,
1888    pages: Option<(usize, usize)>,
1889    ocr_lang: Option<OcrLang>,
1890) -> Result<DoclingDocument, PdfError> {
1891    Pipeline::new()?
1892        .no_table_former(no_table_former)
1893        .no_ocr(no_ocr)
1894        .force_full_page_ocr(force_full_page_ocr)
1895        .no_text_panels(no_text_panels)
1896        .enrichments(enrich)
1897        .pages(pages)
1898        .ocr_lang(ocr_lang)
1899        .convert(bytes, password, name)
1900}
1901
1902#[cfg(feature = "ml")]
1903/// Convenience one-shot image conversion (loads the pipeline per call).
1904pub fn convert_image(bytes: &[u8], name: &str) -> Result<DoclingDocument, PdfError> {
1905    convert_image_with_options(
1906        bytes,
1907        name,
1908        false,
1909        false,
1910        false,
1911        EnrichmentOptions::default(),
1912        None,
1913    )
1914}
1915
1916#[cfg(feature = "ml")]
1917/// Like [`convert_image`], but optionally skips loading/running TableFormer (see
1918/// [`Pipeline::no_table_former`]) and/or layout+OCR+TableFormer entirely (see
1919/// [`Pipeline::no_ocr`]), and/or enables the enrichment passes.
1920pub fn convert_image_with_options(
1921    bytes: &[u8],
1922    name: &str,
1923    no_table_former: bool,
1924    no_ocr: bool,
1925    no_text_panels: bool,
1926    enrich: EnrichmentOptions,
1927    ocr_lang: Option<OcrLang>,
1928) -> Result<DoclingDocument, PdfError> {
1929    Pipeline::new()?
1930        .no_table_former(no_table_former)
1931        .no_ocr(no_ocr)
1932        .no_text_panels(no_text_panels)
1933        .enrichments(enrich)
1934        .ocr_lang(ocr_lang)
1935        .convert_image(bytes, name)
1936}
1937
1938#[cfg(feature = "ml")]
1939/// Convert pre-segmented pages (image + already-known text cells, e.g. METS/hOCR
1940/// scans) through the shared layout + assembly pipeline.
1941pub fn convert_pages(pages: Vec<PdfPage>, name: &str) -> Result<DoclingDocument, PdfError> {
1942    convert_pages_with_options(
1943        pages,
1944        name,
1945        false,
1946        false,
1947        false,
1948        EnrichmentOptions::default(),
1949    )
1950}
1951
1952#[cfg(feature = "ml")]
1953/// Like [`convert_pages`], but optionally skips loading/running TableFormer (see
1954/// [`Pipeline::no_table_former`]) and/or layout+OCR+TableFormer entirely (see
1955/// [`Pipeline::no_ocr`]), and/or enables the enrichment passes.
1956pub fn convert_pages_with_options(
1957    pages: Vec<PdfPage>,
1958    name: &str,
1959    no_table_former: bool,
1960    no_ocr: bool,
1961    no_text_panels: bool,
1962    enrich: EnrichmentOptions,
1963) -> Result<DoclingDocument, PdfError> {
1964    Pipeline::new()?
1965        .no_table_former(no_table_former)
1966        .no_text_panels(no_text_panels)
1967        .no_ocr(no_ocr)
1968        .enrichments(enrich)
1969        .process_pages(pages, name)
1970}
1971
1972#[cfg(feature = "ml")]
1973#[cfg(all(test, feature = "ml"))]
1974mod image_limit_tests {
1975    use super::decode_image_with_max_side;
1976
1977    /// A small valid PNG encoded via the `image` crate (robust vs. a hand-rolled
1978    /// byte literal).
1979    fn png_bytes(w: u32, h: u32) -> Vec<u8> {
1980        use std::io::Cursor;
1981        let img = image::RgbImage::new(w, h);
1982        let mut out = Vec::new();
1983        img.write_to(&mut Cursor::new(&mut out), image::ImageFormat::Png)
1984            .unwrap();
1985        out
1986    }
1987
1988    #[test]
1989    fn normal_image_decodes_under_the_cap() {
1990        let img = decode_image_with_max_side(&png_bytes(8, 8), 30_000).expect("8x8 decodes");
1991        assert_eq!(img.dimensions(), (8, 8));
1992    }
1993
1994    #[test]
1995    fn dimensions_over_the_cap_are_rejected_not_aborted() {
1996        // A per-side cap below the image's declared size must yield a
1997        // recoverable Err, never an allocation-abort — the mechanism that stops
1998        // a crafted image declaring 60000×60000 from OOM-killing the process.
1999        let r = decode_image_with_max_side(&png_bytes(8, 8), 4);
2000        assert!(
2001            r.is_err(),
2002            "decode must fail under the pixel cap, not abort"
2003        );
2004    }
2005}
2006
2007#[cfg(test)]
2008mod median_tests {
2009    #[test]
2010    fn median_of_empty_is_zero_not_a_panic() {
2011        // A crafted table can leave a row/column with zero matched cells; the
2012        // even-count branch would index values[0 - 1] and panic (→ remote crash
2013        // via docling-serve) without the empty guard.
2014        assert_eq!(super::tf_match::median_for_test(&mut []), 0.0);
2015        assert_eq!(super::tf_match::median_for_test(&mut [4.0, 2.0]), 3.0);
2016        assert_eq!(super::tf_match::median_for_test(&mut [5.0, 1.0, 3.0]), 3.0);
2017    }
2018}
2019
2020#[cfg(test)]
2021mod send_check {
2022    /// The Node bindings (`docling-node`) run a shared [`super::Pipeline`] on
2023    /// libuv worker threads (`Arc<Mutex<Pipeline>>`), which is only sound while
2024    /// `Pipeline: Send` holds — this fails to compile if a non-`Send` field
2025    /// (e.g. an `Rc` or a raw pdfium handle) ever lands in the pipeline.
2026    fn assert_send<T: Send>() {}
2027
2028    #[test]
2029    fn pipeline_is_send() {
2030        assert_send::<super::Pipeline>();
2031    }
2032}