Skip to main content

docling_core/
chunker.rs

1//! Document chunking for RAG pipelines — the Rust port of docling-core's
2//! `docling_core.transforms.chunker`.
3//!
4//! Two chunkers, matching docling's semantics output-for-output:
5//!
6//! * [`HierarchicalChunker`] walks the document tree and yields one chunk per
7//!   top-level item (paragraph, whole list, table, picture caption, …), each
8//!   carrying the heading path it sits under. Tables are serialized in
9//!   docling's *triplet* form (`row, column = value`), pictures contribute
10//!   their captions.
11//! * [`HybridChunker`] refines the hierarchical chunks with a tokenizer:
12//!   oversized chunks are split (at item boundaries first, then within the
13//!   text by docling's `semchunk` algorithm), and undersized neighbours that
14//!   share the same headings are merged back together.
15//!
16//! [`contextualize`] renders a chunk to the string an embedding model should
17//! see: the heading path plus the chunk text.
18//!
19//! Anything tokenizer-related is abstracted behind [`ChunkTokenizer`]; a
20//! HuggingFace `tokenizers` implementation ships behind the `chunking` cargo
21//! feature as [`HuggingFaceTokenizer`].
22
23use std::collections::BTreeMap;
24
25use crate::document::{DoclingDocument, Node, Table};
26
27/// What kind of document item a [`ChunkItem`] points at — only what the
28/// hybrid splitting logic needs to know.
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub enum ChunkItemKind {
31    /// A text-ish item (paragraph, list item, code, formula, caption, …).
32    Text,
33    /// A table item (single-item oversized chunks re-split per line, repeating
34    /// docling's header handling).
35    Table,
36    /// A picture item.
37    Picture,
38}
39
40/// One document item contributing to a chunk — the analogue of an entry in
41/// docling's `DocMeta.doc_items`.
42#[derive(Debug, Clone, PartialEq)]
43pub struct ChunkItem {
44    /// The item's ref in [`DoclingDocument::export_to_json`] output
45    /// (`#/texts/12`, `#/tables/0`, `#/pictures/1`, …).
46    pub self_ref: String,
47    pub kind: ChunkItemKind,
48    /// The item serialized *standalone* (docling re-serializes individual
49    /// items when splitting an oversized multi-item chunk — e.g. a nested
50    /// list item flattens to `- text` with no indentation).
51    pub text: String,
52}
53
54/// One chunk — the analogue of docling's `DocChunk` (text + `DocMeta`).
55#[derive(Debug, Clone, PartialEq)]
56pub struct DocChunk {
57    /// The chunk body (markdown-flavoured, unescaped — same text docling puts
58    /// in `DocChunk.text`).
59    pub text: String,
60    /// The heading path above this chunk, outermost first (`DocMeta.headings`;
61    /// `None` when the chunk sits above any heading).
62    pub headings: Option<Vec<String>>,
63    /// The document items the chunk was built from (`DocMeta.doc_items`).
64    pub doc_items: Vec<ChunkItem>,
65}
66
67/// Render a chunk for embedding: the heading path, then the text, joined with
68/// newlines — docling's `BaseChunker.contextualize()`.
69pub fn contextualize(chunk: &DocChunk) -> String {
70    let mut parts: Vec<&str> = Vec::new();
71    if let Some(h) = &chunk.headings {
72        parts.extend(h.iter().map(String::as_str));
73    }
74    parts.push(&chunk.text);
75    parts.join("\n")
76}
77
78// ---------------------------------------------------------------------------
79// Hierarchical chunker
80// ---------------------------------------------------------------------------
81
82/// Structure-driven chunker: one chunk per document item, lists and inline
83/// groups kept whole, heading path tracked as metadata — docling-core's
84/// `HierarchicalChunker` with default parameters.
85#[derive(Debug, Clone, Default)]
86pub struct HierarchicalChunker;
87
88impl HierarchicalChunker {
89    /// Chunk the document.
90    pub fn chunk(&self, doc: &DoclingDocument) -> Vec<DocChunk> {
91        let mut chunks = Vec::new();
92        self.chunk_with(doc, &mut |c| {
93            chunks.push(c);
94            true
95        });
96        chunks
97    }
98
99    /// Stream the chunks: `sink` is called with each chunk as the document
100    /// walk produces it, so a consumer can process (embed, forward) chunks
101    /// without materializing the whole `Vec` first. A `false` return from
102    /// `sink` cancels the walk. [`Self::chunk`] is this with a collecting
103    /// sink — the chunks and their order are identical.
104    pub fn chunk_with(&self, doc: &DoclingDocument, sink: &mut dyn FnMut(DocChunk) -> bool) {
105        let mut w = Walker {
106            alloc: Alloc::default(),
107            headings: BTreeMap::new(),
108            stopped: false,
109            sink,
110        };
111        w.walk(&doc.nodes);
112    }
113}
114
115/// Ref allocator mirroring the numbering `json.rs` gives every item, so
116/// `ChunkItem::self_ref` matches the document's JSON export.
117#[derive(Debug, Default)]
118struct Alloc {
119    texts: usize,
120    groups: usize,
121    tables: usize,
122    pictures: usize,
123    field_regions: usize,
124    field_items: usize,
125}
126
127impl Alloc {
128    fn text(&mut self) -> String {
129        let r = format!("#/texts/{}", self.texts);
130        self.texts += 1;
131        r
132    }
133    fn group(&mut self) -> String {
134        let r = format!("#/groups/{}", self.groups);
135        self.groups += 1;
136        r
137    }
138    fn table(&mut self) -> String {
139        let r = format!("#/tables/{}", self.tables);
140        self.tables += 1;
141        r
142    }
143    fn picture(&mut self) -> String {
144        let r = format!("#/pictures/{}", self.pictures);
145        self.pictures += 1;
146        r
147    }
148    fn field_region(&mut self) -> String {
149        let r = format!("#/field_regions/{}", self.field_regions);
150        self.field_regions += 1;
151        r
152    }
153    fn field_item(&mut self) -> String {
154        let r = format!("#/field_items/{}", self.field_items);
155        self.field_items += 1;
156        r
157    }
158}
159
160struct Walker<'s> {
161    alloc: Alloc,
162    /// Active heading per docling level (title = 0, `section_header` = its
163    /// `level`), pruned like docling's `heading_by_level`.
164    headings: BTreeMap<u8, String>,
165    /// Set once the sink refuses a chunk; the walk unwinds without emitting.
166    stopped: bool,
167    sink: &'s mut dyn FnMut(DocChunk) -> bool,
168}
169
170impl Walker<'_> {
171    fn emit(&mut self, text: String, doc_items: Vec<ChunkItem>) {
172        if self.stopped || text.is_empty() {
173            return;
174        }
175        let headings: Vec<String> = self.headings.values().cloned().collect();
176        self.stopped = !(self.sink)(DocChunk {
177            text,
178            headings: (!headings.is_empty()).then_some(headings),
179            doc_items,
180        });
181    }
182
183    /// Emit a text chunk whose doc items follow docling's inline granularity:
184    /// mixed inline content (a paragraph docling represents as an inline group)
185    /// contributes one item per span, plain text one item.
186    fn emit_inline(&mut self, md_text: &str, self_ref: String) {
187        self.emit_inline_with_runs(md_text, self_ref, &[]);
188    }
189
190    fn emit_inline_with_runs(
191        &mut self,
192        md_text: &str,
193        self_ref: String,
194        runs: &[crate::InlineRun],
195    ) {
196        let body = unescape_text(md_text);
197        if body.is_empty() {
198            return;
199        }
200        let segments: Vec<String> = inline_segments_tagged(md_text)
201            .into_iter()
202            .flat_map(|(text, is_plain)| {
203                if is_plain {
204                    if let Some(split) = split_plain_by_runs(&text, runs) {
205                        return split;
206                    }
207                }
208                vec![text]
209            })
210            .collect();
211        let items: Vec<ChunkItem> = if segments.len() <= 1 {
212            vec![ChunkItem {
213                self_ref,
214                kind: ChunkItemKind::Text,
215                text: body.clone(),
216            }]
217        } else {
218            segments
219                .into_iter()
220                .map(|text| ChunkItem {
221                    self_ref: self_ref.clone(),
222                    kind: ChunkItemKind::Text,
223                    text,
224                })
225                .collect()
226        };
227        self.emit(body, items);
228    }
229
230    fn set_heading(&mut self, doc_level: u8, text: String) {
231        self.headings.retain(|k, _| *k < doc_level);
232        self.headings.insert(doc_level, text);
233    }
234
235    fn walk(&mut self, nodes: &[Node]) {
236        let mut i = 0;
237        while i < nodes.len() {
238            if self.stopped {
239                return;
240            }
241            if matches!(nodes[i], Node::ListItem { .. }) {
242                let start = i;
243                i += 1;
244                loop {
245                    match nodes.get(i) {
246                        Some(Node::ListItem { .. }) => i += 1,
247                        // An empty paragraph between two list items is absorbed
248                        // into the run (mirrors json.rs / markdown.rs).
249                        Some(Node::Paragraph { text })
250                            if text.is_empty()
251                                && matches!(nodes.get(i + 1), Some(Node::ListItem { .. })) =>
252                        {
253                            i += 1
254                        }
255                        _ => break,
256                    }
257                }
258                self.sibling_lists(&nodes[start..i]);
259            } else {
260                self.one(&nodes[i]);
261                i += 1;
262            }
263        }
264    }
265
266    /// Split a run of list items into sibling lists exactly like
267    /// `json.rs::add_sibling_lists`, chunking each list separately (docling
268    /// yields one chunk per `ListGroup`).
269    fn sibling_lists(&mut self, run: &[Node]) {
270        let base = level_of(&run[0]);
271        let mut seg = 0;
272        for k in 0..run.len() {
273            let Node::ListItem {
274                first_in_list,
275                level,
276                ..
277            } = &run[k]
278            else {
279                continue;
280            };
281            if *level != base {
282                continue; // nested item — handled inside `list`
283            }
284            // The backend's list boundary, as in `json.rs::add_sibling_lists`.
285            if k > seg && *first_in_list {
286                self.list(&run[seg..k]);
287                seg = k;
288            }
289        }
290        self.list(&run[seg..]);
291    }
292
293    /// One `ListGroup`: allocate refs in `json.rs::add_list` order (group,
294    /// then per top item its text ref followed by any nested groups) and emit
295    /// a single chunk whose text is the indented markdown list.
296    fn list(&mut self, items: &[Node]) {
297        self.alloc.group();
298        let mut chunk_items = Vec::new();
299        self.list_refs(items, &mut chunk_items);
300        let text = render_list(items);
301        self.emit(text, chunk_items);
302    }
303
304    /// Allocate refs for one list's items (and nested sibling lists), mirroring
305    /// `json.rs::add_list` / `add_sibling_lists` recursion, collecting the
306    /// non-furniture items in allocation (= document) order.
307    fn list_refs(&mut self, items: &[Node], out: &mut Vec<ChunkItem>) {
308        let base = level_of(&items[0]);
309        let mut i = 0;
310        while i < items.len() {
311            let Node::ListItem {
312                ordered,
313                number,
314                text,
315                level,
316                layer,
317                ..
318            } = &items[i]
319            else {
320                i += 1;
321                continue;
322            };
323            if *level > base {
324                i += 1;
325                continue;
326            }
327            let item_ref = self.alloc.text();
328            let mut j = i + 1;
329            while j < items.len() && level_of(&items[j]) > base {
330                j += 1;
331            }
332            let has_nested = j > i + 1;
333            if layer.is_none() {
334                let marker = if *ordered {
335                    format!("{number}.")
336                } else {
337                    "-".to_string()
338                };
339                // An item that carries both inline spans and a nested list is an
340                // empty list item wrapping an inline group in docling's model:
341                // its marker and each inline span are separate doc items.
342                // docling represents a list item whose content is not pure
343                // inline text (it carries a nested list or child pictures) as
344                // an empty item wrapping an inline group: its marker and each
345                // inline span become separate doc items.
346                // docling represents a list item whose content is not pure
347                // inline text (it carries a nested list or its own images) as
348                // an empty item wrapping an inline group: its marker and each
349                // inline span become separate doc items. The chunker's
350                // image placeholder is empty, so the markers are stripped from
351                // the text before segmentation.
352                let has_pics = text.contains("<!-- image -->");
353                let text = strip_image_markers(text);
354                let text = text.as_str();
355                let segments = inline_segments(text);
356                if (has_nested || has_pics) && segments.len() > 1 && text.contains("](") {
357                    out.push(ChunkItem {
358                        self_ref: item_ref.clone(),
359                        kind: ChunkItemKind::Text,
360                        text: format!("{marker} "),
361                    });
362                    for seg in segments {
363                        out.push(ChunkItem {
364                            self_ref: item_ref.clone(),
365                            kind: ChunkItemKind::Text,
366                            text: seg,
367                        });
368                    }
369                } else {
370                    out.push(ChunkItem {
371                        self_ref: item_ref.clone(),
372                        kind: ChunkItemKind::Text,
373                        text: format!("{marker} {}", unescape_text(text)),
374                    });
375                }
376            }
377            // nested items group under this one; each nested sibling list is a
378            // fresh group ref
379            if j > i + 1 {
380                self.nested_sibling_lists(&items[i + 1..j], out);
381            }
382            i = j;
383        }
384    }
385
386    fn nested_sibling_lists(&mut self, run: &[Node], out: &mut Vec<ChunkItem>) {
387        let base = level_of(&run[0]);
388        let mut seg = 0;
389        for k in 0..run.len() {
390            let Node::ListItem {
391                first_in_list,
392                level,
393                ..
394            } = &run[k]
395            else {
396                continue;
397            };
398            if *level != base {
399                continue;
400            }
401            if k > seg && *first_in_list {
402                self.alloc.group();
403                self.list_refs(&run[seg..k], out);
404                seg = k;
405            }
406        }
407        self.alloc.group();
408        self.list_refs(&run[seg..], out);
409    }
410
411    fn one(&mut self, node: &Node) {
412        match node {
413            Node::Heading { level, text } => {
414                let doc_level = if *level == 1 {
415                    0
416                } else {
417                    level.saturating_sub(1)
418                };
419                let self_ref = self.alloc.text();
420                // docling stores heading text unformatted: a heading that is one
421                // uniformly formatted span keeps its plain text; a *partially*
422                // formatted heading becomes an empty heading whose content is an
423                // inline group — which the chunker then yields as a chunk of its
424                // own (under the freshly-set empty heading).
425                let runs = crate::inline_runs_from_markdown(text);
426                if runs.len() <= 1 {
427                    let plain = runs
428                        .first()
429                        .map(|r| r.text.clone())
430                        .unwrap_or_else(|| text.clone());
431                    self.set_heading(doc_level, unescape_text(&plain));
432                } else {
433                    self.set_heading(doc_level, String::new());
434                    let body = unescape_text(text);
435                    self.emit(
436                        body.clone(),
437                        vec![ChunkItem {
438                            self_ref,
439                            kind: ChunkItemKind::Text,
440                            text: body,
441                        }],
442                    );
443                }
444            }
445            Node::Paragraph { text } => {
446                let t = text.trim();
447                let self_ref = self.alloc.text();
448                // A whole-paragraph display equation is a formula item; docling's
449                // chunk serializer re-wraps the raw latex in `$$…$$`.
450                if let Some(inner) = t
451                    .strip_prefix("$$")
452                    .and_then(|s| s.strip_suffix("$$"))
453                    .filter(|s| !s.is_empty())
454                {
455                    let body = format!("$${inner}$$");
456                    self.emit(
457                        body.clone(),
458                        vec![ChunkItem {
459                            self_ref,
460                            kind: ChunkItemKind::Text,
461                            text: body,
462                        }],
463                    );
464                    return;
465                }
466                self.emit_inline(text, self_ref);
467            }
468            // A standalone caption chunks like any text item (docling's chunker
469            // does not treat `caption` specially outside a table/picture).
470            Node::Caption { text, .. } => {
471                let self_ref = self.alloc.text();
472                self.emit_inline(text, self_ref);
473            }
474            Node::LabeledText { .. } => self.one(&node.labeled_as_paragraph()),
475            Node::CheckboxItem { checked, text } => {
476                let self_ref = self.alloc.text();
477                let mark = if *checked { "- [x] " } else { "- [ ] " };
478                let body = format!("{mark}{}", unescape_text(text));
479                self.emit(
480                    body.clone(),
481                    vec![ChunkItem {
482                        self_ref,
483                        kind: ChunkItemKind::Text,
484                        text: body,
485                    }],
486                );
487            }
488            // An enriched display formula chunks like docling's formula item:
489            // the LaTeX is the chunk text.
490            Node::Formula { latex, .. } => {
491                let self_ref = self.alloc.text();
492                let body = format!("$${}$$", latex);
493                self.emit(
494                    body.clone(),
495                    vec![ChunkItem {
496                        self_ref,
497                        kind: ChunkItemKind::Text,
498                        text: body,
499                    }],
500                );
501            }
502            Node::Code { text, .. } => {
503                let self_ref = self.alloc.text();
504                let body = format!("```\n{}\n```", unescape_text(text));
505                self.emit(
506                    body.clone(),
507                    vec![ChunkItem {
508                        self_ref,
509                        kind: ChunkItemKind::Text,
510                        text: body,
511                    }],
512                );
513            }
514            Node::Table(t) => {
515                let self_ref = self.alloc.table();
516                let body = triplet_table_text(t);
517                self.emit(
518                    body.clone(),
519                    vec![ChunkItem {
520                        self_ref,
521                        kind: ChunkItemKind::Table,
522                        text: body,
523                    }],
524                );
525            }
526            Node::Picture { caption, .. } => {
527                let cap = caption.as_deref().filter(|c| !c.is_empty());
528                let cap_item = cap.map(|c| ChunkItem {
529                    self_ref: self.alloc.text(),
530                    kind: ChunkItemKind::Text,
531                    text: unescape_text(c),
532                });
533                self.alloc.picture();
534                // The picture itself serializes to the (empty) chunking image
535                // placeholder, and its caption is already consumed by the
536                // caption chunk — so only the caption text is emitted.
537                if let Some(cap_item) = cap_item {
538                    let body = cap_item.text.clone();
539                    self.emit(body, vec![cap_item]);
540                }
541            }
542            Node::Chart {
543                kind,
544                table,
545                caption,
546                ..
547            } => {
548                let cap = caption.as_deref().filter(|c| !c.is_empty());
549                let cap_item = cap.map(|c| ChunkItem {
550                    self_ref: self.alloc.text(),
551                    kind: ChunkItemKind::Text,
552                    text: unescape_text(c),
553                });
554                let pic_ref = self.alloc.picture();
555                // caption, humanized classification, then the chart's data grid
556                // as a (padded) markdown table — docling's picture serializer
557                // parts, joined with blank lines.
558                let mut parts: Vec<String> = Vec::new();
559                if let Some(ci) = &cap_item {
560                    parts.push(ci.text.clone());
561                }
562                parts.push(humanize_label(kind));
563                let grid = crate::markdown::render_table(table, false);
564                if !grid.is_empty() {
565                    parts.push(unescape_text(&grid));
566                }
567                let body = parts.join("\n\n");
568                // Re-serialized standalone (the hybrid window join), the picture
569                // carries its caption itself, while the caption *item* renders
570                // empty — docling's markdown serializer emits caption-label text
571                // only through the picture.
572                let pic_item = ChunkItem {
573                    self_ref: pic_ref,
574                    kind: ChunkItemKind::Picture,
575                    text: body.clone(),
576                };
577                let items = match cap_item {
578                    Some(mut ci) => {
579                        ci.text = String::new();
580                        vec![ci, pic_item]
581                    }
582                    None => vec![pic_item],
583                };
584                self.emit(body, items);
585            }
586            // A group on a non-body layer (a hidden spreadsheet sheet) carries
587            // no chunkable content, like every other non-body item.
588            Node::Group { layer: Some(_), .. } => {}
589            Node::Group { children, .. } => {
590                // A generic group is a structural container: docling recurses
591                // into it rather than chunking it whole.
592                self.alloc.group();
593                self.walk(children);
594            }
595            // A key-value graph holds no text items, so no chunk comes of it.
596            Node::KeyValueGraph { .. } => {}
597            Node::FieldRegion { items } => {
598                // Each field part (marker / key / value) is its own text item,
599                // and docling chunks each one individually.
600                self.alloc.field_region();
601                for item in items {
602                    self.alloc.field_item();
603                    for part in [&item.marker, &item.key, &item.value].into_iter().flatten() {
604                        let self_ref = self.alloc.text();
605                        let body = unescape_text(part);
606                        self.emit(
607                            body.clone(),
608                            vec![ChunkItem {
609                                self_ref,
610                                kind: ChunkItemKind::Text,
611                                text: body,
612                            }],
613                        );
614                    }
615                }
616            }
617            Node::InlineGroup { md_text, runs, .. } => {
618                let self_ref = self.alloc.text();
619                self.emit_inline_with_runs(md_text, self_ref, runs);
620            }
621            Node::TextDump(text) => {
622                let self_ref = self.alloc.text();
623                let body = unescape_text(text);
624                self.emit(
625                    body.clone(),
626                    vec![ChunkItem {
627                        self_ref,
628                        kind: ChunkItemKind::Text,
629                        text: body,
630                    }],
631                );
632            }
633            // Layout provenance and comment annotations are transparent.
634            Node::Located { inner, .. }
635            | Node::Prov { inner, .. }
636            | Node::Track { inner, .. }
637            | Node::Commented { inner, .. } => self.one(inner),
638            // A PDF page header/footer is a furniture-layer JSON text item:
639            // never chunked (docling's chunker reads the body layer), but it
640            // takes a `#/texts/N` number, so later refs keep matching the JSON.
641            Node::PageFurniture { .. } | Node::FurnitureText { .. } => {
642                self.alloc.text();
643            }
644            // A picture's children are JSON items docling's chunker never
645            // reaches (it iterates without `traverse_pictures`): allocate
646            // their refs only — one text each, plus the list group `json.rs`
647            // opens for each list item (each starts its own list) — so later
648            // refs keep matching the JSON.
649            Node::PictureChildren(children) => {
650                for child in children {
651                    match child {
652                        // A header/footer child is numbered exactly as a
653                        // top-level one is, whatever the JSON does with it.
654                        Node::PageFurniture { .. } => self.one(child),
655                        Node::ListItem { .. } => {
656                            self.alloc.group();
657                            self.alloc.text();
658                        }
659                        _ => {
660                            self.alloc.text();
661                        }
662                    }
663                }
664            }
665            // Non-body layers and doclang-only nodes don't reach the chunker
666            // (nor the JSON body).
667            Node::CommentSection { .. }
668            | Node::Furniture { .. }
669            | Node::PageBreak
670            | Node::PageInfo { .. }
671            | Node::DoclangOnly(_) => {}
672            // Runs are grouped by `walk`; a stray single item (hand-built
673            // document, `Located` wrapper) still chunks instead of panicking.
674            Node::ListItem { .. } => self.sibling_lists(std::slice::from_ref(node)),
675        }
676    }
677}
678
679fn level_of(node: &Node) -> u8 {
680    match node {
681        Node::ListItem { level, .. } => *level,
682        _ => 0,
683    }
684}
685
686/// Render one sibling list as its markdown chunk text (indented items, same
687/// rules as the full markdown serializer's list rendering).
688fn render_list(items: &[Node]) -> String {
689    let mut lines: Vec<String> = Vec::new();
690    for item in items {
691        let Node::ListItem {
692            ordered,
693            number,
694            text,
695            level,
696            layer,
697            ..
698        } = item
699        else {
700            continue;
701        };
702        if layer.is_some() {
703            continue;
704        }
705        let indent = "    ".repeat(*level as usize);
706        let marker = if *ordered {
707            format!("{number}.")
708        } else {
709            "-".to_string()
710        };
711        lines.push(format!(
712            "{indent}{marker} {}",
713            unescape_text(&strip_image_markers(text))
714        ));
715    }
716    lines.join("\n")
717}
718
719/// Strip the Markdown image placeholders a list item's own images fold into
720/// its text (the chunking serializer's `image_placeholder` is empty, so
721/// docling's chunk text carries no marker), collapsing the newlines that
722/// carried them.
723fn strip_image_markers(text: &str) -> String {
724    if !text.contains("<!-- image -->") {
725        return text.to_string();
726    }
727    let cleaned: Vec<&str> = text
728        .split('\n')
729        .map(str::trim_end)
730        .filter(|l| *l != "<!-- image -->")
731        .collect();
732    cleaned.join("\n").trim_end().to_string()
733}
734
735/// docling-core's `_humanize_text`: underscores to spaces, first letter
736/// capitalized (`line_chart` → `Line chart`).
737fn humanize_label(label: &str) -> String {
738    let text = label.replace('_', " ");
739    let mut chars = text.chars();
740    match chars.next() {
741        Some(f) => f.to_uppercase().collect::<String>() + chars.as_str(),
742        None => text,
743    }
744}
745
746/// docling's `TripletTableSerializer` over `export_to_dataframe` semantics: the
747/// leading rows carrying column-header cells become the dataframe's column
748/// names (multiple header rows join per column with `.`; no header rows at all
749/// yield pandas' integer column names), the rest are data rows; the dataframe
750/// is then rendered as `row, column = value` sentences (with the header-only /
751/// single-column special cases and the plain-text flatten fallback).
752fn triplet_table_text(t: &Table) -> String {
753    let rows: Vec<Vec<String>> = t
754        .rows
755        .iter()
756        .enumerate()
757        .map(|(ri, r)| (0..r.len()).map(|ci| cell_chunk_text(t, ri, ci)).collect())
758        .collect();
759    let num_rows = rows.len();
760    let num_cols = rows.iter().map(Vec::len).max().unwrap_or(0);
761    if num_rows == 0 || num_cols == 0 {
762        return String::new();
763    }
764    let cell = |r: usize, c: usize| -> &str {
765        rows.get(r)
766            .and_then(|row| row.get(c))
767            .map(String::as_str)
768            .unwrap_or("")
769    };
770
771    // The header block is the leading run of rows on which a column-header
772    // cell *starts* (docling-core#756): the grid replicates a spanning header
773    // into every row it covers, and the rows beneath it are data, not more
774    // header. Unlike the Markdown serializer there is no "no flags -> row 0"
775    // fallback here: pandas then gets integer column names.
776    let num_headers = {
777        let derived;
778        let cells: &[crate::TableCell] = match &t.cells {
779            Some(c) if !c.is_empty() => c,
780            _ => {
781                derived = t.derive_cells();
782                &derived
783            }
784        };
785        (0..num_rows)
786            .take_while(|&r| cells.iter().any(|c| c.column_header && c.start_row == r))
787            .count()
788    };
789
790    // Column names: header-row texts joined per column with '.', or the integer
791    // positions when there are no header rows.
792    let columns: Vec<String> = if num_headers > 0 {
793        (0..num_cols)
794            .map(|c| {
795                let mut name = String::new();
796                for r in 0..num_headers {
797                    if !name.is_empty() {
798                        name.push('.');
799                    }
800                    name.push_str(cell(r, c));
801                }
802                name
803            })
804            .collect()
805    } else {
806        (0..num_cols).map(|c| c.to_string()).collect()
807    };
808    let data_rows = num_headers..num_rows;
809    let n_data = data_rows.len();
810
811    // Header-only table: emit the header texts directly.
812    if n_data == 0 {
813        return columns
814            .iter()
815            .map(|s| s.trim())
816            .filter(|s| !s.is_empty())
817            .collect::<Vec<_>>()
818            .join(". ");
819    }
820
821    let data = |r: usize, c: usize| -> &str { cell(num_headers + r, c) };
822    let text = if num_cols == 1 {
823        // Single-column: the first data row is the column name, the rest are
824        // values (a single data row emits its cell text alone).
825        let col_name = data(0, 0).trim().to_string();
826        if n_data == 1 {
827            col_name
828        } else {
829            (1..n_data)
830                .map(|r| format!("{col_name} = {}", data(r, 0).trim()))
831                .collect::<Vec<_>>()
832                .join(". ")
833        }
834    } else {
835        // Triplets over the dataframe with the column names copied as row 0.
836        let mut parts = Vec::new();
837        for r in 0..n_data {
838            for (c, col_name) in columns.iter().enumerate().skip(1) {
839                parts.push(format!(
840                    "{}, {} = {}",
841                    data(r, 0).trim(),
842                    col_name.trim(),
843                    data(r, c).trim()
844                ));
845            }
846        }
847        parts.join(". ")
848    };
849    if !text.is_empty() {
850        return text;
851    }
852
853    // Last-resort flatten: the data rows' non-blank cells joined with '. '
854    // (the header rows are the dataframe's columns, so they are not included).
855    (0..n_data)
856        .flat_map(|r| (0..num_cols).map(move |c| (r, c)))
857        .map(|(r, c)| data(r, c).trim())
858        .filter(|s| !s.is_empty())
859        .collect::<Vec<_>>()
860        .join(". ")
861}
862
863/// Split a markdown-flavoured text into docling's inline-item granularity: a
864/// hyperlink / formatted span / inline formula is its own document item in
865/// docling's model, with plain text runs between them. Returns the standalone
866/// serialization of each item (`[text](url)`, `**bold**`, a formula re-wrapped
867/// as `$$latex$$`, plain text), or a single-element vector when the text is one
868/// uniform item. The single space docling's serializer inserts between inline
869/// items is stripped from the adjacent plain runs.
870fn inline_segments(md: &str) -> Vec<String> {
871    inline_segments_tagged(md)
872        .into_iter()
873        .map(|(t, _)| t)
874        .collect()
875}
876
877/// Like [`inline_segments`], with each segment tagged `true` when it came from
878/// plain (unmarked) text — those may still need splitting at run boundaries
879/// invisible in markdown (underline, soft breaks).
880fn inline_segments_tagged(md: &str) -> Vec<(String, bool)> {
881    let chars: Vec<char> = md.chars().collect();
882    let n = chars.len();
883    let find = |from: usize, pat: &str| -> Option<usize> {
884        let hay: String = chars[from..].iter().collect();
885        hay.find(pat).map(|p| from + hay[..p].chars().count())
886    };
887    let mut out: Vec<(String, bool)> = Vec::new();
888    let mut plain = String::new();
889    let mut after_span = false;
890
891    fn flush(
892        out: &mut Vec<(String, bool)>,
893        plain: &mut String,
894        before_span: bool,
895        after_span: bool,
896    ) {
897        let mut p = std::mem::take(plain);
898        if after_span {
899            if let Some(rest) = p.strip_prefix(' ') {
900                p = rest.to_string();
901            }
902        }
903        if before_span {
904            if let Some(rest) = p.strip_suffix(' ') {
905                p = rest.to_string();
906            }
907        }
908        if !p.is_empty() {
909            out.push((unescape_text(&p), true));
910        }
911    }
912
913    let mut i = 0;
914    while i < n {
915        let rest: String = chars[i..].iter().collect();
916        // A hyperlink span (not an image): the whole `[text](url)` is one item.
917        // The URL may itself contain balanced parentheses (`/Duck_(film)`),
918        // so the closing `)` is found by paren depth, not first match.
919        if chars[i] == '[' && !rest.starts_with("[](") {
920            // The label may itself contain balanced brackets (`[[ 1 ]](#ref)`
921            // has label `[ 1 ]`), but an *unbalanced* `[` means this bracket is
922            // plain text preceding a real link (`[ [*note*](url) ]`).
923            let balanced = |c: usize| {
924                let mut d = 0i32;
925                for &ch in &chars[i + 1..c] {
926                    match ch {
927                        '[' => d += 1,
928                        ']' => d -= 1,
929                        _ => {}
930                    }
931                }
932                d == 0
933            };
934            if let Some(close) = find(i + 1, "](").filter(|&c| balanced(c)) {
935                let mut depth = 0usize;
936                let mut url_end = None;
937                for (k, &c) in chars.iter().enumerate().skip(close + 2) {
938                    match c {
939                        '(' => depth += 1,
940                        ')' => {
941                            if depth == 0 {
942                                url_end = Some(k);
943                                break;
944                            }
945                            depth -= 1;
946                        }
947                        _ => {}
948                    }
949                }
950                if let Some(endp) = url_end {
951                    flush(&mut out, &mut plain, true, after_span);
952                    out.push((
953                        unescape_text(&chars[i..=endp].iter().collect::<String>()),
954                        false,
955                    ));
956                    i = endp + 1;
957                    after_span = true;
958                    continue;
959                }
960            }
961        }
962        // A formatted span; longest markers first. An inline code span is a
963        // code item in docling's model, whose standalone form is a fenced block.
964        let mut matched = false;
965        for marker in ["***", "**", "*", "~~", "`"] {
966            if rest.starts_with(marker) {
967                let mlen = marker.chars().count();
968                if let Some(end) = find(i + mlen, marker) {
969                    // A whitespace-only span (`* *` from a literal asterisk in
970                    // running text next to a real italic) is not a docling run:
971                    // treat the marker as plain text.
972                    let inner_blank = chars[i + mlen..end].iter().all(|c| c.is_whitespace());
973                    if end > i + mlen && !inner_blank {
974                        flush(&mut out, &mut plain, true, after_span);
975                        if marker == "`" {
976                            let inner: String = chars[i + 1..end].iter().collect();
977                            out.push((format!("```\n{}\n```", unescape_text(&inner)), false));
978                        } else {
979                            out.push((
980                                unescape_text(&chars[i..end + mlen].iter().collect::<String>()),
981                                false,
982                            ));
983                        }
984                        i = end + mlen;
985                        after_span = true;
986                        matched = true;
987                    }
988                }
989                break;
990            }
991        }
992        if matched {
993            continue;
994        }
995        // A literal `$$` inside running text is not an inline formula: copy it
996        // through as plain characters.
997        if rest.starts_with("$$") {
998            plain.push_str("$$");
999            i += 2;
1000            continue;
1001        }
1002        // An inline formula: standalone it re-serializes in display form.
1003        if chars[i] == '$' {
1004            if let Some(end) = find(i + 1, "$") {
1005                if end > i + 1 {
1006                    flush(&mut out, &mut plain, true, after_span);
1007                    let latex: String = chars[i + 1..end].iter().collect();
1008                    out.push((format!("$${latex}$$"), false));
1009                    i = end + 1;
1010                    after_span = true;
1011                    continue;
1012                }
1013            }
1014        }
1015        plain.push(chars[i]);
1016        i += 1;
1017    }
1018    flush(&mut out, &mut plain, false, after_span);
1019    if out.is_empty() {
1020        out.push((unescape_text(md), true));
1021    }
1022    out
1023}
1024
1025/// Split a plain markdown segment at run boundaries the markdown cannot show
1026/// (an underlined run, a `<sub>`/`<sup>` run): when a consecutive window of
1027/// two or more unmarked runs exactly covers the segment, each run is its own
1028/// document item.
1029fn split_plain_by_runs(segment: &str, runs: &[crate::InlineRun]) -> Option<Vec<String>> {
1030    let target = segment.trim();
1031    if target.is_empty() {
1032        return None;
1033    }
1034    let plainish =
1035        |r: &crate::InlineRun| !r.bold && !r.italic && !r.strike && !r.code && !r.formula;
1036    let fully_plain =
1037        |r: &crate::InlineRun| plainish(r) && !r.underline && r.script == crate::Script::Baseline;
1038    let unmarked: Vec<(&str, bool)> = runs
1039        .iter()
1040        .filter(|r| plainish(r))
1041        .map(|r| (r.text.as_str(), fully_plain(r)))
1042        .collect();
1043    for start in 0..unmarked.len() {
1044        let mut rest = target;
1045        let mut taken: Vec<(String, bool)> = Vec::new();
1046        for (t, fully) in &unmarked[start..] {
1047            let t = t.trim();
1048            if t.is_empty() {
1049                continue;
1050            }
1051            match rest.strip_prefix(t) {
1052                Some(r) => {
1053                    taken.push((unescape_text(t), *fully));
1054                    rest = r.trim_start();
1055                    if rest.is_empty() {
1056                        break;
1057                    }
1058                }
1059                None => break,
1060            }
1061        }
1062        if rest.is_empty() && taken.len() >= 2 {
1063            // Consecutive fully-plain pieces (no underline / sub / sup) are one
1064            // annotation in docling's model — `simplify_text_elements` merges
1065            // them joined with a space, so they never split between themselves.
1066            // Only an underline / sub / sup run is a genuine boundary.
1067            let mut merged: Vec<(String, bool)> = Vec::new();
1068            for (t, fully) in taken {
1069                match merged.last_mut() {
1070                    Some((last, true)) if fully => {
1071                        last.push(' ');
1072                        last.push_str(&t);
1073                    }
1074                    _ => merged.push((t, fully)),
1075                }
1076            }
1077            if merged.len() >= 2 {
1078                return Some(merged.into_iter().map(|(t, _)| t).collect());
1079            }
1080            return None;
1081        }
1082    }
1083    None
1084}
1085
1086/// A table cell's text for the triplet serializer. A *rich* cell (one carrying
1087/// block content) is re-serialized the way docling's chunking serializer sees
1088/// it: paragraphs joined with blank lines, a nested table as its own triplet
1089/// sentences, pictures as an empty placeholder. Plain cells use the flat text
1090/// with the markdown image placeholder stripped (the chunking serializer's
1091/// `image_placeholder` is empty).
1092fn cell_chunk_text(t: &Table, r: usize, c: usize) -> String {
1093    if let Some(blocks) = t
1094        .cell_blocks
1095        .as_ref()
1096        .and_then(|b| b.get(r))
1097        .and_then(|row| row.get(c))
1098        .filter(|b| !b.is_empty())
1099    {
1100        let mut parts: Vec<String> = Vec::new();
1101        for node in blocks.iter() {
1102            let part = block_chunk_text(node);
1103            if !part.is_empty() {
1104                parts.push(part);
1105            }
1106        }
1107        return parts.join("\n\n");
1108    }
1109    let flat = t
1110        .rows
1111        .get(r)
1112        .and_then(|row| row.get(c))
1113        .map(String::as_str)
1114        .unwrap_or("");
1115    unescape_text(flat)
1116        .replace("<!-- image -->", "")
1117        .trim()
1118        .to_string()
1119}
1120
1121/// One block of a rich cell, serialized for chunking.
1122fn block_chunk_text(node: &Node) -> String {
1123    match node {
1124        Node::Paragraph { text } => unescape_text(text),
1125        Node::LabeledText { .. } => block_chunk_text(&node.labeled_as_paragraph()),
1126        Node::InlineGroup { md_text, .. } => unescape_text(md_text),
1127        Node::Code { text, .. } => format!("```\n{}\n```", unescape_text(text)),
1128        Node::Table(inner) => triplet_table_text(inner),
1129        Node::Picture { caption, .. } => caption
1130            .as_deref()
1131            .filter(|c| !c.is_empty())
1132            .map(unescape_text)
1133            .unwrap_or_default(),
1134        Node::ListItem {
1135            ordered,
1136            number,
1137            text,
1138            ..
1139        } => {
1140            let marker = if *ordered {
1141                format!("{number}.")
1142            } else {
1143                "-".to_string()
1144            };
1145            format!("{marker} {}", unescape_text(text))
1146        }
1147        Node::CheckboxItem { checked, text } => {
1148            let mark = if *checked { "- [x] " } else { "- [ ] " };
1149            format!("{mark}{}", unescape_text(text))
1150        }
1151        Node::Heading { text, .. } => unescape_text(text),
1152        Node::Located { inner, .. }
1153        | Node::Prov { inner, .. }
1154        | Node::Track { inner, .. }
1155        | Node::Commented { inner, .. } => block_chunk_text(inner),
1156        Node::Group { layer: Some(_), .. } => String::new(),
1157        Node::Group { children, .. } => children
1158            .iter()
1159            .map(block_chunk_text)
1160            .filter(|s| !s.is_empty())
1161            .collect::<Vec<_>>()
1162            .join("\n"),
1163        _ => String::new(),
1164    }
1165}
1166
1167/// Reverse the model's baked markdown text escaping — same mapping as the JSON
1168/// exporter (docling chunks carry raw text).
1169fn unescape_text(s: &str) -> String {
1170    s.replace("&lt;", "<")
1171        .replace("&gt;", ">")
1172        .replace("&amp;", "&")
1173        .replace("\\_", "_")
1174}
1175
1176// ---------------------------------------------------------------------------
1177// Hybrid chunker
1178// ---------------------------------------------------------------------------
1179
1180/// Token counting for [`HybridChunker`] — docling's `BaseTokenizer`.
1181pub trait ChunkTokenizer {
1182    /// Number of tokens in `text` (no special tokens).
1183    fn count_tokens(&self, text: &str) -> usize;
1184    /// The chunk budget (docling's `max_tokens`, e.g. 256 for MiniLM).
1185    fn max_tokens(&self) -> usize;
1186}
1187
1188/// Tokenization-aware chunker on top of [`HierarchicalChunker`] — docling's
1189/// `HybridChunker` with default parameters (`merge_peers`,
1190/// `repeat_table_header` on; `omit_header_on_overflow` off).
1191pub struct HybridChunker<T: ChunkTokenizer> {
1192    tokenizer: T,
1193    merge_peers: bool,
1194}
1195
1196impl<T: ChunkTokenizer> HybridChunker<T> {
1197    pub fn new(tokenizer: T) -> Self {
1198        Self {
1199            tokenizer,
1200            merge_peers: true,
1201        }
1202    }
1203
1204    /// Disable merging of undersized same-heading neighbours.
1205    pub fn with_merge_peers(mut self, merge_peers: bool) -> Self {
1206        self.merge_peers = merge_peers;
1207        self
1208    }
1209
1210    pub fn max_tokens(&self) -> usize {
1211        self.tokenizer.max_tokens()
1212    }
1213
1214    /// Chunk the document.
1215    pub fn chunk(&self, doc: &DoclingDocument) -> Vec<DocChunk> {
1216        let mut chunks = Vec::new();
1217        self.chunk_with(doc, &mut |c| {
1218            chunks.push(c);
1219            true
1220        });
1221        chunks
1222    }
1223
1224    /// Stream the chunks: each hierarchical chunk is split against the token
1225    /// budget as the document walk produces it, and the peer merge flushes a
1226    /// merged chunk to `sink` as soon as its window closes (a chunk with
1227    /// different headings arrives, or the budget fills). A `false` return from
1228    /// `sink` cancels the chunking. [`Self::chunk`] is this with a collecting
1229    /// sink — the chunks and their order are identical.
1230    pub fn chunk_with(&self, doc: &DoclingDocument, sink: &mut dyn FnMut(DocChunk) -> bool) {
1231        let mut merger = PeerMerger::default();
1232        let mut alive = true;
1233        HierarchicalChunker.chunk_with(doc, &mut |c| {
1234            for split in self.split_by_doc_items(c) {
1235                for chunk in self.split_using_plain_text(split) {
1236                    if !alive {
1237                        return false;
1238                    }
1239                    alive = if self.merge_peers {
1240                        self.merge_push(&mut merger, chunk, sink)
1241                    } else {
1242                        sink(chunk)
1243                    };
1244                }
1245            }
1246            alive
1247        });
1248        if alive {
1249            self.merge_flush(&mut merger, sink);
1250        }
1251    }
1252
1253    fn count_chunk_tokens(&self, chunk: &DocChunk) -> usize {
1254        self.tokenizer.count_tokens(&contextualize(chunk))
1255    }
1256
1257    /// docling's `_make_chunk_from_doc_items`: single-item chunks keep their
1258    /// text; multi-item windows re-join the items' standalone serializations.
1259    fn window_chunk(&self, chunk: &DocChunk, start: usize, end: usize) -> DocChunk {
1260        let doc_items: Vec<ChunkItem> = chunk.doc_items[start..=end].to_vec();
1261        let text = if chunk.doc_items.len() == 1 {
1262            chunk.text.clone()
1263        } else {
1264            doc_items
1265                .iter()
1266                .filter(|it| !it.text.is_empty())
1267                .map(|it| it.text.as_str())
1268                .collect::<Vec<_>>()
1269                .join("\n")
1270        };
1271        DocChunk {
1272            text,
1273            headings: chunk.headings.clone(),
1274            doc_items,
1275        }
1276    }
1277
1278    fn split_by_doc_items(&self, chunk: DocChunk) -> Vec<DocChunk> {
1279        if chunk.doc_items.is_empty() {
1280            return vec![chunk];
1281        }
1282        let max = self.max_tokens();
1283        let num_items = chunk.doc_items.len();
1284        let mut chunks = Vec::new();
1285        let mut window_start = 0usize;
1286        let mut window_end = 0usize; // inclusive
1287        while window_end < num_items {
1288            let mut new_chunk = self.window_chunk(&chunk, window_start, window_end);
1289            if self.count_chunk_tokens(&new_chunk) <= max {
1290                if window_end < num_items - 1 {
1291                    window_end += 1;
1292                    continue;
1293                } else {
1294                    window_end = num_items; // last loop
1295                }
1296            } else if window_start == window_end {
1297                // One item that doesn't fit: keep it; the plain-text splitter
1298                // takes over.
1299                window_end += 1;
1300                window_start = window_end;
1301            } else {
1302                // The window without its last item fit; flush that and start a
1303                // new window at the current item.
1304                new_chunk = self.window_chunk(&chunk, window_start, window_end - 1);
1305                window_start = window_end;
1306            }
1307            chunks.push(new_chunk);
1308        }
1309        chunks
1310    }
1311
1312    fn split_using_plain_text(&self, chunk: DocChunk) -> Vec<DocChunk> {
1313        let total = self.count_chunk_tokens(&chunk);
1314        let max = self.max_tokens();
1315        if total <= max {
1316            return vec![chunk];
1317        }
1318        let text_len = self.tokenizer.count_tokens(&chunk.text);
1319        let other_len = total - text_len;
1320        if other_len >= max {
1321            // Headings alone exceed the budget: drop them and retry.
1322            let stripped = DocChunk {
1323                headings: None,
1324                ..chunk
1325            };
1326            return self.split_using_plain_text(stripped);
1327        }
1328        let available = max - other_len;
1329
1330        let segments =
1331            if chunk.doc_items.len() == 1 && chunk.doc_items[0].kind == ChunkItemKind::Table {
1332                // Table: split line-based, repeating headers. The triplet
1333                // serializer has no header lines, so this is a line-preserving
1334                // split of the table text. (docling constructs the line chunker
1335                // with the *tokenizer's* max_tokens — the `max_tokens=available`
1336                // argument is silently dropped by pydantic — so the line budget is
1337                // the full window, not `available`.)
1338                let lines: Vec<String> = chunk
1339                    .text
1340                    .split('\n')
1341                    .filter(|l| !l.trim().is_empty())
1342                    .map(|l| l.to_string())
1343                    .collect();
1344                line_chunk_text(&lines, &self.tokenizer, max)
1345            } else {
1346                semchunk(&chunk.text, available, &self.tokenizer)
1347            };
1348        segments
1349            .into_iter()
1350            .map(|s| DocChunk {
1351                text: s,
1352                headings: chunk.headings.clone(),
1353                doc_items: chunk.doc_items.clone(),
1354            })
1355            .collect()
1356    }
1357
1358    /// One step of docling's `_merge_chunks_with_matching_metadata`, streamed:
1359    /// extend the window with `chunk` when its headings match the window's and
1360    /// the merged candidate stays within budget, otherwise flush the window to
1361    /// `sink` and start a new one at `chunk`. Returns `false` once the sink
1362    /// cancels.
1363    fn merge_push(
1364        &self,
1365        m: &mut PeerMerger,
1366        chunk: DocChunk,
1367        sink: &mut dyn FnMut(DocChunk) -> bool,
1368    ) -> bool {
1369        if m.window.is_empty() {
1370            m.window.push(chunk);
1371            return true;
1372        }
1373        let candidate = DocChunk {
1374            text: m
1375                .window
1376                .iter()
1377                .map(|c| c.text.as_str())
1378                .chain([chunk.text.as_str()])
1379                .collect::<Vec<_>>()
1380                .join("\n"),
1381            headings: m.window[0].headings.clone(),
1382            doc_items: m
1383                .window
1384                .iter()
1385                .flat_map(|c| c.doc_items.iter().cloned())
1386                .chain(chunk.doc_items.iter().cloned())
1387                .collect(),
1388        };
1389        if chunk.headings == m.window[0].headings
1390            && self.count_chunk_tokens(&candidate) <= self.max_tokens()
1391        {
1392            m.window.push(chunk);
1393            m.merged = Some(candidate);
1394            true
1395        } else {
1396            let alive = self.merge_flush(m, sink);
1397            m.window.push(chunk);
1398            alive
1399        }
1400    }
1401
1402    /// Flush the merge window: a single chunk passes through unchanged, a
1403    /// multi-chunk window emits its precomputed merge. Returns `false` once
1404    /// the sink cancels.
1405    fn merge_flush(&self, m: &mut PeerMerger, sink: &mut dyn FnMut(DocChunk) -> bool) -> bool {
1406        let alive = if m.window.len() == 1 {
1407            sink(m.window.pop().expect("single-chunk window"))
1408        } else if !m.window.is_empty() {
1409            m.window.clear();
1410            sink(m.merged.take().expect("multi-chunk window has a merge"))
1411        } else {
1412            true
1413        };
1414        m.merged = None;
1415        alive
1416    }
1417}
1418
1419/// The in-flight peer-merge window of [`HybridChunker::chunk_with`].
1420#[derive(Default)]
1421struct PeerMerger {
1422    window: Vec<DocChunk>,
1423    merged: Option<DocChunk>,
1424}
1425
1426// ---------------------------------------------------------------------------
1427// Line-based token chunking (docling's LineBasedTokenChunker, empty prefix)
1428// ---------------------------------------------------------------------------
1429
1430/// Pack lines into chunks of at most `max_tokens`, splitting a line only when
1431/// it exceeds the budget on its own — docling's `LineBasedTokenChunker
1432/// .chunk_text` with an empty prefix (which is what the triplet table
1433/// serializer yields). Reproduces its exact output, including the `\n` it
1434/// prepends to a carried-over segment of an oversized line.
1435fn line_chunk_text<T: ChunkTokenizer>(lines: &[String], tok: &T, max_tokens: usize) -> Vec<String> {
1436    let mut chunks: Vec<String> = Vec::new();
1437    let mut current = String::new();
1438    let mut current_len = 0usize;
1439
1440    for line in lines {
1441        let mut remaining: Vec<char> = line.chars().collect();
1442        loop {
1443            let rem_str: String = remaining.iter().collect();
1444            let line_tokens = tok.count_tokens(&rem_str);
1445            let available = max_tokens.saturating_sub(current_len);
1446
1447            if line_tokens <= available {
1448                current.push_str(&rem_str);
1449                current_len += line_tokens;
1450                break;
1451            }
1452            if line_tokens <= max_tokens {
1453                chunks.push(std::mem::take(&mut current));
1454                current_len = 0;
1455                continue;
1456            }
1457            // Too large even for an empty chunk: split off what fits.
1458            let (mut take, rest) = split_by_token_limit(&remaining, available, tok);
1459            let mut rest = rest;
1460            if take.is_empty() {
1461                if rest.is_empty() {
1462                    break;
1463                }
1464                take = rest[..1].iter().collect();
1465                rest = rest[1..].to_vec();
1466            }
1467            current.push('\n');
1468            current.push_str(&take);
1469            chunks.push(std::mem::take(&mut current));
1470            current_len = 0;
1471            remaining = rest;
1472        }
1473    }
1474    if !current.is_empty() {
1475        chunks.push(current);
1476    }
1477    chunks
1478}
1479
1480/// Binary-search the longest char-prefix of `text` within `token_limit`
1481/// tokens, preferring to break at the last ASCII space — docling's
1482/// `split_by_token_limit`.
1483fn split_by_token_limit<T: ChunkTokenizer>(
1484    text: &[char],
1485    token_limit: usize,
1486    tok: &T,
1487) -> (String, Vec<char>) {
1488    if token_limit == 0 || text.is_empty() {
1489        return (String::new(), text.to_vec());
1490    }
1491    let full: String = text.iter().collect();
1492    if tok.count_tokens(&full) <= token_limit {
1493        return (full, Vec::new());
1494    }
1495    let (mut lo, mut hi) = (0usize, text.len());
1496    let mut best: Option<usize> = None;
1497    while lo <= hi {
1498        let mid = (lo + hi) / 2;
1499        let head: String = text[..mid].iter().collect();
1500        if tok.count_tokens(&head) <= token_limit {
1501            best = Some(mid);
1502            lo = mid + 1;
1503        } else {
1504            if mid == 0 {
1505                break;
1506            }
1507            hi = mid - 1;
1508        }
1509    }
1510    let mut best_idx = match best {
1511        Some(b) if b > 0 => b,
1512        _ => return (String::new(), text.to_vec()),
1513    };
1514    // Snap back to the last space, if that leaves a non-empty head.
1515    if let Some(pos) = text[..best_idx].iter().rposition(|c| *c == ' ') {
1516        if pos > 0 {
1517            best_idx = pos;
1518        }
1519    }
1520    (text[..best_idx].iter().collect(), text[best_idx..].to_vec())
1521}
1522
1523// ---------------------------------------------------------------------------
1524// semchunk port (the plain-text splitter HybridChunker delegates to)
1525// ---------------------------------------------------------------------------
1526
1527/// Semantically meaningful non-whitespace splitters, most desirable first.
1528const NON_WS_SPLITTERS: &[&str] = &[
1529    ".", "?", "!", "*", ";", ",", "(", ")", "[", "]", "\u{201c}", "\u{201d}", "\u{2018}",
1530    "\u{2019}", "'", "\"", "`", ":", "\u{2014}", "\u{2026}", "/", "\\", "\u{2013}", "&", "-",
1531];
1532
1533/// Split `text` into chunks of at most `chunk_size` tokens using the most
1534/// semantically meaningful splitter available — the `semchunk` algorithm
1535/// docling's HybridChunker delegates plain-text splitting to.
1536pub fn semchunk<T: ChunkTokenizer>(text: &str, chunk_size: usize, tok: &T) -> Vec<String> {
1537    let mut cache: std::collections::HashMap<String, usize> = std::collections::HashMap::new();
1538    let mut counter = |s: &str| -> usize {
1539        if let Some(n) = cache.get(s) {
1540            return *n;
1541        }
1542        let n = tok.count_tokens(s);
1543        cache.insert(s.to_string(), n);
1544        n
1545    };
1546    let chunks = semchunk_rec(text, chunk_size, &mut counter);
1547    // top-level: drop empty / all-whitespace chunks
1548    chunks
1549        .into_iter()
1550        .filter(|c| !c.is_empty() && !c.chars().all(char::is_whitespace))
1551        .collect()
1552}
1553
1554/// One recursion level of semchunk: split, merge windows back up to size, and
1555/// recurse into oversized splits.
1556fn semchunk_rec(
1557    text: &str,
1558    chunk_size: usize,
1559    counter: &mut dyn FnMut(&str) -> usize,
1560) -> Vec<String> {
1561    let (splitter, splitter_is_ws, splits) = split_text(text);
1562
1563    let split_lens: Vec<usize> = splits.iter().map(|s| s.chars().count()).collect();
1564    let mut cum_lens = Vec::with_capacity(splits.len() + 1);
1565    cum_lens.push(0usize);
1566    for l in &split_lens {
1567        cum_lens.push(cum_lens.last().unwrap() + l);
1568    }
1569    let num_splits_plus_one = splits.len() + 1;
1570
1571    let mut chunks: Vec<String> = Vec::new();
1572    let mut skips: std::collections::HashSet<usize> = std::collections::HashSet::new();
1573
1574    for i in 0..splits.len() {
1575        if skips.contains(&i) {
1576            continue;
1577        }
1578        let split = &splits[i];
1579        if counter(split) > chunk_size {
1580            let inner = semchunk_rec(split, chunk_size, counter);
1581            chunks.extend(inner);
1582        } else {
1583            let (end, merged) = merge_splits(
1584                &splits,
1585                &cum_lens,
1586                chunk_size,
1587                &splitter,
1588                counter,
1589                i,
1590                num_splits_plus_one,
1591            );
1592            for j in (i + 1)..end {
1593                skips.insert(j);
1594            }
1595            chunks.push(merged);
1596        }
1597        // Re-attach a non-whitespace splitter to the last chunk (or emit it as
1598        // its own chunk if it doesn't fit).
1599        let is_last = i == splits.len() - 1 || ((i + 1)..splits.len()).all(|j| skips.contains(&j));
1600        if !splitter_is_ws && !is_last {
1601            let with_splitter = format!(
1602                "{}{}",
1603                chunks.last().map(String::as_str).unwrap_or(""),
1604                splitter
1605            );
1606            if counter(&with_splitter) <= chunk_size {
1607                if let Some(last) = chunks.last_mut() {
1608                    *last = with_splitter;
1609                } else {
1610                    chunks.push(with_splitter);
1611                }
1612            } else {
1613                chunks.push(splitter.clone());
1614            }
1615        }
1616    }
1617    chunks
1618}
1619
1620/// docling/semchunk's `merge_splits`: extend the window with a cum-length-guided
1621/// binary search until the token budget is hit.
1622fn merge_splits(
1623    splits: &[String],
1624    cum_lens: &[usize],
1625    chunk_size: usize,
1626    splitter: &str,
1627    counter: &mut dyn FnMut(&str) -> usize,
1628    start: usize,
1629    high_init: usize,
1630) -> (usize, String) {
1631    let mut average = 0.2f64;
1632    let mut low = start;
1633    let mut high = high_init;
1634    let offset = cum_lens[start];
1635    let mut target = offset as f64 + (chunk_size as f64 * average);
1636
1637    while low < high {
1638        let i = bisect_left(cum_lens, target, low, high);
1639        let midpoint = i.min(high - 1);
1640        let joined = splits[start..midpoint.max(start)].join(splitter);
1641        let tokens = counter(&joined);
1642        let local_cum = cum_lens[midpoint] - offset;
1643        if local_cum > 0 && tokens > 0 {
1644            average = local_cum as f64 / tokens as f64;
1645            target = offset as f64 + (chunk_size as f64 * average);
1646        }
1647        if tokens > chunk_size {
1648            high = midpoint;
1649        } else {
1650            low = midpoint + 1;
1651        }
1652    }
1653    let end = low - 1;
1654    (end, splits[start..end.max(start)].join(splitter))
1655}
1656
1657fn bisect_left(sorted: &[usize], target: f64, mut low: usize, mut high: usize) -> usize {
1658    while low < high {
1659        let mid = (low + high) / 2;
1660        if (sorted[mid] as f64) < target {
1661            low = mid + 1;
1662        } else {
1663            high = mid;
1664        }
1665    }
1666    low
1667}
1668
1669/// semchunk's `_split_text`: pick the most desirable splitter present.
1670fn split_text(text: &str) -> (String, bool, Vec<String>) {
1671    // Longest run of newlines/carriage returns.
1672    if text.contains('\n') || text.contains('\r') {
1673        let splitter = longest_run(text, |c| c == '\n' || c == '\r');
1674        return (splitter.clone(), true, split_on(text, &splitter));
1675    }
1676    // Longest run of tabs.
1677    if text.contains('\t') {
1678        let splitter = longest_run(text, |c| c == '\t');
1679        return (splitter.clone(), true, split_on(text, &splitter));
1680    }
1681    // Longest run of whitespace.
1682    if text.chars().any(char::is_whitespace) {
1683        let splitter = longest_run(text, char::is_whitespace);
1684        if splitter.chars().count() == 1 {
1685            // Prefer a whitespace char preceded by a semantic splitter.
1686            for preceder in NON_WS_SPLITTERS {
1687                if let Some((ws, parts)) = split_after_preceder(text, preceder) {
1688                    return (ws, true, parts);
1689                }
1690            }
1691        }
1692        return (splitter.clone(), true, split_on(text, &splitter));
1693    }
1694    // Most desirable semantic splitter present.
1695    for s in NON_WS_SPLITTERS {
1696        if text.contains(s) {
1697            return (s.to_string(), false, split_on(text, s));
1698        }
1699    }
1700    // No splitter at all: split into characters.
1701    (
1702        String::new(),
1703        true,
1704        text.chars().map(|c| c.to_string()).collect(),
1705    )
1706}
1707
1708/// The longest maximal run of chars matching `pred` (first one wins ties).
1709fn longest_run(text: &str, pred: impl Fn(char) -> bool) -> String {
1710    let mut best = String::new();
1711    let mut cur = String::new();
1712    for c in text.chars() {
1713        if pred(c) {
1714            cur.push(c);
1715        } else {
1716            if cur.chars().count() > best.chars().count() {
1717                best = cur.clone();
1718            }
1719            cur.clear();
1720        }
1721    }
1722    if cur.chars().count() > best.chars().count() {
1723        best = cur;
1724    }
1725    best
1726}
1727
1728fn split_on(text: &str, splitter: &str) -> Vec<String> {
1729    text.split(splitter).map(str::to_string).collect()
1730}
1731
1732/// Python: `re.search(rf'{p}(\s)', text)` → the first whitespace char preceded
1733/// by `p`; then `re.split(rf'(?<={p}){s}', text)` — split at every occurrence
1734/// of that whitespace char immediately preceded by `p`.
1735fn split_after_preceder(text: &str, preceder: &str) -> Option<(String, Vec<String>)> {
1736    let chars: Vec<char> = text.chars().collect();
1737    let p: Vec<char> = preceder.chars().collect();
1738    let mut ws: Option<char> = None;
1739    for i in p.len()..chars.len() {
1740        if chars[i].is_whitespace() && chars[i - p.len()..i] == p[..] {
1741            ws = Some(chars[i]);
1742            break;
1743        }
1744    }
1745    let ws = ws?;
1746    let mut parts = Vec::new();
1747    let mut cur = String::new();
1748    let mut i = 0usize;
1749    while i < chars.len() {
1750        if chars[i] == ws && i >= p.len() && chars[i - p.len()..i] == p[..] {
1751            parts.push(std::mem::take(&mut cur));
1752            i += 1;
1753            continue;
1754        }
1755        cur.push(chars[i]);
1756        i += 1;
1757    }
1758    parts.push(cur);
1759    Some((ws.to_string(), parts))
1760}
1761
1762// ---------------------------------------------------------------------------
1763// HuggingFace tokenizer (feature `chunking`)
1764// ---------------------------------------------------------------------------
1765
1766#[cfg(feature = "chunking")]
1767mod hf {
1768    use super::ChunkTokenizer;
1769
1770    /// Where `scripts/install/download_dependencies.sh` puts the hybrid
1771    /// chunker's default tokenizer (all-MiniLM-L6-v2's `tokenizer.json`),
1772    /// relative to the process's working directory — the same convention as
1773    /// the `.models/` ONNX files.
1774    pub const DEFAULT_TOKENIZER_PATH: &str = ".models/chunk/tokenizer.json";
1775
1776    /// Resolve the tokenizer path for the hybrid chunker: an explicit path
1777    /// wins; otherwise fall back to [`DEFAULT_TOKENIZER_PATH`] when it exists
1778    /// on disk. Errors with the download instructions when neither is
1779    /// available.
1780    pub fn resolve_tokenizer_path(explicit: Option<&str>) -> Result<String, String> {
1781        if let Some(p) = explicit {
1782            return Ok(p.to_string());
1783        }
1784        let resolved = crate::assets::resolve(DEFAULT_TOKENIZER_PATH);
1785        if std::path::Path::new(&resolved).exists() {
1786            return Ok(resolved);
1787        }
1788        Err(format!(
1789            "the hybrid chunker needs a HuggingFace tokenizer.json: none passed and \
1790             {DEFAULT_TOKENIZER_PATH} does not exist — run \
1791             scripts/install/download_dependencies.sh (or pass an explicit path)"
1792        ))
1793    }
1794
1795    /// [`ChunkTokenizer`] backed by a HuggingFace `tokenizer.json` — the Rust
1796    /// analogue of docling's `HuggingFaceTokenizer` (whose default is
1797    /// `sentence-transformers/all-MiniLM-L6-v2` with `max_tokens` 256).
1798    pub struct HuggingFaceTokenizer {
1799        tok: tokenizers::Tokenizer,
1800        max_tokens: usize,
1801    }
1802
1803    impl HuggingFaceTokenizer {
1804        /// Load the tokenizer from an explicit path, or from
1805        /// [`DEFAULT_TOKENIZER_PATH`] when `path` is `None` (see
1806        /// [`resolve_tokenizer_path`]).
1807        pub fn resolve(path: Option<&str>, max_tokens: usize) -> Result<Self, String> {
1808            Self::from_file(resolve_tokenizer_path(path)?, max_tokens)
1809        }
1810
1811        /// Load a `tokenizer.json`. `max_tokens` is the chunk budget (docling
1812        /// resolves it from the model's `sentence_bert_config.json`; for the
1813        /// default MiniLM model that is 256).
1814        pub fn from_file(
1815            path: impl AsRef<std::path::Path>,
1816            max_tokens: usize,
1817        ) -> Result<Self, String> {
1818            let mut tok = tokenizers::Tokenizer::from_file(path.as_ref())
1819                .map_err(|e| format!("failed to load tokenizer: {e}"))?;
1820            // Counting must see the full text (docling tokenizes without
1821            // truncation, padding, or special tokens — MiniLM's tokenizer.json
1822            // ships with fixed-length padding enabled, which would make every
1823            // short string count as the padded length).
1824            let _ = tok.with_truncation(None);
1825            tok.with_padding(None);
1826            Ok(Self { tok, max_tokens })
1827        }
1828    }
1829
1830    impl ChunkTokenizer for HuggingFaceTokenizer {
1831        fn count_tokens(&self, text: &str) -> usize {
1832            self.tok
1833                .encode(text, false)
1834                .map(|e| e.get_tokens().len())
1835                .unwrap_or(0)
1836        }
1837        fn max_tokens(&self) -> usize {
1838            self.max_tokens
1839        }
1840    }
1841}
1842
1843#[cfg(feature = "chunking")]
1844pub use hf::{resolve_tokenizer_path, HuggingFaceTokenizer, DEFAULT_TOKENIZER_PATH};
1845
1846// ---------------------------------------------------------------------------
1847// Window chunker (feature `chunking`) — docling-rag's Markdown window chunker
1848// ---------------------------------------------------------------------------
1849
1850#[cfg(feature = "chunking")]
1851mod window {
1852    use super::DocChunk;
1853    use pulldown_cmark::{Event, HeadingLevel, Parser, Tag, TagEnd};
1854
1855    /// A contiguous run of body text under a heading path.
1856    #[derive(Debug, Clone, Default)]
1857    pub struct Section {
1858        /// The heading stack in effect for this section, outermost first
1859        /// (e.g. `["Guide", "Setup"]`). Empty for pre-heading / body-only text.
1860        pub heading_path: Vec<String>,
1861        /// The plain words of the section body, markup stripped.
1862        pub words: Vec<String>,
1863    }
1864
1865    impl Section {
1866        /// The heading path rendered as a single context line, e.g.
1867        /// `# Guide > Setup`. Empty string when there is no heading.
1868        pub fn heading_context(&self) -> String {
1869            if self.heading_path.is_empty() {
1870                String::new()
1871            } else {
1872                format!("# {}", self.heading_path.join(" > "))
1873            }
1874        }
1875    }
1876
1877    fn level_index(level: HeadingLevel) -> usize {
1878        match level {
1879            HeadingLevel::H1 => 1,
1880            HeadingLevel::H2 => 2,
1881            HeadingLevel::H3 => 3,
1882            HeadingLevel::H4 => 4,
1883            HeadingLevel::H5 => 5,
1884            HeadingLevel::H6 => 6,
1885        }
1886    }
1887
1888    /// Parse Markdown into heading-bounded sections. A new section starts at
1889    /// every heading; the heading path is maintained as a stack keyed by
1890    /// heading level.
1891    pub fn parse_sections(markdown: &str) -> Vec<Section> {
1892        parse_sections_with_stack(markdown, Vec::new()).0
1893    }
1894
1895    /// [`parse_sections`] with an explicit initial heading stack, returning the
1896    /// final stack — lets a streaming caller carry heading context across
1897    /// pieces. `heading_stack[i]` holds the current heading text at level `i+1`
1898    /// (may be empty when a level was skipped).
1899    pub fn parse_sections_with_stack(
1900        markdown: &str,
1901        initial_stack: Vec<String>,
1902    ) -> (Vec<Section>, Vec<String>) {
1903        let mut heading_stack: Vec<String> = initial_stack;
1904        let mut sections: Vec<Section> = Vec::new();
1905        // Text before the first heading of this piece continues the carried-over
1906        // section, so it keeps the heading path in effect at the split point.
1907        let mut current = Section {
1908            heading_path: heading_stack
1909                .iter()
1910                .filter(|h| !h.is_empty())
1911                .cloned()
1912                .collect(),
1913            words: Vec::new(),
1914        };
1915
1916        let mut in_heading = false;
1917        let mut heading_level = 0usize;
1918        let mut heading_buf = String::new();
1919
1920        let push_words = |section: &mut Section, text: &str| {
1921            for w in text.split_whitespace() {
1922                section.words.push(w.to_string());
1923            }
1924        };
1925
1926        let flush = |sections: &mut Vec<Section>, section: &mut Section| {
1927            if !section.words.is_empty() {
1928                sections.push(std::mem::take(section));
1929            } else {
1930                *section = Section::default();
1931            }
1932        };
1933
1934        for event in Parser::new(markdown) {
1935            match event {
1936                Event::Start(Tag::Heading { level, .. }) => {
1937                    in_heading = true;
1938                    heading_level = level_index(level);
1939                    heading_buf.clear();
1940                }
1941                Event::End(TagEnd::Heading(_)) => {
1942                    in_heading = false;
1943                    // Update the heading stack: set this level, drop anything deeper.
1944                    let idx = heading_level.saturating_sub(1);
1945                    if heading_stack.len() <= idx {
1946                        heading_stack.resize(idx + 1, String::new());
1947                    } else {
1948                        heading_stack.truncate(idx + 1);
1949                    }
1950                    heading_stack[idx] = heading_buf.trim().to_string();
1951                    // A heading begins a new section.
1952                    flush(&mut sections, &mut current);
1953                    current.heading_path = heading_stack
1954                        .iter()
1955                        .filter(|h| !h.is_empty())
1956                        .cloned()
1957                        .collect();
1958                }
1959                Event::Text(t) | Event::Code(t) => {
1960                    if in_heading {
1961                        if !heading_buf.is_empty() {
1962                            heading_buf.push(' ');
1963                        }
1964                        heading_buf.push_str(&t);
1965                    } else {
1966                        push_words(&mut current, &t);
1967                    }
1968                }
1969                // Treat hard/soft breaks and rules as whitespace (words already split).
1970                Event::SoftBreak | Event::HardBreak | Event::Rule => {}
1971                _ => {}
1972            }
1973        }
1974        flush(&mut sections, &mut current);
1975        (sections, heading_stack)
1976    }
1977
1978    /// docling-rag's Markdown **window chunker**: the document is split into
1979    /// heading-bounded [`Section`]s of plain words (markup stripped), and a
1980    /// fixed-size window of [`Self::max_words`] words slides over each section
1981    /// with [`Self::overlap`] fractional overlap. A chunk never crosses a
1982    /// heading boundary; [`Self::contextualize`] prefixes the heading path.
1983    #[derive(Debug, Clone)]
1984    pub struct WindowChunker {
1985        /// Window size in words (docling-rag's default 300).
1986        pub max_words: usize,
1987        /// Fractional overlap between consecutive windows (default 0.05 = 5%).
1988        pub overlap: f32,
1989    }
1990
1991    impl Default for WindowChunker {
1992        fn default() -> Self {
1993            WindowChunker {
1994                max_words: 300,
1995                overlap: 0.05,
1996            }
1997        }
1998    }
1999
2000    impl WindowChunker {
2001        pub fn new(max_words: usize, overlap: f32) -> Self {
2002            WindowChunker { max_words, overlap }
2003        }
2004
2005        /// The window size, kept ≥ 1.
2006        fn word_budget(&self) -> usize {
2007            self.max_words.max(1)
2008        }
2009
2010        /// Number of words carried from one window into the next; capped so
2011        /// the window always advances.
2012        fn overlap_words(&self, budget: usize) -> usize {
2013            let o = (budget as f32 * self.overlap).round() as usize;
2014            o.min(budget.saturating_sub(1))
2015        }
2016
2017        /// Chunk a Markdown document.
2018        pub fn chunk(&self, markdown: &str) -> Vec<DocChunk> {
2019            let mut chunks = Vec::new();
2020            self.chunk_with(markdown, &mut |c| {
2021                chunks.push(c);
2022                true
2023            });
2024            chunks
2025        }
2026
2027        /// Stream the chunks: `sink` receives each window as it is cut, and a
2028        /// `false` return cancels. [`Self::chunk`] is this with a collecting
2029        /// sink — the chunks and their order are identical.
2030        pub fn chunk_with(&self, markdown: &str, sink: &mut dyn FnMut(DocChunk) -> bool) {
2031            let (sections, _) = parse_sections_with_stack(markdown, Vec::new());
2032            for section in &sections {
2033                if !self.pack_section(section, sink) {
2034                    return;
2035                }
2036            }
2037        }
2038
2039        /// Slide the window over one completed section, feeding each chunk to
2040        /// `sink`. Returns `false` once the sink cancels — exposed so a
2041        /// streaming caller (docling-rag) can pack sections as they complete.
2042        pub fn pack_section(
2043            &self,
2044            section: &Section,
2045            sink: &mut dyn FnMut(DocChunk) -> bool,
2046        ) -> bool {
2047            let words = &section.words;
2048            if words.is_empty() {
2049                return true;
2050            }
2051            let budget = self.word_budget();
2052            let step = budget - self.overlap_words(budget); // ≥ 1 by construction
2053            let mut start = 0;
2054            loop {
2055                let end = (start + budget).min(words.len());
2056                let chunk = DocChunk {
2057                    text: words[start..end].join(" "),
2058                    headings: (!section.heading_path.is_empty())
2059                        .then(|| section.heading_path.clone()),
2060                    doc_items: Vec::new(),
2061                };
2062                if !sink(chunk) {
2063                    return false;
2064                }
2065                if end >= words.len() {
2066                    return true;
2067                }
2068                start += step;
2069            }
2070        }
2071
2072        /// Render a window chunk for embedding — docling-rag's rendering: the
2073        /// heading path as a `# Outer > Inner` context line, a blank line,
2074        /// then the body (just the body above any heading). Note this differs
2075        /// from the docling chunkers' [`contextualize`](super::contextualize),
2076        /// matching docling-rag instead.
2077        pub fn contextualize(chunk: &DocChunk) -> String {
2078            match &chunk.headings {
2079                Some(h) if !h.is_empty() => format!("# {}\n\n{}", h.join(" > "), chunk.text),
2080                _ => chunk.text.clone(),
2081            }
2082        }
2083    }
2084
2085    #[cfg(test)]
2086    mod tests {
2087        use super::*;
2088
2089        #[test]
2090        fn splits_on_headings_and_tracks_path() {
2091            let md = "\
2092intro words
2093# Chapter 1
2094para one
2095## Section 1.1
2096para two
2097# Chapter 2
2098para three";
2099            let secs = parse_sections(md);
2100            // pre-heading intro, Chapter 1, Section 1.1, Chapter 2.
2101            assert_eq!(secs.len(), 4);
2102            assert!(secs[0].heading_path.is_empty());
2103            assert_eq!(secs[1].heading_path, vec!["Chapter 1"]);
2104            assert_eq!(secs[2].heading_path, vec!["Chapter 1", "Section 1.1"]);
2105            // A deeper heading is dropped when we return to H1.
2106            assert_eq!(secs[3].heading_path, vec!["Chapter 2"]);
2107        }
2108
2109        #[test]
2110        fn strips_markup_to_plain_words() {
2111            let md = "# T\n\nSome **bold** and `code` and [a link](http://x).";
2112            let secs = parse_sections(md);
2113            let words = &secs[0].words;
2114            assert!(words.contains(&"bold".to_string()));
2115            assert!(words.contains(&"code".to_string()));
2116            assert!(words.contains(&"link".to_string()));
2117            // No markdown punctuation survives as its own token.
2118            assert!(!words.iter().any(|w| w.contains('*') || w.contains('`')));
2119        }
2120
2121        #[test]
2122        fn windows_overlap_and_never_cross_headings() {
2123            let body: Vec<String> = (0..25).map(|i| format!("w{i}")).collect();
2124            let md = format!("# A\n\n{}\n\n# B\n\nshort tail\n", body.join(" "));
2125            let chunker = WindowChunker::new(10, 0.2); // step 8, overlap 2
2126            let chunks = chunker.chunk(&md);
2127            // Section A: 25 words → windows [0..10), [8..18), [16..25).
2128            let a: Vec<_> = chunks
2129                .iter()
2130                .filter(|c| c.headings.as_deref() == Some(&["A".to_string()][..]))
2131                .collect();
2132            assert_eq!(a.len(), 3);
2133            assert!(a[0].text.starts_with("w0 ") && a[0].text.ends_with(" w9"));
2134            assert!(a[1].text.starts_with("w8 "), "overlap carries 2 words");
2135            assert!(a[2].text.ends_with(" w24"));
2136            // Section B stays its own chunk; nothing crosses the heading.
2137            let b: Vec<_> = chunks
2138                .iter()
2139                .filter(|c| c.headings.as_deref() == Some(&["B".to_string()][..]))
2140                .collect();
2141            assert_eq!(b.len(), 1);
2142            assert_eq!(b[0].text, "short tail");
2143            assert_eq!(WindowChunker::contextualize(b[0]), "# B\n\nshort tail");
2144        }
2145
2146        #[test]
2147        fn sink_false_cancels_the_window_walk() {
2148            let md = format!(
2149                "# A\n\n{}\n",
2150                (0..50)
2151                    .map(|i| format!("w{i}"))
2152                    .collect::<Vec<_>>()
2153                    .join(" ")
2154            );
2155            let chunker = WindowChunker::new(10, 0.0);
2156            let mut n = 0;
2157            chunker.chunk_with(&md, &mut |_| {
2158                n += 1;
2159                false
2160            });
2161            assert_eq!(n, 1);
2162        }
2163    }
2164}
2165
2166#[cfg(feature = "chunking")]
2167pub use window::{parse_sections, parse_sections_with_stack, Section, WindowChunker};
2168
2169#[cfg(test)]
2170mod tests {
2171    use super::*;
2172
2173    /// A whitespace "tokenizer" for algorithm tests.
2174    struct WordTok(usize);
2175    impl ChunkTokenizer for WordTok {
2176        fn count_tokens(&self, text: &str) -> usize {
2177            text.split_whitespace().count()
2178        }
2179        fn max_tokens(&self) -> usize {
2180            self.0
2181        }
2182    }
2183
2184    fn doc_with(nodes: Vec<Node>) -> DoclingDocument {
2185        let mut d = DoclingDocument::new("t");
2186        for n in nodes {
2187            d.push(n);
2188        }
2189        d
2190    }
2191
2192    #[test]
2193    fn hierarchical_headings_and_items() {
2194        let doc = doc_with(vec![
2195            Node::Heading {
2196                level: 1,
2197                text: "Title".into(),
2198            },
2199            Node::Paragraph {
2200                text: "Intro".into(),
2201            },
2202            Node::Heading {
2203                level: 2,
2204                text: "Sec".into(),
2205            },
2206            Node::Paragraph {
2207                text: "Body".into(),
2208            },
2209        ]);
2210        let chunks = HierarchicalChunker.chunk(&doc);
2211        assert_eq!(chunks.len(), 2);
2212        assert_eq!(chunks[0].text, "Intro");
2213        assert_eq!(chunks[0].headings.as_deref(), Some(&["Title".into()][..]));
2214        assert_eq!(chunks[0].doc_items[0].self_ref, "#/texts/1");
2215        assert_eq!(
2216            chunks[1].headings.as_deref(),
2217            Some(&["Title".into(), "Sec".into()][..])
2218        );
2219        assert_eq!(contextualize(&chunks[1]), "Title\nSec\nBody");
2220    }
2221
2222    /// A page footer is a JSON text item (furniture layer) the chunker never
2223    /// chunks; the refs of the items after it still match the JSON.
2224    #[test]
2225    fn page_furniture_keeps_refs_aligned_with_json() {
2226        let doc = doc_with(vec![
2227            Node::Paragraph {
2228                text: "Before".into(),
2229            },
2230            Node::PageFurniture {
2231                footer: true,
2232                location: [0, 490, 100, 512],
2233                text: "1.10.2".into(),
2234            },
2235            Node::Paragraph {
2236                text: "After".into(),
2237            },
2238        ]);
2239        let chunks = HierarchicalChunker.chunk(&doc);
2240        assert_eq!(chunks.len(), 2);
2241        assert_eq!(chunks[1].text, "After");
2242        let json: serde_json::Value = serde_json::from_str(&doc.export_to_json()).unwrap();
2243        assert_eq!(json["texts"][2]["text"], "After");
2244        assert_eq!(chunks[1].doc_items[0].self_ref, "#/texts/2");
2245    }
2246
2247    /// A picture's children are never chunked, but they take `#/texts/N`
2248    /// (and list-group) slots in the JSON — the refs after them must still
2249    /// point at the right items.
2250    #[test]
2251    fn picture_children_are_not_chunked_but_keep_refs_aligned() {
2252        let doc = doc_with(vec![
2253            Node::Picture {
2254                caption: None,
2255                caption_href: None,
2256                image: None,
2257                classification: None,
2258                caption_parent: crate::CaptionParent::Item,
2259                caption_location: None,
2260            },
2261            Node::PictureChildren(vec![
2262                Node::Paragraph {
2263                    text: "axis label".into(),
2264                },
2265                Node::ListItem {
2266                    ordered: false,
2267                    number: 0,
2268                    first_in_list: true,
2269                    text: "callout".into(),
2270                    level: 0,
2271                    marker: None,
2272                    location: None,
2273                    dclx: None,
2274                    href: None,
2275                    layer: None,
2276                },
2277            ]),
2278            Node::Paragraph {
2279                text: "Body".into(),
2280            },
2281        ]);
2282        let chunks = HierarchicalChunker.chunk(&doc);
2283        assert_eq!(chunks.len(), 1);
2284        assert_eq!(chunks[0].text, "Body");
2285        let json = doc.export_to_json_value();
2286        assert_eq!(json["texts"][2]["text"], "Body");
2287        assert_eq!(chunks[0].doc_items[0].self_ref, "#/texts/2");
2288    }
2289
2290    #[test]
2291    fn heading_shadowing_prunes_deeper_levels() {
2292        let doc = doc_with(vec![
2293            Node::Heading {
2294                level: 2,
2295                text: "A".into(),
2296            },
2297            Node::Heading {
2298                level: 3,
2299                text: "A.1".into(),
2300            },
2301            Node::Heading {
2302                level: 2,
2303                text: "B".into(),
2304            },
2305            Node::Paragraph { text: "p".into() },
2306        ]);
2307        let chunks = HierarchicalChunker.chunk(&doc);
2308        assert_eq!(chunks[0].headings.as_deref(), Some(&["B".into()][..]));
2309    }
2310
2311    #[test]
2312    fn triplet_table() {
2313        let t = Table {
2314            rows: vec![
2315                vec!["".into(), "Col1".into()],
2316                vec!["Row1".into(), "v".into()],
2317            ],
2318            ..Default::default()
2319        };
2320        assert_eq!(triplet_table_text(&t), "Row1, Col1 = v");
2321        // Single-column: row 0 is the dataframe header, the first data row
2322        // becomes the column name, the rest the values.
2323        let single = Table {
2324            rows: vec![vec!["H".into()], vec!["a".into()], vec!["b".into()]],
2325            ..Default::default()
2326        };
2327        assert_eq!(triplet_table_text(&single), "a = b");
2328    }
2329
2330    #[test]
2331    fn hybrid_merges_small_peers_and_splits_large() {
2332        let doc = doc_with(vec![
2333            Node::Heading {
2334                level: 2,
2335                text: "S".into(),
2336            },
2337            Node::Paragraph { text: "a b".into() },
2338            Node::Paragraph { text: "c d".into() },
2339        ]);
2340        let chunks = HybridChunker::new(WordTok(16)).chunk(&doc);
2341        assert_eq!(chunks.len(), 1, "peers under one heading merge");
2342        assert_eq!(chunks[0].text, "a b\nc d");
2343
2344        let long = "w ".repeat(40).trim().to_string();
2345        let doc = doc_with(vec![Node::Paragraph { text: long }]);
2346        let chunks = HybridChunker::new(WordTok(16)).chunk(&doc);
2347        assert!(chunks.len() > 1, "oversized paragraph splits");
2348        for c in &chunks {
2349            assert!(WordTok(16).count_tokens(&contextualize(c)) <= 16);
2350        }
2351    }
2352
2353    #[test]
2354    fn semchunk_prefers_newlines_then_sentences() {
2355        let tok = WordTok(4);
2356        let out = semchunk("one two three. four five six\nseven eight", 4, &tok);
2357        assert!(out.iter().all(|c| tok.count_tokens(c) <= 4), "{out:?}");
2358    }
2359}