Skip to content

Commit b48dffd

Browse files
artizclaude
andcommitted
fix(pdf): fp32 layout retry for pages the int8 quant loses; model inventory
Same build, same models, same bill: one machine's CPU converts it perfectly, another's dissolves page 1 into orphan one-liners — every table gone — while its CUDA build is fine. The default int8 layout graph keeps its confidences near the 0.5 label thresholds, and a different CPU's quantized kernels (AVX-VNNI vs AVX2, CUDA's mixed fallback) can flip a whole page's detections under them. Verified live: DOCLING_RS_FP32=1 fixed the page on the affected machine. The guard: when a dense digital page (≥15 text cells) ends up with detections covering <50% of its text cells, that one page re-runs on the fp32 graph — lazily loaded, only when the int8 file was auto-selected (explicit DOCLING_LAYOUT_ONNX / DOCLING_RS_FP32 choices are respected) — and whichever detections cover more win, with a warning naming both coverages. Coverage = fraction of non-empty cells some detection claims at docling's 0.2 intersection-over-self (assemble::layout_cell_coverage, unit-tested). The full corpus never trips the guard (0 firings) and stays byte-identical: 5/14 strict + 6/14 ws-normalized, same per-fixture diffs. Diagnostics so "wrong models" stops being a mystery: docling-pdf's new model_inventory() resolves what every stage would load right now (layout int8/fp32 pick, TableFormer's decoder ranking via the extracted tableformer::resolved_paths — shared with load so it can't drift — OCR language pair, pdfium), with found/size per file. docling re-exports it; docling-serve logs it at startup and serves it in /v1/config (documented in openapi.yaml), so two deployments converting differently is one curl diff away. Verified: with the threshold temporarily forced the fallback loads and converts (bill unchanged, 22 table rows); at the real threshold the bill and the whole groundtruth corpus never trigger it. Unit tests for the coverage metric; fmt/clippy clean. Signed-off-by: artiz <artem.kustikov@gmail.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EY5KAiquN4YpVf2PXEQkVT
1 parent fdc3c84 commit b48dffd

8 files changed

Lines changed: 327 additions & 52 deletions

File tree

crates/docling-pdf/src/assemble.rs

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -316,6 +316,37 @@ fn dedup_nested_code(kept: &mut Vec<Region>) {
316316
kept.retain(|_| !*keep.next().unwrap());
317317
}
318318

319+
/// Fraction of the page's non-empty text cells that some detected region
320+
/// claims (>0.2 intersection-over-self, docling's assignment rule). 1.0 for a
321+
/// page without text cells.
322+
///
323+
/// The int8-layout guard keys off this: a dense digital page whose detections
324+
/// cover almost none of its text is the signature of quantized confidences
325+
/// flipping under the 0.5 label thresholds on this CPU's kernels — not of a
326+
/// genuinely empty layout — and is worth re-running on the fp32 graph.
327+
pub fn layout_cell_coverage(regions: &[Region], cells: &[TextCell]) -> f32 {
328+
let mut total = 0usize;
329+
let mut covered = 0usize;
330+
for c in cells {
331+
if c.text.trim().is_empty() {
332+
continue;
333+
}
334+
total += 1;
335+
let ca = area(c.l, c.t, c.r, c.b).max(1.0);
336+
if regions
337+
.iter()
338+
.any(|r| inter(r, c.l, c.t, c.r, c.b) / ca > 0.2)
339+
{
340+
covered += 1;
341+
}
342+
}
343+
if total == 0 {
344+
1.0
345+
} else {
346+
covered as f32 / total as f32
347+
}
348+
}
349+
319350
/// Append `text` regions for cells the layout left uncovered ("orphan cells"),
320351
/// the way docling's `LayoutPostprocessor` does (`create_orphan_clusters`): any
321352
/// non-empty cell that no kept region covers (>50% of the cell's area) becomes a
@@ -1926,6 +1957,38 @@ mod tests {
19261957
use crate::pdfium_backend::{LinkAnnot, PdfPage, TextCell};
19271958
use docling_core::Node;
19281959

1960+
/// The int8-layout guard's coverage metric: cells under detections count,
1961+
/// cells outside don't, whitespace cells are ignored, and a cell-less page
1962+
/// reads as fully covered (nothing to rescue).
1963+
#[test]
1964+
fn layout_cell_coverage_counts_claimed_text_cells() {
1965+
let cell = |text: &str, l: f32, t: f32| TextCell {
1966+
text: text.into(),
1967+
l,
1968+
t,
1969+
r: l + 40.0,
1970+
b: t + 10.0,
1971+
};
1972+
let region = Region {
1973+
label: "text",
1974+
score: 0.9,
1975+
l: 0.0,
1976+
t: 0.0,
1977+
r: 100.0,
1978+
b: 50.0,
1979+
};
1980+
let cells = vec![
1981+
cell("inside", 10.0, 10.0),
1982+
cell("also inside", 10.0, 30.0),
1983+
cell("outside", 10.0, 200.0),
1984+
cell(" ", 10.0, 210.0), // whitespace: not counted at all
1985+
];
1986+
let cov = super::layout_cell_coverage(std::slice::from_ref(&region), &cells);
1987+
assert!((cov - 2.0 / 3.0).abs() < 1e-6, "got {cov}");
1988+
assert_eq!(super::layout_cell_coverage(&[], &[]), 1.0);
1989+
assert_eq!(super::layout_cell_coverage(&[], &cells), 0.0);
1990+
}
1991+
19291992
/// #165: a picture no longer claims cells at 0.2 intersection-over-self.
19301993
/// A line straddling the figure border (≤80 % contained) becomes an orphan
19311994
/// region and survives the contained-regulars drop — before the fix its

crates/docling-pdf/src/layout.rs

Lines changed: 69 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,18 @@ pub struct LayoutModel {
8484
/// `layout_heron_int8.onnx`. Batched calls then fall back to per-page runs
8585
/// instead of failing the conversion.
8686
batch_unsupported: bool,
87+
/// The fp32 graph to escalate a suspicious page to, set only when the
88+
/// *auto-selected* int8 graph loaded (an explicit `DOCLING_LAYOUT_ONNX` /
89+
/// `DOCLING_RS_FP32` choice is respected). Int8 confidences sit close
90+
/// enough to the 0.5 label thresholds that a different CPU's quantized
91+
/// kernels (AVX-VNNI vs AVX2, CUDA's fallback mix) can flip a whole page's
92+
/// detections — observed as a bill page whose tables all dissolved into
93+
/// orphan lines on one machine while converting perfectly on another.
94+
fp32_path: Option<String>,
95+
/// Lazily-loaded session over `fp32_path` — most documents never pay for it.
96+
fp32: Option<Session>,
97+
/// Intra-op threads, kept for the lazy fp32 load.
98+
intra: usize,
8799
}
88100

89101
#[cfg(feature = "ml")]
@@ -107,20 +119,62 @@ impl LayoutModel {
107119
if crate::timing::enabled() {
108120
eprintln!("docling-pdf: layout model: {path}");
109121
}
122+
// Escalation target for the quant-robustness guard: only when the
123+
// int8 graph was picked automatically and the fp32 one is also there.
124+
let fp32_path = if std::env::var("DOCLING_LAYOUT_ONNX").is_err() {
125+
let fp32 = crate::resolve_asset("models/layout_heron.onnx");
126+
(path != fp32 && std::path::Path::new(&fp32).exists()).then_some(fp32)
127+
} else {
128+
None
129+
};
130+
let session = Self::open_session(&path, intra)?;
131+
Ok(Self {
132+
session,
133+
batch_unsupported: false,
134+
fp32_path,
135+
fp32: None,
136+
intra,
137+
})
138+
}
139+
140+
fn open_session(path: &str, intra: usize) -> Result<Session, String> {
110141
let builder = Session::builder()
111142
.map_err(|e| format!("layout: builder: {e}"))?
112143
// Let inference use the available cores (ort otherwise defaults low);
113144
// a large PDF runs this model once per page.
114145
.with_intra_threads(intra)
115146
.map_err(|e| format!("layout: intra_threads: {e}"))?;
116-
let session = crate::ep::apply(builder)
147+
crate::ep::apply(builder)
117148
.map_err(|e| format!("layout: {e}"))?
118-
.commit_from_file(&path)
119-
.map_err(|e| format!("layout: load {path}: {e}"))?;
120-
Ok(Self {
121-
session,
122-
batch_unsupported: false,
123-
})
149+
.commit_from_file(path)
150+
.map_err(|e| format!("layout: load {path}: {e}"))
151+
}
152+
153+
/// Re-run one page through the fp32 graph — the escape hatch for a page
154+
/// whose int8 detections look implausible (see `fp32_path`). `Ok(None)`
155+
/// when there is nothing to escalate to: fp32 already loaded, an explicit
156+
/// model override, or no fp32 file on disk.
157+
pub fn predict_fp32_fallback(
158+
&mut self,
159+
img: &RgbImage,
160+
page_w: f32,
161+
page_h: f32,
162+
) -> Result<Option<Vec<Region>>, String> {
163+
let Some(path) = self.fp32_path.clone() else {
164+
return Ok(None);
165+
};
166+
if self.fp32.is_none() {
167+
if crate::timing::enabled() {
168+
eprintln!("docling-pdf: loading fp32 layout fallback: {path}");
169+
}
170+
self.fp32 = Some(Self::open_session(&path, self.intra)?);
171+
}
172+
let session = self.fp32.as_mut().expect("just loaded");
173+
Ok(Some(
174+
Self::run_on(session, &[(img, page_w, page_h)])?
175+
.pop()
176+
.expect("one result per input page"),
177+
))
124178
}
125179

126180
/// Detect layout regions on a page image. `page_w`/`page_h` are the page size
@@ -183,6 +237,13 @@ impl LayoutModel {
183237
}
184238

185239
fn run_batch(&mut self, pages: &[(&RgbImage, f32, f32)]) -> Result<Vec<Vec<Region>>, String> {
240+
Self::run_on(&mut self.session, pages)
241+
}
242+
243+
fn run_on(
244+
session: &mut Session,
245+
pages: &[(&RgbImage, f32, f32)],
246+
) -> Result<Vec<Vec<Region>>, String> {
186247
if pages.is_empty() {
187248
return Ok(Vec::new());
188249
}
@@ -202,8 +263,7 @@ impl LayoutModel {
202263
}
203264
let input = Tensor::from_array(([batch, 3, SIDE as usize, SIDE as usize], data))
204265
.map_err(|e| format!("layout: input tensor: {e}"))?;
205-
let outputs = self
206-
.session
266+
let outputs = session
207267
.run(ort::inputs!["pixel_values" => input])
208268
.map_err(|e| format!("layout: inference: {e}"))?;
209269
let (lshape, logits) = outputs["logits"]

crates/docling-pdf/src/lib.rs

Lines changed: 102 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -238,6 +238,64 @@ pub(crate) fn resolve_asset(rel: &str) -> String {
238238
/// models are conformance-validated — see docs/PDF_CONFORMANCE.md — and load/run
239239
/// markedly faster on CPU), unless `DOCLING_RS_FP32` opts back into full
240240
/// precision; else the fp32 default.
241+
/// One resolved runtime asset — which file a stage would load right now,
242+
/// given the CWD, the env overrides and the int8/fp32 preference.
243+
#[cfg(feature = "ml")]
244+
#[derive(Debug, Clone)]
245+
pub struct ModelEntry {
246+
/// Pipeline stage, e.g. `layout`, `tableformer.decoder`, `ocr.rec`.
247+
pub stage: &'static str,
248+
/// The resolved path (absolute or CWD-relative, as it will be opened).
249+
pub path: String,
250+
/// Whether the file exists right now.
251+
pub found: bool,
252+
/// File size in bytes (0 when missing) — enough to tell an int8 quant
253+
/// from an fp32 graph, or a stale model from a re-published one, at a
254+
/// glance without hashing gigabytes per request.
255+
pub bytes: u64,
256+
}
257+
258+
/// Resolve the whole runtime model set **without loading anything** — the
259+
/// exact selection each stage performs at load time (layout honors the
260+
/// int8/fp32 preference, TableFormer its decoder ranking, OCR the language
261+
/// pair), plus the pdfium library. docling-serve exposes this at
262+
/// `/v1/config` and logs it at startup, so "the server picked up different
263+
/// models" is one `curl` away instead of a mystery of dissolved tables.
264+
/// Resolution is CWD-relative with an exe-dir fallback, so the answer can
265+
/// legitimately differ between two working directories.
266+
#[cfg(feature = "ml")]
267+
pub fn model_inventory() -> Vec<ModelEntry> {
268+
fn entry(stage: &'static str, path: String) -> ModelEntry {
269+
let meta = std::fs::metadata(&path).ok();
270+
ModelEntry {
271+
stage,
272+
found: meta.is_some(),
273+
bytes: meta.map(|m| m.len()).unwrap_or(0),
274+
path,
275+
}
276+
}
277+
let (enc, dec, bbx) = tableformer::resolved_paths();
278+
let (rec, dict) = ocr::resolve_rec_pair(ocr::OcrLang::from_env());
279+
let pdfium =
280+
std::env::var("PDFIUM_DYNAMIC_LIB_PATH").unwrap_or_else(|_| resolve_asset(".pdfium/lib"));
281+
vec![
282+
entry(
283+
"layout",
284+
model_path(
285+
"DOCLING_LAYOUT_ONNX",
286+
"models/layout_heron.onnx",
287+
"models/layout_heron_int8.onnx",
288+
),
289+
),
290+
entry("tableformer.encoder", enc),
291+
entry("tableformer.decoder", dec),
292+
entry("tableformer.bbox", bbx),
293+
entry("ocr.rec", rec),
294+
entry("ocr.dict", dict),
295+
entry("pdfium", pdfium),
296+
]
297+
}
298+
241299
pub(crate) fn model_path(env: &str, fp32_default: &str, int8_default: &str) -> String {
242300
if let Ok(p) = std::env::var(env) {
243301
return p;
@@ -495,11 +553,54 @@ impl Worker {
495553
page.code_cells.clear();
496554
page.word_cells.clear();
497555
}
556+
// Quant-robustness guard: the default int8 layout graph keeps its
557+
// confidences near the 0.5 label thresholds, and a different CPU's
558+
// quantized kernels can flip a whole page's detections under them —
559+
// tables and paragraphs then dissolve into orphan one-liners while the
560+
// same build converts the page perfectly elsewhere. When a dense
561+
// digital page ends up with detections covering almost none of its
562+
// text cells, re-run that one page on the fp32 graph (lazy-loaded,
563+
// auto-int8 selection only) and keep whichever detections cover more.
564+
let mut regions = regions;
565+
if !page.cells.is_empty() {
566+
let thresholded = |rs: &[layout::Region]| -> Vec<layout::Region> {
567+
rs.iter()
568+
.filter(|r| r.score >= layout::label_threshold(r.label))
569+
.cloned()
570+
.collect()
571+
};
572+
let text_cells = page
573+
.cells
574+
.iter()
575+
.filter(|c| !c.text.trim().is_empty())
576+
.count();
577+
let cov = assemble::layout_cell_coverage(&thresholded(&regions), &page.cells);
578+
if text_cells >= 15 && cov < 0.5 {
579+
let retry = self
580+
.layout
581+
.as_mut()
582+
.expect("layout model loaded unless no_ocr")
583+
.predict_fp32_fallback(&page.image, page.width, page.height)
584+
.map_err(|e| PdfError::Layout(format!("page {}: {e}", n + 1)))?;
585+
if let Some(retry) = retry {
586+
let cov2 = assemble::layout_cell_coverage(&thresholded(&retry), &page.cells);
587+
if cov2 > cov {
588+
eprintln!(
589+
"docling-pdf: page {}: int8 layout covered {:.0}% of the text \
590+
cells; the fp32 retry covers {:.0}% — using it",
591+
n + 1,
592+
cov * 100.0,
593+
cov2 * 100.0
594+
);
595+
regions = retry;
596+
}
597+
}
598+
}
599+
}
498600
// docling's LayoutPostprocessor drops each detection below its label's
499601
// confidence threshold (stricter than the 0.3 base the predictor keeps),
500602
// before any overlap resolution. This removes the low-confidence tables /
501603
// pictures / list-items that otherwise double-emit or mis-classify.
502-
let mut regions = regions;
503604
regions.retain(|r| r.score >= layout::label_threshold(r.label));
504605
// Resolve overlapping detections once, before OCR.
505606
let mut regions = assemble::resolve(regions);

crates/docling-pdf/src/ocr.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ impl OcrLang {
7171
/// pair with a warning rather than failing — the usual missing-optional-asset
7272
/// convention. Explicit `DOCLING_OCR_REC_ONNX` / `DOCLING_OCR_DICT` paths win
7373
/// over all of this; they are a pair, so set both together.
74-
fn resolve_rec_pair(lang: OcrLang) -> (String, String) {
74+
pub(crate) fn resolve_rec_pair(lang: OcrLang) -> (String, String) {
7575
const CH: (&str, &str) = ("models/ocr_rec.onnx", "models/ppocr_keys_v1.txt");
7676
const EN: (&str, &str) = ("models/ocr_rec_en.onnx", "models/en_dict.txt");
7777
let want_ch = lang == OcrLang::Ch;

0 commit comments

Comments
 (0)