graph_r/
text.rs

1// Author: spicyzhug@gmail.com
2// Copyright: lituus-io, all rights reserved
3//
4//! The text rules every index shares: tokenization, case folding, char
5//! n-grams, stopwords and the token estimate behind budgets. One
6//! implementation for graph-r's lexicon, link-r's BM25 and embedder, and
7//! their users — a query term and an indexed term always agree.
8//! Allocation-free apart from the case-folding of non-lowercase tokens.
9
10use std::borrow::Cow;
11
12/// Tokens: maximal runs of Unicode alphanumerics, borrowed from the input.
13/// Not case-folded (see [`normalize`] and [`for_each_term`]).
14pub fn tokens(text: &str) -> impl Iterator<Item = &str> {
15    text.split(|c: char| !c.is_alphanumeric()).filter(|t| !t.is_empty())
16}
17
18/// A token's canonical lowercase form; borrowed when it already is one.
19#[must_use]
20pub fn normalize(token: &str) -> Cow<'_, str> {
21    if token.chars().all(|c| !c.is_uppercase()) {
22        Cow::Borrowed(token)
23    } else if token.is_ascii() {
24        Cow::Owned(token.to_ascii_lowercase())
25    } else {
26        Cow::Owned(token.to_lowercase())
27    }
28}
29
30/// Run `f` on every case-folded token of `text`, through one reused buffer.
31pub fn for_each_term(text: &str, mut f: impl FnMut(&str)) {
32    let mut buf = String::new();
33    for t in tokens(text) {
34        match normalize(t) {
35            Cow::Borrowed(t) => f(t),
36            Cow::Owned(o) => {
37                buf.clear();
38                buf.push_str(&o);
39                f(&buf);
40            }
41        }
42    }
43}
44
45/// Run `f` on each `n`-character window of `token` (on char boundaries); a
46/// token of at most `n` characters is emitted whole. No allocation.
47pub fn for_each_char_ngram(token: &str, n: usize, mut f: impl FnMut(&str)) {
48    if n == 0 || token.is_empty() {
49        return;
50    }
51    let mut ends = token.char_indices().map(|(i, _)| i).chain(std::iter::once(token.len()));
52    let mut starts = token.char_indices().map(|(i, _)| i);
53    // `end` runs `n` characters ahead of `start`.
54    let Some(mut end) = ends.nth(n) else {
55        f(token);
56        return;
57    };
58    for start in starts.by_ref() {
59        f(&token[start..end]);
60        match ends.next() {
61            Some(e) => end = e,
62            None => break,
63        }
64    }
65}
66
67/// A compact English stopword list, sorted for binary search.
68static STOPWORDS: &[&str] = &[
69    "a", "about", "above", "after", "again", "all", "am", "an", "and", "any", "are", "as", "at",
70    "be", "because", "been", "before", "being", "below", "between", "both", "but", "by", "can",
71    "did", "do", "does", "doing", "down", "during", "each", "few", "for", "from", "further", "had",
72    "has", "have", "having", "he", "her", "here", "hers", "him", "his", "how", "i", "if", "in",
73    "into", "is", "it", "its", "itself", "just", "me", "more", "most", "my", "no", "nor", "not",
74    "of", "off", "on", "once", "only", "or", "other", "our", "out", "over", "own", "same", "she",
75    "should", "so", "some", "such", "than", "that", "the", "their", "them", "then", "there",
76    "these", "they", "this", "those", "through", "to", "too", "under", "until", "up", "very",
77    "was", "we", "were", "what", "when", "where", "which", "while", "who", "whom", "why", "will",
78    "with", "you", "your",
79];
80
81/// Whether a case-folded token is a common English stopword.
82#[must_use]
83pub fn is_stopword(term: &str) -> bool {
84    STOPWORDS.binary_search(&term).is_ok()
85}
86
87/// Model tokens in `bytes` of text: the 4-bytes-per-token rule every budget
88/// (graph-r query, prompt assembly, request sizing) uses.
89#[must_use]
90#[inline]
91pub fn estimate_tokens(bytes: usize) -> u64 {
92    (bytes as u64).div_ceil(4)
93}
94
95/// The text bytes a budget of `tokens` admits (the inverse of
96/// [`estimate_tokens`]).
97#[must_use]
98#[inline]
99pub fn bytes_for_tokens(tokens: u64) -> usize {
100    usize::try_from(tokens.saturating_mul(4)).unwrap_or(usize::MAX)
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    #[test]
108    fn tokens_are_zero_copy_unicode_runs() {
109        assert_eq!(
110            tokens("Hello, world!  foo_bar.baz").collect::<Vec<_>>(),
111            ["Hello", "world", "foo", "bar", "baz"]
112        );
113        assert_eq!(tokens("café—über").collect::<Vec<_>>(), ["café", "über"]);
114    }
115
116    #[test]
117    fn terms_fold_case_and_borrow_when_they_can() {
118        assert!(matches!(normalize("already"), Cow::Borrowed(_)));
119        assert_eq!(normalize("ÜBER"), "über");
120        let mut t = Vec::new();
121        for_each_term("Install FooBar v2.1 (Löwe)", |x| t.push(x.to_owned()));
122        assert_eq!(t, ["install", "foobar", "v2", "1", "löwe"]);
123    }
124
125    #[test]
126    fn char_ngrams_window_on_char_boundaries() {
127        let grams = |s: &str, n| {
128            let mut v = Vec::new();
129            for_each_char_ngram(s, n, |g| v.push(g.to_owned()));
130            v
131        };
132        assert_eq!(grams("hello", 3), ["hel", "ell", "llo"]);
133        assert_eq!(grams("hi", 3), ["hi"]);
134        assert_eq!(grams("abc", 3), ["abc"]);
135        assert_eq!(grams("café", 2), ["ca", "af", "fé"]);
136        assert!(grams("", 2).is_empty());
137    }
138
139    #[test]
140    fn stopwords_and_estimate() {
141        assert!(is_stopword("the") && !is_stopword("bigquery"));
142        assert!(STOPWORDS.windows(2).all(|w| w[0] < w[1]));
143        assert_eq!((estimate_tokens(0), estimate_tokens(1), estimate_tokens(8)), (0, 1, 2));
144        assert_eq!(estimate_tokens(bytes_for_tokens(250)), 250);
145    }
146}
graph_r/text.rs

graph_r/
text.rs
