1use std::collections::BTreeMap;
24
25use crate::document::{DoclingDocument, Node, Table};
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub enum ChunkItemKind {
31 Text,
33 Table,
36 Picture,
38}
39
40#[derive(Debug, Clone, PartialEq)]
43pub struct ChunkItem {
44 pub self_ref: String,
47 pub kind: ChunkItemKind,
48 pub text: String,
52}
53
54#[derive(Debug, Clone, PartialEq)]
56pub struct DocChunk {
57 pub text: String,
60 pub headings: Option<Vec<String>>,
63 pub doc_items: Vec<ChunkItem>,
65}
66
67pub 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#[derive(Debug, Clone, Default)]
86pub struct HierarchicalChunker;
87
88impl HierarchicalChunker {
89 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 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#[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 headings: BTreeMap<u8, String>,
165 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 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 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 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; }
284 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 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 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 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 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 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 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 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 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 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 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 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 Node::Group { layer: Some(_), .. } => {}
589 Node::Group { children, .. } => {
590 self.alloc.group();
593 self.walk(children);
594 }
595 Node::KeyValueGraph { .. } => {}
597 Node::FieldRegion { items } => {
598 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 Node::Located { inner, .. }
635 | Node::Prov { inner, .. }
636 | Node::Track { inner, .. }
637 | Node::Commented { inner, .. } => self.one(inner),
638 Node::PageFurniture { .. } | Node::FurnitureText { .. } => {
642 self.alloc.text();
643 }
644 Node::PictureChildren(children) => {
650 for child in children {
651 match child {
652 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 Node::CommentSection { .. }
668 | Node::Furniture { .. }
669 | Node::PageBreak
670 | Node::PageInfo { .. }
671 | Node::DoclangOnly(_) => {}
672 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
686fn 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
719fn 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
735fn 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
746fn 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 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 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 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 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 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 (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
863fn inline_segments(md: &str) -> Vec<String> {
871 inline_segments_tagged(md)
872 .into_iter()
873 .map(|(t, _)| t)
874 .collect()
875}
876
877fn 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 if chars[i] == '[' && !rest.starts_with("[](") {
920 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 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 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 if rest.starts_with("$$") {
998 plain.push_str("$$");
999 i += 2;
1000 continue;
1001 }
1002 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
1025fn 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 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
1086fn 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
1121fn 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
1167fn unescape_text(s: &str) -> String {
1170 s.replace("<", "<")
1171 .replace(">", ">")
1172 .replace("&", "&")
1173 .replace("\\_", "_")
1174}
1175
1176pub trait ChunkTokenizer {
1182 fn count_tokens(&self, text: &str) -> usize;
1184 fn max_tokens(&self) -> usize;
1186}
1187
1188pub 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 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 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 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 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; 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; }
1296 } else if window_start == window_end {
1297 window_end += 1;
1300 window_start = window_end;
1301 } else {
1302 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 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 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 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 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#[derive(Default)]
1421struct PeerMerger {
1422 window: Vec<DocChunk>,
1423 merged: Option<DocChunk>,
1424}
1425
1426fn 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 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
1480fn 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 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
1523const NON_WS_SPLITTERS: &[&str] = &[
1529 ".", "?", "!", "*", ";", ",", "(", ")", "[", "]", "\u{201c}", "\u{201d}", "\u{2018}",
1530 "\u{2019}", "'", "\"", "`", ":", "\u{2014}", "\u{2026}", "/", "\\", "\u{2013}", "&", "-",
1531];
1532
1533pub 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 chunks
1549 .into_iter()
1550 .filter(|c| !c.is_empty() && !c.chars().all(char::is_whitespace))
1551 .collect()
1552}
1553
1554fn 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 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
1620fn 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
1669fn split_text(text: &str) -> (String, bool, Vec<String>) {
1671 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 if text.contains('\t') {
1678 let splitter = longest_run(text, |c| c == '\t');
1679 return (splitter.clone(), true, split_on(text, &splitter));
1680 }
1681 if text.chars().any(char::is_whitespace) {
1683 let splitter = longest_run(text, char::is_whitespace);
1684 if splitter.chars().count() == 1 {
1685 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 for s in NON_WS_SPLITTERS {
1696 if text.contains(s) {
1697 return (s.to_string(), false, split_on(text, s));
1698 }
1699 }
1700 (
1702 String::new(),
1703 true,
1704 text.chars().map(|c| c.to_string()).collect(),
1705 )
1706}
1707
1708fn 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
1732fn 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#[cfg(feature = "chunking")]
1767mod hf {
1768 use super::ChunkTokenizer;
1769
1770 pub const DEFAULT_TOKENIZER_PATH: &str = ".models/chunk/tokenizer.json";
1775
1776 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 pub struct HuggingFaceTokenizer {
1799 tok: tokenizers::Tokenizer,
1800 max_tokens: usize,
1801 }
1802
1803 impl HuggingFaceTokenizer {
1804 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 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 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#[cfg(feature = "chunking")]
1851mod window {
1852 use super::DocChunk;
1853 use pulldown_cmark::{Event, HeadingLevel, Parser, Tag, TagEnd};
1854
1855 #[derive(Debug, Clone, Default)]
1857 pub struct Section {
1858 pub heading_path: Vec<String>,
1861 pub words: Vec<String>,
1863 }
1864
1865 impl Section {
1866 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 pub fn parse_sections(markdown: &str) -> Vec<Section> {
1892 parse_sections_with_stack(markdown, Vec::new()).0
1893 }
1894
1895 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 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 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 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 Event::SoftBreak | Event::HardBreak | Event::Rule => {}
1971 _ => {}
1972 }
1973 }
1974 flush(&mut sections, &mut current);
1975 (sections, heading_stack)
1976 }
1977
1978 #[derive(Debug, Clone)]
1984 pub struct WindowChunker {
1985 pub max_words: usize,
1987 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 fn word_budget(&self) -> usize {
2007 self.max_words.max(1)
2008 }
2009
2010 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 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 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 §ions {
2033 if !self.pack_section(section, sink) {
2034 return;
2035 }
2036 }
2037 }
2038
2039 pub fn pack_section(
2043 &self,
2044 section: &Section,
2045 sink: &mut dyn FnMut(DocChunk) -> bool,
2046 ) -> bool {
2047 let words = §ion.words;
2048 if words.is_empty() {
2049 return true;
2050 }
2051 let budget = self.word_budget();
2052 let step = budget - self.overlap_words(budget); 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 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 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 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 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); let chunks = chunker.chunk(&md);
2127 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 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 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 #[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 #[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 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}