diff --git a/crates/chunkmd/src/lib.rs b/crates/chunkmd/src/lib.rs index 8b1403d..86b0528 100644 --- a/crates/chunkmd/src/lib.rs +++ b/crates/chunkmd/src/lib.rs @@ -33,7 +33,7 @@ pub fn chunk(markdown: &str, budgets: &Budgets) -> Vec { let mut chunks: Vec = Vec::new(); for page in pages { for packed in pack::page(page.text, budgets) { - for (piece, tokens) in pack::enforce_hard_max(&packed, budgets.hard_max) { + for (piece, tokens) in pack::enforce_hard_max(&packed.text, budgets.hard_max) { if tokens < budgets.min_tokens { continue; } @@ -42,11 +42,15 @@ pub fn chunk(markdown: &str, budgets: &Budgets) -> Vec { // that have headings but no page anchors. Page zero counts as // no page at all, which is what the Python's truth test makes // of it. + // The passage's own position, not the page's: an + // unpaginated document is a single page beginning at zero, + // so asking at the page put every chunk in the book under + // whichever heading came first in the file. let chapter = page .number .filter(|number| *number != 0) .and_then(|number| pages::chapter_of_page(&chapters, number)) - .or_else(|| by_offset.at(page.offset)); + .or_else(|| by_offset.at(page.offset + packed.offset)); let idx = chunks.len() as u32; let notes = footnote::looks_like_citations(&piece); @@ -110,6 +114,30 @@ mod tests { assert_eq!(chunks[0].chapter.as_deref(), Some("The Second Estate")); } + #[test] + fn every_chunk_answers_with_the_heading_it_sits_under() { + // An unpaginated document is one page beginning at zero, so asking + // the page where a chunk was rather than asking the chunk put the + // whole book under whichever heading came first in the file. + let markdown = "# Chapter One\n\nBody text long enough to be worth \ + keeping as a passage of its own.\n\n# Bibliography\n\n\ + Bloch, Marc. *Feudal Society*. London: Routledge, 1961.\n\n\ + Duby, Georges. *The Three Orders*. Chicago: Chicago, 1980.\n"; + let budgets = Budgets { + target: 40, + overlap: 0, + min_tokens: 5, + ..Budgets::default() + }; + let chunks = chunk(markdown, &budgets); + let under: Vec<&str> = chunks + .iter() + .map(|chunk| chunk.chapter.as_deref().unwrap_or("—")) + .collect(); + assert!(under.contains(&"Chapter One"), "{under:?}"); + assert!(under.contains(&"Bibliography"), "{under:?}"); + } + #[test] fn a_chapter_is_found_by_offset_when_there_are_no_page_markers() { let markdown = "# The Second Estate\n\nalpha beta gamma"; diff --git a/crates/chunkmd/src/pack.rs b/crates/chunkmd/src/pack.rs index 3a44918..c90791c 100644 --- a/crates/chunkmd/src/pack.rs +++ b/crates/chunkmd/src/pack.rs @@ -23,6 +23,20 @@ pub struct Budgets { pub hard_max: usize, } +/// A passage as packed, and where in the page it began. +/// +/// The offset is what lets a chunk be placed back in the document it came +/// from — which is how it finds the heading it sits under. Without it every +/// chunk of an unpaginated document answers with the first heading in the +/// file. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Packed { + /// The passage. + pub text: String, + /// Byte offset within the page's text where the passage begins. + pub offset: usize, +} + impl Default for Budgets { fn default() -> Self { Self { @@ -35,14 +49,17 @@ impl Default for Budgets { } /// Pack one page's text into chunk-sized pieces. -pub fn page(text: &str, budgets: &Budgets) -> Vec { +pub fn page(text: &str, budgets: &Budgets) -> Vec { let paragraphs = textkit::paragraphs(text); if paragraphs.is_empty() { return Vec::new(); } + // Where a piece began, so that what is packed can still be placed. The + // pieces are slices of `text`, so the arithmetic is exact. + let at = |piece: &str| (piece.as_ptr() as usize).saturating_sub(text.as_ptr() as usize); - let mut chunks = Vec::new(); - let mut buf: Vec = Vec::new(); + let mut chunks: Vec = Vec::new(); + let mut buf: Vec<(String, usize)> = Vec::new(); let mut buf_tokens = 0usize; for paragraph in paragraphs { @@ -56,41 +73,53 @@ pub fn page(text: &str, budgets: &Budgets) -> Vec { for sentence in textkit::sentences(paragraph) { let sentence_tokens = textkit::count(sentence); if buf_tokens + sentence_tokens > budgets.target && !buf.is_empty() { - chunks.push(buf.join(" ")); + chunks.push(joined(&buf, " ")); buf = carry_tail(&buf, budgets.overlap); - buf_tokens = textkit::count(&buf.join(" ")); + buf_tokens = textkit::count(&joined(&buf, " ").text); } - buf.push(sentence.to_string()); + buf.push((sentence.to_string(), at(sentence))); buf_tokens += sentence_tokens; } continue; } if buf_tokens + paragraph_tokens > budgets.target && !buf.is_empty() { - chunks.push(buf.join("\n\n")); + chunks.push(joined(&buf, "\n\n")); buf = carry_tail(&buf, budgets.overlap); - buf_tokens = textkit::count(&buf.join("\n\n")); + buf_tokens = textkit::count(&joined(&buf, "\n\n").text); } - buf.push(paragraph.to_string()); + buf.push((paragraph.to_string(), at(paragraph))); buf_tokens += paragraph_tokens; } if !buf.is_empty() { - chunks.push(buf.join("\n\n")); + chunks.push(joined(&buf, "\n\n")); } chunks } +/// Join a buffer's pieces, keeping where the first of them began. +fn joined(buf: &[(String, usize)], separator: &str) -> Packed { + Packed { + text: buf + .iter() + .map(|(piece, _)| piece.as_str()) + .collect::>() + .join(separator), + offset: buf.first().map(|(_, at)| *at).unwrap_or_default(), + } +} + /// Keep the last pieces of a buffer, up to `overlap` tokens of context. /// /// Always keeps at least one piece, even when that piece alone is over /// budget: an overlap of nothing is a chunk boundary with no bridge across /// it, which is the thing overlap exists to prevent. -fn carry_tail(buf: &[String], overlap: usize) -> Vec { - let mut out: Vec = Vec::new(); +fn carry_tail(buf: &[(String, usize)], overlap: usize) -> Vec<(String, usize)> { + let mut out: Vec<(String, usize)> = Vec::new(); let mut total = 0usize; for piece in buf.iter().rev() { - let tokens = textkit::count(piece); + let tokens = textkit::count(&piece.0); if total + tokens > overlap && !out.is_empty() { break; } @@ -125,6 +154,10 @@ pub fn enforce_hard_max(text: &str, hard_max: usize) -> Vec<(String, usize)> { mod tests { use super::*; + fn texts(packed: Vec) -> Vec { + packed.into_iter().map(|one| one.text).collect() + } + fn budgets(target: usize, overlap: usize) -> Budgets { Budgets { target, @@ -135,20 +168,20 @@ mod tests { #[test] fn short_text_is_one_chunk() { - let chunks = page("one\n\ntwo", &Budgets::default()); + let chunks = texts(page("one\n\ntwo", &Budgets::default())); assert_eq!(chunks, ["one\n\ntwo"]); } #[test] fn paragraphs_are_rejoined_with_a_blank_line() { - let chunks = page("alpha\n\n\n\nbeta", &Budgets::default()); + let chunks = texts(page("alpha\n\n\n\nbeta", &Budgets::default())); assert_eq!(chunks, ["alpha\n\nbeta"]); } #[test] fn a_full_buffer_flushes_and_carries_a_tail() { let text = "one two three\n\nfour five six\n\nseven eight nine"; - let chunks = page(text, &budgets(6, 4)); + let chunks = texts(page(text, &budgets(6, 4))); assert!(chunks.len() > 1, "{chunks:?}"); // Every chunk after the first repeats the end of the one before it. for pair in chunks.windows(2) { @@ -165,7 +198,7 @@ mod tests { #[test] fn an_oversized_paragraph_is_cut_at_sentences() { let paragraph = "Alpha one. Beta two. Gamma three. Delta four. Epsilon five."; - let chunks = page( + let chunks = texts(page( paragraph, &Budgets { target: 6, @@ -173,7 +206,7 @@ mod tests { min_tokens: 1, hard_max: 8, }, - ); + )); assert!(chunks.len() > 1, "{chunks:?}"); assert!(chunks[0].contains(' ')); } @@ -185,7 +218,7 @@ mod tests { #[test] fn the_tail_keeps_one_piece_even_when_it_is_over_budget() { - let buf = vec!["a longer paragraph than the overlap allows".to_string()]; + let buf = vec![("a longer paragraph than the overlap allows".to_string(), 0)]; assert_eq!(carry_tail(&buf, 1).len(), 1); } diff --git a/crates/chunkmd/src/pages.rs b/crates/chunkmd/src/pages.rs index 290ee31..ff315c6 100644 --- a/crates/chunkmd/src/pages.rs +++ b/crates/chunkmd/src/pages.rs @@ -25,8 +25,10 @@ pub struct Page<'a> { pub number: Option, /// The page's text, trimmed. pub text: &'a str, - /// Byte offset in the source markdown where the page's body starts, used - /// to find the chapter when there are no page markers to go on. + /// Byte offset in the source markdown where `text` begins, used to find + /// the chapter a passage sits under. Exact rather than approximate: it is + /// added to a passage's own offset within the page, and a heading missed + /// by a few bytes is the wrong chapter for everything after it. pub offset: usize, } @@ -51,6 +53,12 @@ pub struct Chapter { /// Nothing surfaced the loss: the book stayed ingested, with a single chunk. /// So the prologue is emitted with no page number, since we genuinely do not /// know its page, and the offset-based chapter lookup still tags it. +/// Where `part` begins inside `whole`, which it must be a slice of. +fn offset_of(whole: &str, part: &str) -> usize { + (part.as_ptr() as usize).saturating_sub(whole.as_ptr() as usize) +} + +/// Cut `markdown` at its page markers. pub fn split_by_page(markdown: &str) -> Vec> { let markers: Vec<_> = PAGE.find_iter(markdown).collect(); if markers.is_empty() { @@ -67,7 +75,7 @@ pub fn split_by_page(markdown: &str) -> Vec> { pages.push(Page { number: None, text: head, - offset: 0, + offset: offset_of(markdown, head), }); } for (at, marker) in markers.iter().enumerate() { @@ -80,7 +88,7 @@ pub fn split_by_page(markdown: &str) -> Vec> { pages.push(Page { number: Some(textkit::decimal(&markdown[marker.start()..marker.end()])), text, - offset: start, + offset: offset_of(markdown, text), }); } }