Skip to main content

kevy_text/
token.rs

1//! Script-aware, dictionary-free tokenization (RFC D1).
2//!
3//! - Latin/ASCII alphanumeric runs → lowercased word tokens (length
4//!   ≥ 2; single characters are noise).
5//! - CJK (unified ideographs, kana, hangul syllables) → adjacent
6//!   BIGRAMS; a lone CJK character emits a unigram so single-char
7//!   texts remain findable.
8//! - Tokens never cross a script boundary.
9
10/// Pluggable tokenizer (v2.7 ships only [`tokenize`], the default).
11pub trait Tokenizer {
12    /// Produce tokens for `text` (UTF-8; invalid bytes are skipped).
13    fn tokens(&self, text: &[u8]) -> Vec<Vec<u8>>;
14}
15
16/// The default dictionary-free tokenizer.
17#[derive(Debug, Clone, Copy, Default)]
18pub struct KevyTokenizer;
19
20impl Tokenizer for KevyTokenizer {
21    fn tokens(&self, text: &[u8]) -> Vec<Vec<u8>> {
22        tokenize(text)
23    }
24}
25
26fn is_cjk(c: char) -> bool {
27    matches!(c,
28        '\u{4E00}'..='\u{9FFF}'   // CJK unified ideographs
29        | '\u{3400}'..='\u{4DBF}' // extension A
30        | '\u{3040}'..='\u{309F}' // hiragana
31        | '\u{30A0}'..='\u{30FF}' // katakana
32        | '\u{AC00}'..='\u{D7AF}' // hangul syllables
33    )
34}
35
36/// Tokenize per the default rules (see module doc).
37pub fn tokenize(text: &[u8]) -> Vec<Vec<u8>> {
38    let s = String::from_utf8_lossy(text);
39    let mut out = Vec::new();
40    let mut word = String::new();
41    let mut prev_cjk: Option<char> = None;
42    let mut cjk_run = 0usize;
43
44    let flush_word = |word: &mut String, out: &mut Vec<Vec<u8>>| {
45        if word.chars().count() >= 2 {
46            out.push(word.to_lowercase().into_bytes());
47        }
48        word.clear();
49    };
50
51    for c in s.chars() {
52        if is_cjk(c) {
53            flush_word(&mut word, &mut out);
54            if let Some(p) = prev_cjk {
55                let mut bi = String::with_capacity(8);
56                bi.push(p);
57                bi.push(c);
58                out.push(bi.into_bytes());
59            }
60            prev_cjk = Some(c);
61            cjk_run += 1;
62        } else {
63            // end of a CJK run: a lone character still emits itself
64            if cjk_run == 1
65                && let Some(p) = prev_cjk
66            {
67                out.push(p.to_string().into_bytes());
68            }
69            prev_cjk = None;
70            cjk_run = 0;
71            if c.is_alphanumeric() {
72                word.push(c);
73            } else {
74                flush_word(&mut word, &mut out);
75            }
76        }
77    }
78    if cjk_run == 1
79        && let Some(p) = prev_cjk
80    {
81        out.push(p.to_string().into_bytes());
82    }
83    flush_word(&mut word, &mut out);
84    out
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    fn toks(s: &str) -> Vec<String> {
92        tokenize(s.as_bytes())
93            .into_iter()
94            .map(|t| String::from_utf8(t).unwrap())
95            .collect()
96    }
97
98    #[test]
99    fn latin_words_lowercased_min_two() {
100        assert_eq!(toks("Hello, Rust world! a I"), vec!["hello", "rust", "world"]);
101        assert_eq!(toks("v2.7-alpha"), vec!["v2", "alpha"]);
102    }
103
104    #[test]
105    fn cjk_bigrams_no_dictionary() {
106        assert_eq!(toks("全文检索"), vec!["全文", "文检", "检索"]);
107        assert_eq!(toks("猫"), vec!["猫"], "lone CJK char = unigram");
108        assert_eq!(toks("ひらがな"), vec!["ひら", "らが", "がな"]);
109        assert_eq!(toks("한국어"), vec!["한국", "국어"]);
110    }
111
112    #[test]
113    fn mixed_scripts_never_cross() {
114        assert_eq!(
115            toks("Rust搜索engine引擎x"),
116            vec!["rust", "搜索", "engine", "引擎"],
117            "boundaries split; trailing single latin char dropped"
118        );
119        assert_eq!(toks("日本語abc漢字"), vec!["日本", "本語", "abc", "漢字"]);
120    }
121
122    #[test]
123    fn junk_and_empty() {
124        assert!(toks("").is_empty());
125        assert!(toks("!!! ...").is_empty());
126        let bad = tokenize(&[0xFF, 0xFE, b'o', b'k', b'a', b'y']);
127        assert_eq!(bad, vec![b"okay".to_vec()]);
128    }
129}