funkpower/catena
1
0
Fork 0
mirror of https://git.pumpin.online/funkpower/catena.git synced 2026-10-22 18:57:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

chunkmd: page and chapter aware chunking, roles, footnote flag

Browse Source
This commit is contained in:
funkpower
2026-09-19 11:22:56 -04:00
parent acf8605d37
commit e0473fb8ae
45 changed files with 2054 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/chunkmd/Cargo.toml
+18
View File
@@ -0,0 +1,18 @@
[package]
name = "chunkmd"
description = "Cuts markdown into retrievable passages that carry their page, chapter and role."
version.workspace = true
edition.workspace = true
rust-version.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
clap.workspace = true
corpus-core.workspace = true
fancy-regex.workspace = true
regex.workspace = true
textkit.workspace = true
[lints]
workspace = true
crates/chunkmd/src/footnote.rs
+179
View File
@@ -0,0 +1,179 @@
//! Spotting a chunk that is really a page of numbered notes.
//!
//! Footnotes and endnotes are citations in disguise: they read as body text,
//! rank as body text, and answer a query with a publisher and a page number.
//! Flagging them lets a consumer treat them as the bibliography they are
//! without moving them out of the body, where their page provenance lives.
use std::sync::LazyLock;
use fancy_regex::Regex as FancyRegex;
use regex::Regex;
/// A private-use codepoint marking the start of a note's text.
///
/// EPUB3 note definitions carry a structural signal — `epub:type="footnote"`
/// and the ARIA `doc-footnote` role — that is reliable in a way the visible
/// marker never is: `[65]`, `fn39.`, a bare superscript, all of them and
/// none of them. A converter that can see that structure marks it with this
/// character, which survives conversion to markdown verbatim. When it is
/// present, nothing below has to guess.
///
/// It is stripped from the text a chunk carries, after the flag is set.
pub const SENTINEL: char = '\u{e000}';
/// How far past its marker a note's citation signal may sit, in characters.
///
/// Citation signals sit near the front of an entry. The cap stops a marker
/// followed by pages of prose from matching a signal far away. HTML class
/// soup inflates entries, so it is generous.
const ENTRY_WINDOW: usize = 500;
/// At least this many numbered-note lines have to carry a citation signal.
///
/// Two, because one avoids a false positive from a stray "1. " list item,
/// and a real page of notes almost always has more.
const MIN_CITED_NOTES: usize = 2;
/// The start of a numbered note, in any of the shapes conversion leaves
/// behind. A marker only counts at a line start or after a run of two or
/// more spaces — several notes can share one physical line — so that a
/// superscript *reference* to a note, mid-sentence, does not qualify.
static NOTE_MARKER: LazyLock<FancyRegex> = LazyLock::new(|| {
FancyRegex::new(
r"(?xm)
(?: ^[ \t]{0,8} (?:-[ \t]+)? # line start, optional '- ' list bullet
| (?<=\s\s) # or after a 2+ space run
| (?<=\ -\ ) # or after ' - ' (inline list separator)
)
(?:
\d{1,3} \\?[.)] [ \t]+ \S # plain '12. ' or escaped '12\. '
| \[ \d{1,3} \] [ \t]* \S # bracket-numbered '[12] Author'
| \d{1,3} [ \t]+ \S # bare '12 ' from a page bottom
| [¹²³⁰-⁹]+ \s* \S # unicode superscript marker
| (?:fn|ftn|note) \d{1,4} [.)] [ \t]+ \S # id-style 'fn39. '
| (?:<span[^>]*>\s*)* # outer span wrap
<a\b[^>]*> \s* (?:<span[^>]*>\s*)* (?:<sup>\s*)?
\d{1,3} \s* [.)]?
\s* (?:</sup>\s*)? (?:</span>\s*)* </a>
(?:\s*</span>)* \s* [.)]? # anchor forms
| \[ \s* (?:<span[^>]*>)? \s* \d{1,3}
\s* (?:</span>)? \s* \] \( \#[^)]* \) # [42](#...)
| <sup> \s* \d{1,3} \s* </sup> # bare sup note
)",
)
.expect("a valid pattern")
});
/// What makes a note bibliographic rather than discursive.
static CITE_SIGNAL: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(
r"(?ix)
\((?:[^()]*\b)?(?:1[5-9]|20)\d{2}\) # parenthesised year: (Oxford, 1998)
| \bpp?\.\s*\d # p. 12 / pp. 12-34
| \bvol\.\s*\d # vol. 2
| \bIbid\b | \bop\.\s*cit\b | \bloc\.\s*cit\b
| \bet\s+al\b
| \b(?:University|Clarendon|Academic)\s+Press\b | \bPress,\s
| \beds?\.\s | \btrans\.\s
| \*[^*\n]{3,80}\*\s*,?\s*\d # short form: *Title*, 183
| <span[^>]*italic[^>]*> # HTML-italic title survives conversion
| ,\s+\d{1,4}\s*[.;] # trailing page: '..., 183.'
| https?:// # web citations (URLs, DOIs)",
)
.expect("a valid pattern")
});
/// Whether a chunk carries footnote or endnote citation material.
///
/// The structural signal first, when the source had one. Otherwise each note
/// entry runs from its marker to the next one, capped at a few hundred
/// characters, and enough of them have to look like citations.
pub fn looks_like_citations(text: &str) -> bool {
if text.contains(SENTINEL) {
return true;
}
let starts: Vec<usize> = NOTE_MARKER
.find_iter(text)
.filter_map(Result::ok)
.map(|found| found.start())
.collect();
if starts.len() < MIN_CITED_NOTES {
return false;
}
let mut cited = 0;
for (at, start) in starts.iter().enumerate() {
let end = starts.get(at + 1).copied().unwrap_or(text.len());
let entry: String = text[*start..end].chars().take(ENTRY_WINDOW).collect();
if CITE_SIGNAL.is_match(&entry) {
cited += 1;
if cited >= MIN_CITED_NOTES {
return true;
}
}
}
false
}
/// Remove the structural marker from text on its way out.
pub fn strip_sentinel(text: &str) -> String {
text.replace(SENTINEL, "")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn the_structural_marker_settles_it_on_its_own() {
assert!(looks_like_citations(&format!("{SENTINEL}See above.")));
}
#[test]
fn two_numbered_notes_with_citation_signals_count() {
let text = "12. Arendt, The Human Condition (Chicago, 1958), 44.\n\
13. Bloch, Feudal Society, vol. 2, p. 17.";
assert!(looks_like_citations(text));
}
#[test]
fn one_note_is_not_enough() {
assert!(!looks_like_citations(
"12. Arendt, The Human Condition (Chicago, 1958), 44."
));
}
#[test]
fn a_numbered_list_without_citations_is_not_notes() {
let text = "1. Boil the water.\n2. Add the salt.\n3. Wait.";
assert!(!looks_like_citations(text));
}
#[test]
fn a_mid_sentence_reference_is_not_a_note() {
let text = "As shown elsewhere 12. and again 13. the argument holds \
(Oxford, 1998) and p. 12.";
assert!(!looks_like_citations(text));
}
#[test]
fn a_signal_far_past_its_marker_does_not_count() {
let filler = "word ".repeat(200);
let text = format!("12. {filler}(Oxford, 1998)\n13. {filler}p. 44");
assert!(!looks_like_citations(&text));
}
#[test]
fn anchor_wrapped_markers_are_recognised() {
let text = "<a href=\"#ref1\"><sup>1</sup></a>. Arendt, (Chicago, 1958), 44.\n\
<a href=\"#ref2\"><sup>2</sup></a>. Bloch, p. 17.";
assert!(looks_like_citations(text));
}
#[test]
fn the_marker_is_removed_from_the_text_that_ships() {
assert_eq!(strip_sentinel(&format!("{SENTINEL}note")), "note");
}
}
crates/chunkmd/src/lib.rs
+187
View File
@@ -0,0 +1,187 @@
//! Cuts markdown into retrievable passages that carry their page, chapter
//! and role.
//!
//! The order of the passes is load-bearing and is not obvious:
//!
//! 1. Split on page markers, so a chunk never spans a page and its page
//! number is exact rather than inferred.
//! 2. Pack paragraphs toward a token target within each page, carrying an
//! overlap forward.
//! 3. Enforce the hard maximum on token boundaries, whatever the structure
//! said.
//! 4. Drop what is too small to be worth an embedding.
//! 5. Flag the note chunks, *before* the structural marker is stripped out
//! of their text.
//! 6. Classify roles across the whole document at once, because smoothing a
//! stray chunk needs to see its neighbours.
pub mod footnote;
pub mod pack;
pub mod pages;
pub mod role;
use corpus_core::{Chunk, chunk::flags};
pub use pack::Budgets;
/// Cut `markdown` into chunks.
pub fn chunk(markdown: &str, budgets: &Budgets) -> Vec<Chunk> {
let pages = pages::split_by_page(markdown);
let chapters = pages::chapters(markdown);
let by_page = pages::page_chapters(&chapters);
let by_offset = pages::ChapterIndex::new(&chapters);
let mut chunks: Vec<Chunk> = Vec::new();
for page in pages {
for packed in pack::page(page.text, budgets) {
for piece in pack::enforce_hard_max(&packed, budgets.hard_max) {
let tokens = textkit::count(&piece);
if tokens < budgets.min_tokens {
continue;
}
// A page-based chapter when the document has page markers;
// the byte offset otherwise, which is what carries EPUBs
// that have headings but no page anchors.
let chapter = page
.number
.and_then(|number| by_page.get(&number).copied())
.or_else(|| by_offset.at(page.offset));
let idx = chunks.len() as u32;
let notes = footnote::looks_like_citations(&piece);
let mut chunk = Chunk::new(idx, footnote::strip_sentinel(&piece), tokens as u32);
chunk.page = page.number;
chunk.chapter = chapter.map(str::to_string);
if notes {
chunk.flag(flags::FOOTNOTE_CITATIONS);
}
chunks.push(chunk);
}
}
}
let mut roles: Vec<_> = chunks
.iter()
.map(|chunk| role::classify(&chunk.text, chunk.chapter.as_deref()))
.collect();
role::smooth(&mut roles);
for (chunk, role) in chunks.iter_mut().zip(roles) {
chunk.role = role;
}
chunks
}
#[cfg(test)]
mod tests {
use super::*;
use corpus_core::Role;
fn small() -> Budgets {
Budgets {
target: 20,
overlap: 4,
min_tokens: 1,
hard_max: 40,
}
}
#[test]
fn chunks_are_numbered_from_zero_without_gaps() {
let markdown = "# One\n\nalpha beta\n\n<!-- Page 2 -->\n\ngamma delta\n\nepsilon zeta";
let chunks = chunk(markdown, &small());
for (at, chunk) in chunks.iter().enumerate() {
assert_eq!(chunk.idx, at as u32);
}
}
#[test]
fn a_chunk_carries_the_page_it_came_from() {
let markdown = "<!-- Page 7 -->\n\nalpha beta gamma";
let chunks = chunk(markdown, &small());
assert_eq!(chunks[0].page, Some(7));
}
#[test]
fn a_chunk_carries_its_chapter() {
let markdown = "<!-- Page 1 -->\n\n# The Second Estate\n\nalpha beta gamma";
let chunks = chunk(markdown, &small());
assert_eq!(chunks[0].chapter.as_deref(), Some("The Second Estate"));
}
#[test]
fn a_chapter_is_found_by_offset_when_there_are_no_page_markers() {
let markdown = "# The Second Estate\n\nalpha beta gamma";
let chunks = chunk(markdown, &small());
assert_eq!(chunks[0].page, None);
assert_eq!(chunks[0].chapter.as_deref(), Some("The Second Estate"));
}
#[test]
fn a_chunk_below_the_minimum_is_dropped() {
let budgets = Budgets {
min_tokens: 100,
..small()
};
assert!(chunk("alpha beta", &budgets).is_empty());
}
#[test]
fn nothing_exceeds_the_hard_maximum() {
let markdown = "word ".repeat(500);
let budgets = Budgets {
target: 30,
overlap: 5,
min_tokens: 1,
hard_max: 35,
};
for chunk in chunk(&markdown, &budgets) {
assert!(chunk.tokens <= 35, "{} tokens", chunk.tokens);
}
}
#[test]
fn a_note_chunk_is_flagged_and_its_marker_is_stripped() {
let markdown = format!(
"{}12. Arendt, The Human Condition (Chicago, 1958), 44.\n\n\
13. Bloch, Feudal Society, vol. 2, p. 17.",
footnote::SENTINEL
);
let chunks = chunk(&markdown, &small());
assert!(chunks[0].has_flag(flags::FOOTNOTE_CITATIONS));
assert!(!chunks[0].text.contains(footnote::SENTINEL));
}
#[test]
fn roles_are_smoothed_across_the_whole_document() {
// Two bibliography-shaped pages with an ambiguous one between them.
let entries = "Arendt, Hannah. The Human Condition. Chicago: University Press, 1958.\n\n\
Bloch, Marc. Feudal Society. London: Routledge, 1961.\n\n\
Cassirer, Ernst. The Myth of the State. New Haven: Yale Press, 1946.";
let markdown = format!(
"<!-- Page 1 -->\n\n{entries}\n\n<!-- Page 2 -->\n\nhttps://example.org/a\n\n\
<!-- Page 3 -->\n\n{entries}"
);
let chunks = chunk(
&markdown,
&Budgets {
target: 200,
overlap: 0,
min_tokens: 1,
hard_max: 400,
},
);
assert!(chunks.len() >= 3, "{chunks:#?}");
assert!(
chunks.iter().all(|chunk| chunk.role == Role::Bibliography),
"{:?}",
chunks.iter().map(|c| c.role).collect::<Vec<_>>()
);
}
#[test]
fn empty_input_produces_no_chunks() {
assert!(chunk("", &Budgets::default()).is_empty());
assert!(chunk(" \n\n ", &Budgets::default()).is_empty());
}
}
crates/chunkmd/src/main.rs
+130
View File
@@ -0,0 +1,130 @@
//! Cuts markdown into retrievable passages that carry their page, chapter
//! and role.
use std::io::{BufWriter, Read, Write};
use std::path::PathBuf;
use std::process::ExitCode;
use chunkmd::Budgets;
use clap::Parser;
use corpus_core::{exit, ndjson};
/// Cut markdown into chunks, one JSON record per line.
///
/// Reads markdown on stdin, or from a path. Page markers (`<!-- Page 41 -->`)
/// and headings become the page and chapter on every chunk, so a passage can
/// always be cited back to where it came from.
#[derive(Parser)]
#[command(
version,
about,
long_about = None,
after_help = "\
Example:
mdscrub book.md | chunkmd --target 1024 --overlap 128 > chunks.ndjson
The defaults suit a model with a 512-token window; a larger model wants
--target 1024 --overlap 128 --min-tokens 80 --hard-max 1280."
)]
struct Args {
/// Markdown file to read. Reads stdin when absent.
path: Option<PathBuf>,
/// Flush a chunk at this many tokens.
#[arg(long, default_value_t = Budgets::default().target, value_name = "TOKENS")]
target: usize,
/// Carry this much of each chunk into the next one.
#[arg(long, default_value_t = Budgets::default().overlap, value_name = "TOKENS")]
overlap: usize,
/// Drop chunks smaller than this.
#[arg(long, default_value_t = Budgets::default().min_tokens, value_name = "TOKENS")]
min_tokens: usize,
/// Never emit a chunk larger than this.
#[arg(long, default_value_t = Budgets::default().hard_max, value_name = "TOKENS")]
hard_max: usize,
}
fn main() -> ExitCode {
let args = Args::parse();
let budgets = Budgets {
target: args.target,
overlap: args.overlap,
min_tokens: args.min_tokens,
hard_max: args.hard_max,
};
if let Err(why) = sane(&budgets) {
eprintln!("chunkmd: {why}");
return ExitCode::from(exit::USAGE);
}
let markdown = match read(args.path.as_deref()) {
Ok(markdown) => markdown,
Err(e) => {
eprintln!("chunkmd: {e}");
return ExitCode::from(exit::IO);
}
};
let mut out = BufWriter::new(std::io::stdout().lock());
for chunk in chunkmd::chunk(&markdown, &budgets) {
if let Err(e) = ndjson::write_line(&mut out, &chunk) {
// A consumer that stopped reading is not a failure: `| head` is a
// normal way to look at a stream.
if e.is_broken_pipe() {
return ExitCode::from(exit::OK);
}
eprintln!("chunkmd: {e}");
return ExitCode::from(exit::IO);
}
}
match out.flush() {
Ok(()) => ExitCode::from(exit::OK),
Err(e) if e.kind() == std::io::ErrorKind::BrokenPipe => ExitCode::from(exit::OK),
Err(e) => {
eprintln!("chunkmd: {e}");
ExitCode::from(exit::IO)
}
}
}
/// Budgets that cannot produce anything sensible, rejected before reading
/// input rather than after.
fn sane(budgets: &Budgets) -> Result<(), String> {
if budgets.target == 0 || budgets.hard_max == 0 {
return Err("--target and --hard-max must be above zero".into());
}
if budgets.overlap >= budgets.target {
return Err(format!(
"--overlap {} must be below --target {}, or a chunk is all overlap",
budgets.overlap, budgets.target
));
}
if budgets.hard_max < budgets.target {
return Err(format!(
"--hard-max {} must be at least --target {}",
budgets.hard_max, budgets.target
));
}
if budgets.min_tokens > budgets.hard_max {
return Err(format!(
"--min-tokens {} is above --hard-max {}, so every chunk would be dropped",
budgets.min_tokens, budgets.hard_max
));
}
Ok(())
}
fn read(path: Option<&std::path::Path>) -> std::io::Result<String> {
match path {
Some(path) => std::fs::read_to_string(path),
None => {
let mut markdown = String::new();
std::io::stdin().read_to_string(&mut markdown)?;
Ok(markdown)
}
}
}
crates/chunkmd/src/pack.rs
+186
View File
@@ -0,0 +1,186 @@
//! Packing paragraphs into chunks of about the size an embedding model wants.
//!
//! Page first, paragraph second, sentence only when a paragraph is itself too
//! big. Every flush carries a tail of the previous chunk forward, so a
//! sentence that answers a question is not orphaned from the sentence that
//! asked it.
/// The token budgets a chunker works to.
///
/// The defaults suit a model with a 512-token comfortable window. A model
/// with a larger one wants roughly double everything — the numbers below are
/// ratios that have held up, not constants with a life of their own.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Budgets {
/// Where a chunk is flushed, in tokens.
pub target: usize,
/// How much of the previous chunk is carried into the next one.
pub overlap: usize,
/// Chunks smaller than this are dropped: a fragment retrieves badly and
/// costs an embedding.
pub min_tokens: usize,
/// No chunk may exceed this, whatever structure says.
pub hard_max: usize,
}
impl Default for Budgets {
fn default() -> Self {
Self {
target: 512,
overlap: 64,
min_tokens: 50,
hard_max: 768,
}
}
}
/// Pack one page's text into chunk-sized pieces.
pub fn page(text: &str, budgets: &Budgets) -> Vec<String> {
let paragraphs = textkit::paragraphs(text);
if paragraphs.is_empty() {
return Vec::new();
}
let mut chunks = Vec::new();
let mut buf: Vec<String> = Vec::new();
let mut buf_tokens = 0usize;
for paragraph in paragraphs {
let paragraph_tokens = textkit::count(paragraph);
if paragraph_tokens > budgets.hard_max {
// A paragraph nothing can hold: break it at sentences instead.
// Note that a flush from here joins with a space, while the final
// flush below joins with a blank line — the seam a chunk was cut
// at is visible in how its pieces are rejoined.
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(" "));
buf = carry_tail(&buf, budgets.overlap);
buf_tokens = textkit::count(&buf.join(" "));
}
buf.push(sentence.to_string());
buf_tokens += sentence_tokens;
}
continue;
}
if buf_tokens + paragraph_tokens > budgets.target && !buf.is_empty() {
chunks.push(buf.join("\n\n"));
buf = carry_tail(&buf, budgets.overlap);
buf_tokens = textkit::count(&buf.join("\n\n"));
}
buf.push(paragraph.to_string());
buf_tokens += paragraph_tokens;
}
if !buf.is_empty() {
chunks.push(buf.join("\n\n"));
}
chunks
}
/// 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<String> {
let mut out: Vec<String> = Vec::new();
let mut total = 0usize;
for piece in buf.iter().rev() {
let tokens = textkit::count(piece);
if total + tokens > overlap && !out.is_empty() {
break;
}
out.insert(0, piece.clone());
total += tokens;
}
out
}
/// Split anything still over the hard maximum, on token boundaries.
///
/// Defence in depth. The structural passes above should have handled it; this
/// is what guarantees the promise regardless.
pub fn enforce_hard_max(text: &str, hard_max: usize) -> Vec<String> {
textkit::split_at_token_limit(text, hard_max)
}
#[cfg(test)]
mod tests {
use super::*;
fn budgets(target: usize, overlap: usize) -> Budgets {
Budgets {
target,
overlap,
..Budgets::default()
}
}
#[test]
fn short_text_is_one_chunk() {
let chunks = 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());
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));
assert!(chunks.len() > 1, "{chunks:?}");
// Every chunk after the first repeats the end of the one before it.
for pair in chunks.windows(2) {
let earlier_tail = pair[0].rsplit("\n\n").next().expect("a paragraph");
assert!(
pair[1].starts_with(earlier_tail),
"{:?} does not carry {:?}",
pair[1],
earlier_tail
);
}
}
#[test]
fn an_oversized_paragraph_is_cut_at_sentences() {
let paragraph = "Alpha one. Beta two. Gamma three. Delta four. Epsilon five.";
let chunks = page(
paragraph,
&Budgets {
target: 6,
overlap: 2,
min_tokens: 1,
hard_max: 8,
},
);
assert!(chunks.len() > 1, "{chunks:?}");
assert!(chunks[0].contains(' '));
}
#[test]
fn an_empty_page_packs_to_nothing() {
assert!(page(" \n\n ", &Budgets::default()).is_empty());
}
#[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()];
assert_eq!(carry_tail(&buf, 1).len(), 1);
}
#[test]
fn nothing_leaves_the_packer_over_the_hard_maximum() {
let long = "word ".repeat(2000);
for piece in enforce_hard_max(&long, 100) {
assert!(textkit::count(&piece) <= 100);
}
}
}
crates/chunkmd/src/pages.rs
+294
View File
@@ -0,0 +1,294 @@
//! Finding where the pages and the chapters are.
//!
//! Converted markdown carries page boundaries as `<!-- Page 41 -->` comments
//! and chapter boundaries as headings. Both are needed on every chunk: a
//! passage a reader cannot locate in the book is a passage they cannot check.
use std::collections::HashMap;
use std::sync::LazyLock;
use regex::Regex;
static PAGE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"<!--\s*Page\s+(\d+)\s*-->").expect("a valid pattern"));
static HEADING: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"(?m)^(#{1,3})\s+(.+)$").expect("a valid pattern"));
/// One page's worth of text, and where it started in the source.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Page<'a> {
/// The page number from the marker, absent for text that came before the
/// first marker.
pub number: Option<u32>,
/// 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.
pub offset: usize,
}
/// A heading and the span of the document it covers.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Chapter {
/// The heading text, trimmed.
pub title: String,
/// Page the heading falls on, when the document has page markers.
pub start_page: Option<u32>,
/// Page the next heading falls on, or the last page.
pub end_page: Option<u32>,
/// Byte offset of the heading in the source markdown.
pub start_offset: usize,
}
/// Split markdown on its page markers.
///
/// Text *before* the first marker is still the book. Starting the walk at the
/// first marker silently deleted it — which cost one EPUB 362k of its 364k
/// characters, because it carried exactly one stray anchor near the very end.
/// 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.
pub fn split_by_page(markdown: &str) -> Vec<Page<'_>> {
let markers: Vec<_> = PAGE.find_iter(markdown).collect();
if markers.is_empty() {
return vec![Page {
number: None,
text: markdown,
offset: 0,
}];
}
let mut pages = Vec::new();
let head = trim(&markdown[..markers[0].start()]);
if !head.is_empty() {
pages.push(Page {
number: None,
text: head,
offset: 0,
});
}
for (at, marker) in markers.iter().enumerate() {
let start = marker.end();
let end = markers
.get(at + 1)
.map_or(markdown.len(), |next| next.start());
let text = trim(&markdown[start..end]);
if !text.is_empty() {
pages.push(Page {
number: Some(page_number(&markdown[marker.start()..marker.end()])),
text,
offset: start,
});
}
}
pages
}
/// Extract the headings that act as chapters, with their page and byte spans.
///
/// A document that has H1s is chaptered by its H1s; one that only has H2s is
/// chaptered by those. Mixing levels would make a chapter list where a
/// section and the chapter containing it sit side by side.
pub fn chapters(markdown: &str) -> Vec<Chapter> {
let mut headings: Vec<(usize, String, usize)> = Vec::new();
for found in HEADING.captures_iter(markdown) {
let title = trim(found.get(2).expect("group 2").as_str());
if title.is_empty() {
continue;
}
headings.push((
found.get(1).expect("group 1").len(),
title.to_string(),
found.get(0).expect("group 0").start(),
));
}
if headings.is_empty() {
return Vec::new();
}
let shallowest = headings
.iter()
.map(|(level, _, _)| *level)
.min()
.expect("a heading");
headings.retain(|(level, _, _)| *level == shallowest);
let page_at = PageIndex::new(markdown);
let mut chapters = Vec::with_capacity(headings.len());
for (at, (_, title, offset)) in headings.iter().enumerate() {
let end = headings
.get(at + 1)
.map_or(markdown.len(), |(_, _, next)| *next);
chapters.push(Chapter {
title: title.clone(),
start_page: page_at.at(*offset),
end_page: page_at.at(end),
start_offset: *offset,
});
}
chapters
}
/// Which page a byte offset falls on.
pub struct PageIndex(Vec<(usize, u32)>);
impl PageIndex {
/// Index the page markers in `markdown`.
pub fn new(markdown: &str) -> Self {
Self(
PAGE.find_iter(markdown)
.map(|marker| (marker.start(), page_number(marker.as_str())))
.collect(),
)
}
/// The last page marker at or before `offset`.
pub fn at(&self, offset: usize) -> Option<u32> {
let found = self.0.partition_point(|(start, _)| *start <= offset);
found.checked_sub(1).map(|at| self.0[at].1)
}
}
/// Which chapter a byte offset falls in.
pub struct ChapterIndex<'a>(&'a [Chapter]);
impl<'a> ChapterIndex<'a> {
/// Index a chapter list, which must be in document order.
pub fn new(chapters: &'a [Chapter]) -> Self {
Self(chapters)
}
/// The last chapter starting at or before `offset`.
pub fn at(&self, offset: usize) -> Option<&'a str> {
let found = self.0.partition_point(|c| c.start_offset <= offset);
found.checked_sub(1).map(|at| self.0[at].title.as_str())
}
}
/// Map each page number to the chapter covering it.
///
/// Where chapters overlap — which happens when two headings land on one page
/// — the later one wins, because that is the one a reader scrolling down
/// would be under.
pub fn page_chapters(chapters: &[Chapter]) -> HashMap<u32, &str> {
let mut out = HashMap::new();
for chapter in chapters {
let Some(start) = chapter.start_page else {
continue;
};
// A zero end page means the same thing as no end page: this chapter
// is one page long as far as anyone can tell.
let end = match chapter.end_page {
Some(0) | None => start,
Some(end) => end,
};
if end < start {
continue;
}
for page in start..=end {
out.insert(page, chapter.title.as_str());
}
}
out
}
/// Read the digits out of a page marker.
///
/// Digits are matched by the Unicode class, so a marker could in principle
/// carry non-ASCII decimal digits; those count the same way here as they
/// would in Python. A number too large for the field saturates rather than
/// wrapping — see `parity/KNOWN.md`.
fn page_number(marker: &str) -> u32 {
let mut number: u32 = 0;
for digit in marker.chars().filter_map(|c| c.to_digit(10)) {
number = number.saturating_mul(10).saturating_add(digit);
}
number
}
fn trim(text: &str) -> &str {
text.trim_matches(textkit::is_space)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn markdown_with_no_markers_is_one_unnumbered_page() {
let pages = split_by_page("just text");
assert_eq!(pages.len(), 1);
assert_eq!(pages[0].number, None);
assert_eq!(pages[0].text, "just text");
}
#[test]
fn text_before_the_first_marker_is_kept() {
let pages = split_by_page("prologue\n\n<!-- Page 19 -->\n\nbody");
assert_eq!(pages.len(), 2);
assert_eq!(pages[0].number, None);
assert_eq!(pages[0].text, "prologue");
assert_eq!(pages[1].number, Some(19));
assert_eq!(pages[1].text, "body");
}
#[test]
fn an_empty_page_is_dropped() {
let pages = split_by_page("<!-- Page 1 -->\n\n<!-- Page 2 -->\n\nbody");
assert_eq!(pages.len(), 1);
assert_eq!(pages[0].number, Some(2));
}
#[test]
fn markers_may_be_spelled_loosely() {
let pages = split_by_page("<!--Page 7-->\n\nbody");
assert_eq!(pages[0].number, Some(7));
}
#[test]
fn chapters_come_from_the_shallowest_heading_level_present() {
let markdown = "## Alpha\n\ntext\n\n### Sub\n\ntext\n\n## Beta\n\ntext";
let found = chapters(markdown);
assert_eq!(
found.iter().map(|c| c.title.as_str()).collect::<Vec<_>>(),
["Alpha", "Beta"]
);
}
#[test]
fn a_heading_with_nothing_after_it_is_not_a_chapter() {
assert!(chapters("# \n\n").is_empty());
}
#[test]
fn a_hash_on_its_own_line_takes_the_next_line_as_its_title() {
// The whitespace between the hashes and the title may include the
// newline, so a bare `#` adopts what follows it. Faithful to the
// Python, and the case does not arise in converted markdown.
let found = chapters("# \n\ntext");
assert_eq!(found.len(), 1);
assert_eq!(found[0].title, "text");
}
#[test]
fn chapters_carry_the_pages_they_span() {
let markdown = "<!-- Page 1 -->\n\n# One\n\ntext\n\n<!-- Page 4 -->\n\n# Two\n\ntext";
let found = chapters(markdown);
assert_eq!(found[0].start_page, Some(1));
assert_eq!(found[0].end_page, Some(4));
assert_eq!(found[1].start_page, Some(4));
let by_page = page_chapters(&found);
assert_eq!(by_page.get(&1), Some(&"One"));
assert_eq!(by_page.get(&3), Some(&"One"));
assert_eq!(by_page.get(&4), Some(&"Two"));
}
#[test]
fn an_offset_before_the_first_heading_has_no_chapter() {
let found = chapters("front\n\n# One\n\ntext");
let index = ChapterIndex::new(&found);
assert_eq!(index.at(0), None);
assert_eq!(index.at(10), Some("One"));
}
}
crates/chunkmd/src/role.rs
+317
View File
@@ -0,0 +1,317 @@
//! Telling body text from the apparatus around it.
//!
//! Retrieval wants body text. An index, a table of contents and a
//! bibliography all match a query beautifully and answer nothing — they are
//! lists of the words a reader is searching for. Marking them lets a search
//! exclude them while keeping them addressable for the consumers that do want
//! them: citation extraction reads bibliographies on purpose.
//!
//! Two signals, in order. The chapter title, when the document's headings
//! said what the section is; then the shape of the content, for back-matter
//! buried inside a body chapter — common in EPUBs whose bibliography heading
//! is not at the level the chapter map used.
use std::sync::LazyLock;
use corpus_core::Role;
use regex::Regex;
fn pattern(source: &str) -> Regex {
Regex::new(source).expect("a valid pattern")
}
/// Chapter titles that give the section away, tried in this order.
static BY_TITLE: LazyLock<Vec<(Role, Regex)>> = LazyLock::new(|| {
vec![
(
Role::Bibliography,
pattern(
r"(?i)(bibliograph|works\s+cited|literature\s+cited|further\s+reading|^\s*references\s*$|^\s*(end)?notes\s*$)",
),
),
(Role::Index, pattern(r"(?i)^\s*index(\s+of\s+\w+)?\s*$")),
(Role::Toc, pattern(r"(?i)^\s*(table\s+of\s+)?contents\s*$")),
(
Role::Frontmatter,
pattern(
r"(?i)(copyright|title\s*page|colophon|^\s*dedication\s*$|^\s*epigraph\s*$|about\s+the\s+(author|publisher)|acknowledg)",
),
),
]
});
/// A standalone bibliography heading inside the chunk's own text —
/// "Bibliography", "## Works Cited", "Selected Bibliography".
static INTEXT_BIBLIO_HEADING: LazyLock<Regex> = LazyLock::new(|| {
pattern(
r"(?im)^\s*#{0,6}\s*(select(?:ed)?\s+)?(bibliography|works\s+cited|references|sources)\s*$",
)
});
/// One bibliography entry: "Lastname, First…", or the dash run that stands
/// for "the same author again".
static BIBLIO_ENTRY: LazyLock<Regex> =
LazyLock::new(|| pattern(r"^(?:[A-Z][\w'’\-]+,\s+\S|[—–-]{2,3}\s*[.,])"));
/// The continuation of a wrapped entry: "New York: Springer, 2013."
static BIBLIO_CONTINUATION: LazyLock<Regex> =
LazyLock::new(|| pattern(r"\b(University|Press|Publish|Books|Verlag|Editions|& Co|Inc\.)\b"));
static YEAR: LazyLock<Regex> = LazyLock::new(|| pattern(r"\b(1[5-9]\d{2}|20\d{2})\b"));
/// An index line: a short lead, then a list of page numbers.
static INDEX_LINE: LazyLock<Regex> =
LazyLock::new(|| pattern(r"^.{1,70}?,\s+[ivxlc\d]{1,4}(?:[,–-]\s*[ivxlc\d]{1,4})*\s*$"));
/// A contents line: dot leaders to a page number, or a link to an internal
/// anchor, which is what a converted EPUB's table of contents becomes.
static TOC_LINE: LazyLock<Regex> = LazyLock::new(|| pattern(r"(\.{3,}\s*\d{1,4}\s*$|\]\(#)"));
/// A wrapped entry long enough that its publisher name is probably a
/// coincidence rather than a citation, in characters.
const CONTINUATION_MAX: usize = 200;
/// How many entry-shaped paragraphs make a bibliography, and what share of
/// the paragraphs they have to be.
const BIBLIO_MIN_ENTRIES: usize = 3;
const BIBLIO_SHARE: f64 = 0.55;
/// The same, for an in-text bibliography heading, which is evidence on its
/// own and so needs less corroboration.
const BIBLIO_MIN_ENTRIES_WITH_HEADING: usize = 2;
/// How many list-shaped lines make an index or a table of contents.
const LIST_MIN_LINES: usize = 5;
const LIST_SHARE: f64 = 0.5;
/// Classify one chunk.
pub fn classify(text: &str, chapter: Option<&str>) -> Role {
if let Some(chapter) = chapter.filter(|title| !title.is_empty()) {
for (role, pattern) in BY_TITLE.iter() {
if pattern.is_match(chapter) {
return *role;
}
}
}
// Content-shape fallbacks. Cheap line scans; a chunk is a thousand
// tokens at the outside.
let lines: Vec<&str> = text
.split('\n')
.map(trim)
.filter(|line| !line.is_empty())
.collect();
if lines.is_empty() {
return Role::Body;
}
let paragraphs = textkit::paragraphs(text);
let entryish = paragraphs
.iter()
.filter(|paragraph| looks_like_entry(paragraph))
.count();
let enough_entries =
entryish >= BIBLIO_MIN_ENTRIES && entryish as f64 >= BIBLIO_SHARE * paragraphs.len() as f64;
let headed =
INTEXT_BIBLIO_HEADING.is_match(text) && entryish >= BIBLIO_MIN_ENTRIES_WITH_HEADING;
if enough_entries || headed {
return Role::Bibliography;
}
if dominates(
lines
.iter()
.filter(|line| INDEX_LINE.is_match(line))
.count(),
lines.len(),
) {
return Role::Index;
}
if dominates(
lines.iter().filter(|line| TOC_LINE.is_match(line)).count(),
lines.len(),
) {
return Role::Toc;
}
Role::Body
}
fn looks_like_entry(paragraph: &str) -> bool {
let year = YEAR.is_match(paragraph);
let publisher = BIBLIO_CONTINUATION.is_match(paragraph);
if BIBLIO_ENTRY.is_match(paragraph) && (year || publisher) {
return true;
}
publisher && year && paragraph.chars().count() < CONTINUATION_MAX
}
fn dominates(matching: usize, total: usize) -> bool {
matching >= LIST_MIN_LINES && matching as f64 >= LIST_SHARE * total as f64
}
/// Fix classification stragglers.
///
/// A run of at most two body chunks sandwiched between chunks of the same
/// non-body role inherits that role. This catches the boundary chunks of a
/// bibliography or an index, whose own content shape is ambiguous because
/// the entries there are URL-heavy, or split across the seam.
pub fn smooth(roles: &mut [Role]) {
const MAX_RUN: usize = 2;
let mut at = 0;
while at < roles.len() {
if roles[at] != Role::Body {
at += 1;
continue;
}
let mut end = at;
while end < roles.len() && roles[end] == Role::Body {
end += 1;
}
let before = at.checked_sub(1).map(|before| roles[before]);
let after = roles.get(end).copied();
if let Some(role) = before
&& end - at <= MAX_RUN
&& before == after
&& role != Role::Body
{
roles[at..end].fill(role);
}
at = end;
}
}
fn trim(text: &str) -> &str {
text.trim_matches(textkit::is_space)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_chapter_title_decides_when_it_can() {
for (title, role) in [
("Selected Bibliography", Role::Bibliography),
("Works Cited", Role::Bibliography),
("References", Role::Bibliography),
("Notes", Role::Bibliography),
("Endnotes", Role::Bibliography),
("Index", Role::Index),
("Index of Names", Role::Index),
("Table of Contents", Role::Toc),
("Contents", Role::Toc),
("Copyright", Role::Frontmatter),
("Acknowledgements", Role::Frontmatter),
("About the Author", Role::Frontmatter),
("III. The Second Estate", Role::Body),
] {
assert_eq!(
classify("ordinary text", Some(title)),
role,
"for {title:?}"
);
}
}
#[test]
fn a_title_that_merely_mentions_references_in_a_sentence_is_not_a_bibliography() {
assert_eq!(
classify("text", Some("What the references tell us")),
Role::Body
);
}
#[test]
fn entry_shaped_paragraphs_make_a_bibliography_without_a_title() {
let text = "Arendt, Hannah. The Human Condition. Chicago: University Press, 1958.\n\n\
Bloch, Marc. Feudal Society. London: Routledge, 1961.\n\n\
Cassirer, Ernst. The Myth of the State. New Haven: Yale Press, 1946.";
assert_eq!(classify(text, None), Role::Bibliography);
}
#[test]
fn an_in_text_heading_lowers_the_bar_to_two_entries() {
let text = "## Works Cited\n\n\
Arendt, Hannah. The Human Condition. Chicago: University Press, 1958.\n\n\
Bloch, Marc. Feudal Society. London: Routledge, 1961.";
assert_eq!(classify(text, None), Role::Bibliography);
}
#[test]
fn page_number_lists_make_an_index() {
let text = "agriculture, 12, 34, 102\n\
bishops, 7\n\
charters, 88, 91\n\
demesne, 14-16\n\
estates, 3, 44\n\
fealty, 201";
assert_eq!(classify(text, None), Role::Index);
}
#[test]
fn dot_leaders_make_a_table_of_contents() {
let text = "One. The Setting......1\n\
Two. The Quarrel......22\n\
Three. The Break......45\n\
Four. The Peace......70\n\
Five. After......99";
assert_eq!(classify(text, None), Role::Toc);
}
#[test]
fn ordinary_prose_is_body() {
let text = "The second estate held its land in exchange for service, and \
the service was military. That much is agreed; what follows is not.";
assert_eq!(classify(text, None), Role::Body);
}
#[test]
fn empty_text_is_body() {
assert_eq!(classify(" \n ", None), Role::Body);
}
#[test]
fn a_short_body_run_inside_a_bibliography_is_absorbed() {
let mut roles = vec![
Role::Bibliography,
Role::Body,
Role::Body,
Role::Bibliography,
];
smooth(&mut roles);
assert_eq!(roles, [Role::Bibliography; 4]);
}
#[test]
fn a_long_body_run_is_left_alone() {
let mut roles = vec![
Role::Bibliography,
Role::Body,
Role::Body,
Role::Body,
Role::Bibliography,
];
let before = roles.clone();
smooth(&mut roles);
assert_eq!(roles, before);
}
#[test]
fn a_run_between_different_roles_is_left_alone() {
let mut roles = vec![Role::Bibliography, Role::Body, Role::Index];
let before = roles.clone();
smooth(&mut roles);
assert_eq!(roles, before);
}
#[test]
fn body_at_either_end_is_left_alone() {
let mut roles = vec![Role::Body, Role::Index, Role::Body];
let before = roles.clone();
smooth(&mut roles);
assert_eq!(roles, before);
}
}