Merge pull request #170 from artiz/claude/serve-ui-redesign

serve: docling.rs page shell for the playground, and an OpenAPI descr…

Author: artiz

README.md

@@ -117,7 +117,9 @@ cold call on the test fixtures); a semaphore bounds concurrent conversions.
 Markdown responses stream (chunked transfer); `/health` + `/ready` suit
 container probes, and SIGTERM drains in-flight requests before exit.
 `GET /` serves API docs plus an interactive test form — upload or URL in,
-streamed result out, with extracted pictures rendered below the text:
+streamed result out, with extracted pictures rendered below the text — and
+`GET /openapi.yaml` describes the whole API (OpenAPI 3.1), so Swagger UI,
+Redoc or a client generator can be pointed straight at a running server:
 
 <p align="center">
   <img src="docs/assets/serve-form.png" alt="docling-serve test form: a converted image with the Markdown result and a gallery of extracted pictures" width="720">

crates/docling-pdf/src/assemble.rs

@@ -316,6 +316,37 @@ fn dedup_nested_code(kept: &mut Vec<Region>) {
     kept.retain(|_| !*keep.next().unwrap());
 }
 
+/// Fraction of the page's non-empty text cells that some detected region
+/// claims (>0.2 intersection-over-self, docling's assignment rule). 1.0 for a
+/// page without text cells.
+///
+/// The int8-layout guard keys off this: a dense digital page whose detections
+/// cover almost none of its text is the signature of quantized confidences
+/// flipping under the 0.5 label thresholds on this CPU's kernels — not of a
+/// genuinely empty layout — and is worth re-running on the fp32 graph.
+pub fn layout_cell_coverage(regions: &[Region], cells: &[TextCell]) -> f32 {
+    let mut total = 0usize;
+    let mut covered = 0usize;
+    for c in cells {
+        if c.text.trim().is_empty() {
+            continue;
+        }
+        total += 1;
+        let ca = area(c.l, c.t, c.r, c.b).max(1.0);
+        if regions
+            .iter()
+            .any(|r| inter(r, c.l, c.t, c.r, c.b) / ca > 0.2)
+        {
+            covered += 1;
+        }
+    }
+    if total == 0 {
+        1.0
+    } else {
+        covered as f32 / total as f32
+    }
+}
+
 /// Append `text` regions for cells the layout left uncovered ("orphan cells"),
 /// the way docling's `LayoutPostprocessor` does (`create_orphan_clusters`): any
 /// non-empty cell that no kept region covers (>50% of the cell's area) becomes a
@@ -1926,6 +1957,38 @@ mod tests {
     use crate::pdfium_backend::{LinkAnnot, PdfPage, TextCell};
     use docling_core::Node;
 
+    /// The int8-layout guard's coverage metric: cells under detections count,
+    /// cells outside don't, whitespace cells are ignored, and a cell-less page
+    /// reads as fully covered (nothing to rescue).
+    #[test]
+    fn layout_cell_coverage_counts_claimed_text_cells() {
+        let cell = |text: &str, l: f32, t: f32| TextCell {
+            text: text.into(),
+            l,
+            t,
+            r: l + 40.0,
+            b: t + 10.0,
+        };
+        let region = Region {
+            label: "text",
+            score: 0.9,
+            l: 0.0,
+            t: 0.0,
+            r: 100.0,
+            b: 50.0,
+        };
+        let cells = vec![
+            cell("inside", 10.0, 10.0),
+            cell("also inside", 10.0, 30.0),
+            cell("outside", 10.0, 200.0),
+            cell("   ", 10.0, 210.0), // whitespace: not counted at all
+        ];
+        let cov = super::layout_cell_coverage(std::slice::from_ref(&region), &cells);
+        assert!((cov - 2.0 / 3.0).abs() < 1e-6, "got {cov}");
+        assert_eq!(super::layout_cell_coverage(&[], &[]), 1.0);
+        assert_eq!(super::layout_cell_coverage(&[], &cells), 0.0);
+    }
+
     /// #165: a picture no longer claims cells at 0.2 intersection-over-self.
     /// A line straddling the figure border (≤80 % contained) becomes an orphan
     /// region and survives the contained-regulars drop — before the fix its

crates/docling-pdf/src/layout.rs

@@ -84,6 +84,18 @@ pub struct LayoutModel {
     /// `layout_heron_int8.onnx`. Batched calls then fall back to per-page runs
     /// instead of failing the conversion.
     batch_unsupported: bool,
+    /// The fp32 graph to escalate a suspicious page to, set only when the
+    /// *auto-selected* int8 graph loaded (an explicit `DOCLING_LAYOUT_ONNX` /
+    /// `DOCLING_RS_FP32` choice is respected). Int8 confidences sit close
+    /// enough to the 0.5 label thresholds that a different CPU's quantized
+    /// kernels (AVX-VNNI vs AVX2, CUDA's fallback mix) can flip a whole page's
+    /// detections — observed as a bill page whose tables all dissolved into
+    /// orphan lines on one machine while converting perfectly on another.
+    fp32_path: Option<String>,
+    /// Lazily-loaded session over `fp32_path` — most documents never pay for it.
+    fp32: Option<Session>,
+    /// Intra-op threads, kept for the lazy fp32 load.
+    intra: usize,
 }
 
 #[cfg(feature = "ml")]
@@ -107,20 +119,62 @@ impl LayoutModel {
         if crate::timing::enabled() {
             eprintln!("docling-pdf: layout model: {path}");
         }
+        // Escalation target for the quant-robustness guard: only when the
+        // int8 graph was picked automatically and the fp32 one is also there.
+        let fp32_path = if std::env::var("DOCLING_LAYOUT_ONNX").is_err() {
+            let fp32 = crate::resolve_asset("models/layout_heron.onnx");
+            (path != fp32 && std::path::Path::new(&fp32).exists()).then_some(fp32)
+        } else {
+            None
+        };
+        let session = Self::open_session(&path, intra)?;
+        Ok(Self {
+            session,
+            batch_unsupported: false,
+            fp32_path,
+            fp32: None,
+            intra,
+        })
+    }
+
+    fn open_session(path: &str, intra: usize) -> Result<Session, String> {
         let builder = Session::builder()
             .map_err(|e| format!("layout: builder: {e}"))?
             // Let inference use the available cores (ort otherwise defaults low);
             // a large PDF runs this model once per page.
             .with_intra_threads(intra)
             .map_err(|e| format!("layout: intra_threads: {e}"))?;
-        let session = crate::ep::apply(builder)
+        crate::ep::apply(builder)
             .map_err(|e| format!("layout: {e}"))?
-            .commit_from_file(&path)
-            .map_err(|e| format!("layout: load {path}: {e}"))?;
-        Ok(Self {
-            session,
-            batch_unsupported: false,
-        })
+            .commit_from_file(path)
+            .map_err(|e| format!("layout: load {path}: {e}"))
+    }
+
+    /// Re-run one page through the fp32 graph — the escape hatch for a page
+    /// whose int8 detections look implausible (see `fp32_path`). `Ok(None)`
+    /// when there is nothing to escalate to: fp32 already loaded, an explicit
+    /// model override, or no fp32 file on disk.
+    pub fn predict_fp32_fallback(
+        &mut self,
+        img: &RgbImage,
+        page_w: f32,
+        page_h: f32,
+    ) -> Result<Option<Vec<Region>>, String> {
+        let Some(path) = self.fp32_path.clone() else {
+            return Ok(None);
+        };
+        if self.fp32.is_none() {
+            if crate::timing::enabled() {
+                eprintln!("docling-pdf: loading fp32 layout fallback: {path}");
+            }
+            self.fp32 = Some(Self::open_session(&path, self.intra)?);
+        }
+        let session = self.fp32.as_mut().expect("just loaded");
+        Ok(Some(
+            Self::run_on(session, &[(img, page_w, page_h)])?
+                .pop()
+                .expect("one result per input page"),
+        ))
     }
 
     /// Detect layout regions on a page image. `page_w`/`page_h` are the page size
@@ -183,6 +237,13 @@ impl LayoutModel {
     }
 
     fn run_batch(&mut self, pages: &[(&RgbImage, f32, f32)]) -> Result<Vec<Vec<Region>>, String> {
+        Self::run_on(&mut self.session, pages)
+    }
+
+    fn run_on(
+        session: &mut Session,
+        pages: &[(&RgbImage, f32, f32)],
+    ) -> Result<Vec<Vec<Region>>, String> {
         if pages.is_empty() {
             return Ok(Vec::new());
         }
@@ -202,8 +263,7 @@ impl LayoutModel {
         }
         let input = Tensor::from_array(([batch, 3, SIDE as usize, SIDE as usize], data))
             .map_err(|e| format!("layout: input tensor: {e}"))?;
-        let outputs = self
-            .session
+        let outputs = session
             .run(ort::inputs!["pixel_values" => input])
             .map_err(|e| format!("layout: inference: {e}"))?;
         let (lshape, logits) = outputs["logits"]

crates/docling-pdf/src/lib.rs

@@ -232,12 +232,70 @@ pub(crate) fn resolve_asset(rel: &str) -> String {
     rel.to_string()
 }
 
+/// One resolved runtime asset — which file a stage would load right now,
+/// given the CWD, the env overrides and the int8/fp32 preference.
 #[cfg(feature = "ml")]
+#[derive(Debug, Clone)]
+pub struct ModelEntry {
+    /// Pipeline stage, e.g. `layout`, `tableformer.decoder`, `ocr.rec`.
+    pub stage: &'static str,
+    /// The resolved path (absolute or CWD-relative, as it will be opened).
+    pub path: String,
+    /// Whether the file exists right now.
+    pub found: bool,
+    /// File size in bytes (0 when missing) — enough to tell an int8 quant
+    /// from an fp32 graph, or a stale model from a re-published one, at a
+    /// glance without hashing gigabytes per request.
+    pub bytes: u64,
+}
+
+/// Resolve the whole runtime model set **without loading anything** — the
+/// exact selection each stage performs at load time (layout honors the
+/// int8/fp32 preference, TableFormer its decoder ranking, OCR the language
+/// pair), plus the pdfium library. docling-serve exposes this at
+/// `/v1/config` and logs it at startup, so "the server picked up different
+/// models" is one `curl` away instead of a mystery of dissolved tables.
+/// Resolution is CWD-relative with an exe-dir fallback, so the answer can
+/// legitimately differ between two working directories.
+#[cfg(feature = "ml")]
+pub fn model_inventory() -> Vec<ModelEntry> {
+    fn entry(stage: &'static str, path: String) -> ModelEntry {
+        let meta = std::fs::metadata(&path).ok();
+        ModelEntry {
+            stage,
+            found: meta.is_some(),
+            bytes: meta.map(|m| m.len()).unwrap_or(0),
+            path,
+        }
+    }
+    let (enc, dec, bbx) = tableformer::resolved_paths();
+    let (rec, dict) = ocr::resolve_rec_pair(ocr::OcrLang::from_env());
+    let pdfium =
+        std::env::var("PDFIUM_DYNAMIC_LIB_PATH").unwrap_or_else(|_| resolve_asset(".pdfium/lib"));
+    vec![
+        entry(
+            "layout",
+            model_path(
+                "DOCLING_LAYOUT_ONNX",
+                "models/layout_heron.onnx",
+                "models/layout_heron_int8.onnx",
+            ),
+        ),
+        entry("tableformer.encoder", enc),
+        entry("tableformer.decoder", dec),
+        entry("tableformer.bbox", bbx),
+        entry("ocr.rec", rec),
+        entry("ocr.dict", dict),
+        entry("pdfium", pdfium),
+    ]
+}
+
 /// Resolve a model path: an explicit env override always wins; otherwise the
 /// INT8 variant of the default path when it exists on disk (the quantized
 /// models are conformance-validated — see docs/PDF_CONFORMANCE.md — and load/run
 /// markedly faster on CPU), unless `DOCLING_RS_FP32` opts back into full
 /// precision; else the fp32 default.
+#[cfg(feature = "ml")]
 pub(crate) fn model_path(env: &str, fp32_default: &str, int8_default: &str) -> String {
     if let Ok(p) = std::env::var(env) {
         return p;
@@ -495,11 +553,54 @@ impl Worker {
             page.code_cells.clear();
             page.word_cells.clear();
         }
+        // Quant-robustness guard: the default int8 layout graph keeps its
+        // confidences near the 0.5 label thresholds, and a different CPU's
+        // quantized kernels can flip a whole page's detections under them —
+        // tables and paragraphs then dissolve into orphan one-liners while the
+        // same build converts the page perfectly elsewhere. When a dense
+        // digital page ends up with detections covering almost none of its
+        // text cells, re-run that one page on the fp32 graph (lazy-loaded,
+        // auto-int8 selection only) and keep whichever detections cover more.
+        let mut regions = regions;
+        if !page.cells.is_empty() {
+            let thresholded = |rs: &[layout::Region]| -> Vec<layout::Region> {
+                rs.iter()
+                    .filter(|r| r.score >= layout::label_threshold(r.label))
+                    .cloned()
+                    .collect()
+            };
+            let text_cells = page
+                .cells
+                .iter()
+                .filter(|c| !c.text.trim().is_empty())
+                .count();
+            let cov = assemble::layout_cell_coverage(&thresholded(&regions), &page.cells);
+            if text_cells >= 15 && cov < 0.5 {
+                let retry = self
+                    .layout
+                    .as_mut()
+                    .expect("layout model loaded unless no_ocr")
+                    .predict_fp32_fallback(&page.image, page.width, page.height)
+                    .map_err(|e| PdfError::Layout(format!("page {}: {e}", n + 1)))?;
+                if let Some(retry) = retry {
+                    let cov2 = assemble::layout_cell_coverage(&thresholded(&retry), &page.cells);
+                    if cov2 > cov {
+                        eprintln!(
+                            "docling-pdf: page {}: int8 layout covered {:.0}% of the text \
+                             cells; the fp32 retry covers {:.0}% — using it",
+                            n + 1,
+                            cov * 100.0,
+                            cov2 * 100.0
+                        );
+                        regions = retry;
+                    }
+                }
+            }
+        }
         // docling's LayoutPostprocessor drops each detection below its label's
         // confidence threshold (stricter than the 0.3 base the predictor keeps),
         // before any overlap resolution. This removes the low-confidence tables /
         // pictures / list-items that otherwise double-emit or mis-classify.
-        let mut regions = regions;
         regions.retain(|r| r.score >= layout::label_threshold(r.label));
         // Resolve overlapping detections once, before OCR.
         let mut regions = assemble::resolve(regions);

crates/docling-pdf/src/ocr.rs

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