textkit: token counting and the two segmenters

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"

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

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"

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};

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"]);
+    }
+}

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<u32> {
+    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<String> {
+    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());
+    }
+}