From 0208b14b5dfbe3d54839660a28830fbcd972b460 Mon Sep 17 00:00:00 2001 From: funkpower Date: Sat, 19 Sep 2026 11:04:27 -0400 Subject: [PATCH] textkit: token counting and the two segmenters --- Cargo.lock | 121 +++++++++++++++++++++++ Cargo.toml | 2 + crates/textkit/Cargo.toml | 15 +++ crates/textkit/src/lib.rs | 17 ++++ crates/textkit/src/segment.rs | 181 ++++++++++++++++++++++++++++++++++ crates/textkit/src/tokens.rs | 118 ++++++++++++++++++++++ 6 files changed, 454 insertions(+) create mode 100644 crates/textkit/Cargo.toml create mode 100644 crates/textkit/src/lib.rs create mode 100644 crates/textkit/src/segment.rs create mode 100644 crates/textkit/src/tokens.rs diff --git a/Cargo.lock b/Cargo.lock index 292d1e5..8cda8ba 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,53 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "aho-corasick" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + +[[package]] +name = "bit-set" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3" +dependencies = [ + "bit-vec", +] + +[[package]] +name = "bit-vec" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" + +[[package]] +name = "bstr" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bb31b46c14244e20ee9984b11bf5c992b91fb6939fea616e3512c8baecdbe5f" +dependencies = [ + "memchr", + "regex-automata", + "serde_core", +] + [[package]] name = "corpus-core" version = "0.1.0" @@ -18,12 +65,29 @@ version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" +[[package]] +name = "fancy-regex" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72cf461f865c862bb7dc573f643dd6a2b6842f7c30b07882b56bd148cc2761b8" +dependencies = [ + "bit-set", + "regex-automata", + "regex-syntax", +] + [[package]] name = "itoa" version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + [[package]] name = "memchr" version = "2.8.3" @@ -68,6 +132,41 @@ dependencies = [ "syn", ] +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + [[package]] name = "schemars" version = "1.2.2" @@ -158,6 +257,13 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "textkit" +version = "0.1.0" +dependencies = [ + "tiktoken-rs", +] + [[package]] name = "thiserror" version = "2.0.20" @@ -178,6 +284,21 @@ dependencies = [ "syn", ] +[[package]] +name = "tiktoken-rs" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "027853bbf8c7763b77c5c595f1c271c7d536ced7d6f83452911b944621e57fc2" +dependencies = [ + "anyhow", + "base64", + "bstr", + "fancy-regex", + "lazy_static", + "regex", + "rustc-hash", +] + [[package]] name = "unicode-ident" version = "1.0.26" diff --git a/Cargo.toml b/Cargo.toml index f745776..a5c05e3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,6 +11,7 @@ repository = "https://git.pumpin.online/funkpower/catena" [workspace.dependencies] corpus-core = { path = "crates/corpus-core", version = "0.1.0" } +textkit = { path = "crates/textkit", version = "0.1.0" } anyhow = "1" clap = { version = "4", features = ["derive"] } @@ -18,6 +19,7 @@ schemars = "1" serde = { version = "1", features = ["derive"] } serde_json = "1" thiserror = "2" +tiktoken-rs = "0.12" [profile.release] strip = true diff --git a/crates/textkit/Cargo.toml b/crates/textkit/Cargo.toml new file mode 100644 index 0000000..842dbf9 --- /dev/null +++ b/crates/textkit/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "textkit" +description = "Tokenising and segmenting text, the way the catena tools agree to do it." +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +tiktoken-rs.workspace = true + +[lints.rust] +unsafe_code = "forbid" +missing_docs = "warn" diff --git a/crates/textkit/src/lib.rs b/crates/textkit/src/lib.rs new file mode 100644 index 0000000..dbcbe9f --- /dev/null +++ b/crates/textkit/src/lib.rs @@ -0,0 +1,17 @@ +//! Tokenising and segmenting text, the way the catena tools agree to do it. +//! +//! Two jobs, both of which have to give the same answer everywhere or chunk +//! boundaries move: counting tokens the way the embedding stack counts them, +//! and cutting text into sentences and paragraphs the way the existing +//! pipelines cut it. +//! +//! The segmenters reproduce Python's `re` semantics deliberately, down to +//! which characters count as whitespace, because their output decides where a +//! chunk ends and a chunk boundary that moves is a chunk that no longer +//! matches what is already in an index. + +pub mod segment; +pub mod tokens; + +pub use segment::{is_space, paragraphs, sentences}; +pub use tokens::{count, decode_lossy, encode, split_at_token_limit}; diff --git a/crates/textkit/src/segment.rs b/crates/textkit/src/segment.rs new file mode 100644 index 0000000..491e9ec --- /dev/null +++ b/crates/textkit/src/segment.rs @@ -0,0 +1,181 @@ +//! Cutting text into sentences and paragraphs. +//! +//! Both segmenters reproduce what Python's `re` module does with +//! `(?<=[.!?])\s+` and `\n\s*\n`, because the chunk boundaries already in +//! existing indexes were cut that way. Where the obvious Rust spelling would +//! differ, the difference is spelled out below rather than smoothed over. + +/// Whether a character is whitespace, by Python's definition. +/// +/// Python's `str.isspace` and its `\s` class cover the Unicode `White_Space` +/// property *and* the four ASCII separators `U+001C`–`U+001F`, which Rust's +/// [`char::is_whitespace`] does not. Trimming with the narrower definition +/// leaves a separator on the front of a piece, which changes its token count, +/// which moves a chunk boundary. +pub fn is_space(c: char) -> bool { + c.is_whitespace() || matches!(c, '\u{1c}'..='\u{1f}') +} + +/// Trim leading and trailing whitespace, by Python's definition. +fn trim(s: &str) -> &str { + s.trim_matches(is_space) +} + +/// Drop the empty pieces and trim what is left, which is what every caller of +/// a Python `re.split` here does immediately afterwards. +fn tidy(pieces: Vec<&str>) -> Vec<&str> { + pieces + .into_iter() + .map(trim) + .filter(|piece| !piece.is_empty()) + .collect() +} + +/// Split text into sentences at whitespace that follows `.`, `!` or `?`. +/// +/// Crude on purpose — it does not know that `Mr.` or `p. 41` are not +/// sentence ends. This is the fallback path for a paragraph too large to +/// chunk whole, where a wrong cut costs a little context and a clever cut +/// that disagrees with the existing index costs a re-embedding. +pub fn sentences(text: &str) -> Vec<&str> { + let mut pieces = Vec::new(); + let mut piece_start = 0; + let mut previous = None; + let mut chars = text.char_indices().peekable(); + + while let Some((at, c)) = chars.next() { + if !is_space(c) { + previous = Some(c); + continue; + } + // Whitespace is consumed a whole run at a time, because `\s+` is + // greedy: the separator is the run, not its first character. + let mut run_end = at + c.len_utf8(); + while let Some(&(next_at, next)) = chars.peek() { + if !is_space(next) { + break; + } + run_end = next_at + next.len_utf8(); + chars.next(); + } + if matches!(previous, Some('.' | '!' | '?')) { + pieces.push(&text[piece_start..at]); + piece_start = run_end; + } + previous = None; + } + + pieces.push(&text[piece_start..]); + tidy(pieces) +} + +/// Split text into paragraphs at blank lines. +/// +/// A separator is a newline, any run of whitespace, and a newline — so one +/// blank line and four of them separate equally, and trailing spaces on an +/// otherwise blank line do not stop it being blank. +pub fn paragraphs(text: &str) -> Vec<&str> { + let mut pieces = Vec::new(); + let mut piece_start = 0; + let mut chars = text.char_indices().peekable(); + + while let Some((at, c)) = chars.next() { + if c != '\n' { + continue; + } + // The greedy `\s*` in `\n\s*\n` backtracks to the *last* newline in + // the whitespace run, so that is where the next paragraph begins. + let mut last_newline = None; + while let Some(&(next_at, next)) = chars.peek() { + if !is_space(next) { + break; + } + if next == '\n' { + last_newline = Some(next_at); + } + chars.next(); + } + if let Some(newline) = last_newline { + pieces.push(&text[piece_start..at]); + piece_start = newline + 1; + } + } + + pieces.push(&text[piece_start..]); + tidy(pieces) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn sentences_split_after_terminators() { + assert_eq!( + sentences("One. Two! Three? Four"), + ["One.", "Two!", "Three?", "Four"] + ); + } + + #[test] + fn a_terminator_without_following_whitespace_is_not_a_boundary() { + assert_eq!(sentences("etc.and on"), ["etc.and on"]); + assert_eq!(sentences("3.14 is pi"), ["3.14 is pi"]); + } + + #[test] + fn the_whole_whitespace_run_is_the_separator() { + assert_eq!(sentences("One.\n\n Two."), ["One.", "Two."]); + } + + #[test] + fn leading_whitespace_is_not_a_boundary() { + assert_eq!(sentences(" One. Two."), ["One.", "Two."]); + } + + #[test] + fn empty_and_blank_text_yield_nothing() { + assert!(sentences("").is_empty()); + assert!(sentences(" \n ").is_empty()); + assert!(paragraphs("").is_empty()); + assert!(paragraphs("\n\n\n").is_empty()); + } + + #[test] + fn paragraphs_split_on_a_blank_line() { + assert_eq!(paragraphs("one\n\ntwo"), ["one", "two"]); + } + + #[test] + fn a_single_newline_is_not_a_paragraph_break() { + assert_eq!(paragraphs("one\ntwo"), ["one\ntwo"]); + } + + #[test] + fn a_blank_line_of_spaces_still_separates() { + assert_eq!(paragraphs("one\n \t \ntwo"), ["one", "two"]); + } + + #[test] + fn many_blank_lines_separate_once() { + assert_eq!(paragraphs("one\n\n\n\n two"), ["one", "two"]); + } + + #[test] + fn carriage_returns_survive_the_split() { + assert_eq!(paragraphs("one\r\n\r\ntwo"), ["one", "two"]); + } + + #[test] + fn ascii_separators_count_as_whitespace() { + // Python's str.isspace is true for these; Rust's is_whitespace is not. + assert!(is_space('\u{1c}')); + assert_eq!(sentences("One.\u{1c}Two."), ["One.", "Two."]); + } + + #[test] + fn multi_byte_characters_do_not_shift_the_boundaries() { + assert_eq!(sentences("Café. Naïve."), ["Café.", "Naïve."]); + assert_eq!(paragraphs("Café\n\nNaïve"), ["Café", "Naïve"]); + } +} diff --git a/crates/textkit/src/tokens.rs b/crates/textkit/src/tokens.rs new file mode 100644 index 0000000..5e4035f --- /dev/null +++ b/crates/textkit/src/tokens.rs @@ -0,0 +1,118 @@ +//! Token counting, against the same tokeniser the embedding stack uses. + +use tiktoken_rs::{CoreBPE, cl100k_base_singleton}; + +/// The tokeniser every tool in this workspace counts with. +/// +/// `cl100k_base`, loaded from data compiled into the binary — no download, no +/// cache directory, nothing to go stale. Counting against a different +/// tokeniser than the one the index was built with silently moves every chunk +/// boundary, which is why this is one function and not a parameter. +fn bpe() -> &'static CoreBPE { + cl100k_base_singleton() +} + +/// Token ids for `text`, with special-token sequences treated as ordinary +/// text rather than as control tokens. +/// +/// Text out of a PDF can contain anything, including the literal string +/// `<|endoftext|>`; encoding it as a control token would be both wrong and a +/// way for a document to reach into the model's input. +pub fn encode(text: &str) -> Vec { + bpe().encode_ordinary(text) +} + +/// How many tokens `text` is. +pub fn count(text: &str) -> usize { + encode(text).len() +} + +/// Text for `tokens`, with undecodable byte sequences replaced. +/// +/// Cutting a token list at an arbitrary index can split a multi-byte +/// character, so decoding a slice has to be lossy or it would fail on exactly +/// the inputs that need it. +pub fn decode_lossy(tokens: &[u32]) -> String { + match bpe().decode_bytes(tokens) { + Ok(bytes) => String::from_utf8_lossy(&bytes).into_owned(), + Err(_) => String::new(), + } +} + +/// Split `text` into pieces of at most `limit` tokens, cutting on token +/// boundaries. +/// +/// The last resort of a chunker: it guarantees a maximum but respects no +/// other structure, so it should only ever see text that the structural +/// passes could not get under the limit. Returns `text` unchanged when it +/// already fits. A `limit` of zero yields nothing, since no piece could +/// satisfy it. +pub fn split_at_token_limit(text: &str, limit: usize) -> Vec { + if limit == 0 { + return Vec::new(); + } + let ids = encode(text); + if ids.len() <= limit { + return vec![text.to_string()]; + } + ids.chunks(limit).map(decode_lossy).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn counts_agree_with_the_python_tokeniser() { + // Values from `tiktoken.get_encoding("cl100k_base").encode(s, + // disallowed_special=())`. + for (text, expected) in [ + ("", 0), + ("hello world", 2), + ("The quick brown fox jumps over the lazy dog.", 10), + ("\n\n", 1), + ("naïve café", 4), + ] { + assert_eq!(count(text), expected, "counting {text:?}"); + } + } + + #[test] + fn a_special_token_spelling_is_just_text() { + // Seven ordinary tokens, not the one control token it names. + assert_eq!(count("<|endoftext|>"), 7); + } + + #[test] + fn text_under_the_limit_comes_back_untouched() { + assert_eq!(split_at_token_limit("hello world", 10), ["hello world"]); + } + + #[test] + fn splitting_preserves_every_character() { + let text = "The quick brown fox jumps over the lazy dog, repeatedly and at length."; + let pieces = split_at_token_limit(text, 3); + assert!(pieces.len() > 1); + assert_eq!(pieces.concat(), text); + } + + #[test] + fn no_piece_exceeds_the_limit() { + let text = "alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu"; + for piece in split_at_token_limit(text, 4) { + assert!(count(&piece) <= 4, "piece {piece:?} is over the limit"); + } + } + + #[test] + fn a_zero_limit_yields_nothing() { + assert!(split_at_token_limit("anything", 0).is_empty()); + } + + #[test] + fn splitting_a_multi_byte_character_is_lossy_not_fatal() { + // Every emoji is several tokens, so a limit of one cuts inside one. + let pieces = split_at_token_limit("🙂🙂🙂", 1); + assert!(!pieces.is_empty()); + } +}