mirror of
https://git.pumpin.online/funkpower/catena.git
synced 2026-10-22 18:57:14 +00:00
textkit: token counting and the two segmenters
This commit is contained in:
funkpower
@@ -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"
|
||||
@@ -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};
|
||||
@@ -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"]);
|
||||
}
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user