khive_text/
tokenizer.rs

1//! Tokenizer implementations: whitespace, CJK-character, keyword, identifier, unicode-word.
2
3use crate::identifier::{is_identifier, split_identifier};
4use crate::lang::is_cjk_char;
5use crate::Tokenizer;
6
7// ---------------------------------------------------------------------------
8// WhitespaceTokenizer
9// ---------------------------------------------------------------------------
10
11/// Splits on ASCII whitespace, trims leading/trailing ASCII punctuation from
12/// each token, and drops empty results.
13#[derive(Debug, Default, Clone)]
14pub struct WhitespaceTokenizer;
15
16impl Tokenizer for WhitespaceTokenizer {
17    fn tokenize(&self, text: &str) -> Vec<String> {
18        text.split_whitespace()
19            .map(|w| {
20                w.trim_matches(|c: char| c.is_ascii_punctuation())
21                    .to_string()
22            })
23            .filter(|t| !t.is_empty())
24            .collect()
25    }
26}
27
28// ---------------------------------------------------------------------------
29// CjkCharTokenizer
30// ---------------------------------------------------------------------------
31
32/// Emits each CJK character as its own token; non-CJK runs split on whitespace with ASCII punctuation stripped.
33#[derive(Debug, Default, Clone)]
34pub struct CjkCharTokenizer;
35
36impl Tokenizer for CjkCharTokenizer {
37    fn tokenize(&self, text: &str) -> Vec<String> {
38        let mut tokens = Vec::new();
39        let mut latin_buf = String::new();
40
41        let flush = |buf: &mut String, out: &mut Vec<String>| {
42            for part in buf.split_whitespace() {
43                let p = part.trim_matches(|c: char| c.is_ascii_punctuation());
44                if !p.is_empty() {
45                    out.push(p.to_string());
46                }
47            }
48            buf.clear();
49        };
50
51        for ch in text.chars() {
52            if is_cjk_char(ch) {
53                flush(&mut latin_buf, &mut tokens);
54                tokens.push(ch.to_string());
55            } else {
56                latin_buf.push(ch);
57            }
58        }
59        flush(&mut latin_buf, &mut tokens);
60        tokens
61    }
62}
63
64// ---------------------------------------------------------------------------
65// KeywordTokenizer
66// ---------------------------------------------------------------------------
67
68/// Returns the entire (whitespace-trimmed) input as a single token.
69/// Empty input → empty vec.
70#[derive(Debug, Default, Clone)]
71pub struct KeywordTokenizer;
72
73impl Tokenizer for KeywordTokenizer {
74    fn tokenize(&self, text: &str) -> Vec<String> {
75        let t = text.trim();
76        if t.is_empty() {
77            vec![]
78        } else {
79            vec![t.to_string()]
80        }
81    }
82}
83
84// ---------------------------------------------------------------------------
85// IdentifierTokenizer
86// ---------------------------------------------------------------------------
87
88/// Identifier-aware tokenizer: emits lowercased original + split parts for identifiers,
89/// falls back to `WhitespaceTokenizer` for plain words.
90#[derive(Debug, Clone)]
91pub struct IdentifierTokenizer {
92    pub min_part_len: usize,
93}
94
95impl Default for IdentifierTokenizer {
96    fn default() -> Self {
97        Self { min_part_len: 1 }
98    }
99}
100
101impl Tokenizer for IdentifierTokenizer {
102    fn tokenize(&self, text: &str) -> Vec<String> {
103        let ws = WhitespaceTokenizer;
104        let mut result: Vec<String> = Vec::new();
105
106        for raw_word in text.split_whitespace() {
107            let word = raw_word.trim_matches(|c: char| c.is_ascii_punctuation());
108            if word.is_empty() {
109                continue;
110            }
111            if is_identifier(word) {
112                let lower = word.to_lowercase();
113                let parts = split_identifier(word, self.min_part_len);
114                result.push(lower.clone());
115                for part in parts {
116                    if part != lower {
117                        result.push(part);
118                    }
119                }
120            } else {
121                result.extend(ws.tokenize(word));
122            }
123        }
124        result
125    }
126}
127
128// ---------------------------------------------------------------------------
129// UnicodeWordTokenizer (feature = "unicode")
130// ---------------------------------------------------------------------------
131
132/// Splits on Unicode word boundaries using `unicode_segmentation`.
133///
134/// Enable with `features = ["unicode"]`.
135#[cfg(feature = "unicode")]
136#[derive(Debug, Default, Clone)]
137pub struct UnicodeWordTokenizer;
138
139#[cfg(feature = "unicode")]
140impl Tokenizer for UnicodeWordTokenizer {
141    fn tokenize(&self, text: &str) -> Vec<String> {
142        use unicode_segmentation::UnicodeSegmentation;
143        UnicodeSegmentation::unicode_words(text)
144            .map(|w| w.to_string())
145            .collect()
146    }
147}
148
149// ---------------------------------------------------------------------------
150// Tests
151// ---------------------------------------------------------------------------
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156
157    // --- WhitespaceTokenizer ---
158
159    #[test]
160    fn whitespace_normal() {
161        let t = WhitespaceTokenizer;
162        assert_eq!(t.tokenize("hello world"), vec!["hello", "world"]);
163    }
164
165    #[test]
166    fn whitespace_empty() {
167        assert!(WhitespaceTokenizer.tokenize("").is_empty());
168        assert!(WhitespaceTokenizer.tokenize("   ").is_empty());
169    }
170
171    #[test]
172    fn whitespace_strips_outer_punct() {
173        let t = WhitespaceTokenizer;
174        assert_eq!(t.tokenize("hello, world!"), vec!["hello", "world"]);
175        assert_eq!(t.tokenize("...foo..."), vec!["foo"]);
176    }
177
178    #[test]
179    fn whitespace_preserves_inner_punct() {
180        let t = WhitespaceTokenizer;
181        assert_eq!(t.tokenize("fine-tuning v1.2"), vec!["fine-tuning", "v1.2"]);
182    }
183
184    #[test]
185    fn whitespace_all_punct_dropped() {
186        assert!(WhitespaceTokenizer.tokenize("... , ...").is_empty());
187    }
188
189    // --- CjkCharTokenizer ---
190
191    #[test]
192    fn cjk_spec_example() {
193        let got = CjkCharTokenizer.tokenize("使用LoRA进行fine-tuning");
194        assert_eq!(got, vec!["使", "用", "LoRA", "进", "行", "fine-tuning"]);
195    }
196
197    #[test]
198    fn cjk_empty() {
199        assert!(CjkCharTokenizer.tokenize("").is_empty());
200    }
201
202    #[test]
203    fn cjk_only_cjk() {
204        assert_eq!(CjkCharTokenizer.tokenize("中文"), vec!["中", "文"]);
205    }
206
207    #[test]
208    fn cjk_only_latin() {
209        assert_eq!(
210            CjkCharTokenizer.tokenize("hello world"),
211            vec!["hello", "world"]
212        );
213    }
214
215    #[test]
216    fn cjk_punct_stripped_from_latin_run() {
217        let got = CjkCharTokenizer.tokenize("你好, world!");
218        assert_eq!(got, vec!["你", "好", "world"]);
219    }
220
221    #[test]
222    fn cjk_hangul_unigrams() {
223        assert_eq!(CjkCharTokenizer.tokenize("가나"), vec!["가", "나"]);
224    }
225
226    // --- KeywordTokenizer ---
227
228    #[test]
229    fn keyword_whole_phrase() {
230        assert_eq!(
231            KeywordTokenizer.tokenize("attention is all you need"),
232            vec!["attention is all you need"]
233        );
234    }
235
236    #[test]
237    fn keyword_empty() {
238        assert!(KeywordTokenizer.tokenize("").is_empty());
239        assert!(KeywordTokenizer.tokenize("   ").is_empty());
240    }
241
242    #[test]
243    fn keyword_trims_outer_whitespace() {
244        assert_eq!(KeywordTokenizer.tokenize("  hello  "), vec!["hello"]);
245    }
246
247    #[test]
248    fn keyword_single_word() {
249        assert_eq!(KeywordTokenizer.tokenize("LoRA"), vec!["LoRA"]);
250    }
251
252    // --- IdentifierTokenizer ---
253
254    #[test]
255    fn identifier_spec_example() {
256        let t = IdentifierTokenizer::default();
257        // "LoRA" is_identifier → lora + lo + ra; "attention" is plain → pass-through
258        let got = t.tokenize("LoRA attention");
259        assert_eq!(got, vec!["lora", "lo", "ra", "attention"]);
260    }
261
262    #[test]
263    fn identifier_empty() {
264        assert!(IdentifierTokenizer::default().tokenize("").is_empty());
265    }
266
267    #[test]
268    fn identifier_plain_passthrough() {
269        assert_eq!(
270            IdentifierTokenizer::default().tokenize("attention"),
271            vec!["attention"]
272        );
273    }
274
275    #[test]
276    fn identifier_snake_case() {
277        let got = IdentifierTokenizer::default().tokenize("fine_tuning");
278        assert_eq!(got, vec!["fine_tuning", "fine", "tuning"]);
279    }
280
281    #[test]
282    fn identifier_kebab() {
283        let got = IdentifierTokenizer::default().tokenize("fine-tuning");
284        assert_eq!(got, vec!["fine-tuning", "fine", "tuning"]);
285    }
286
287    #[test]
288    fn identifier_min_part_len_filters_short() {
289        let t = IdentifierTokenizer { min_part_len: 3 };
290        // "LoRA" parts: "lo" (2), "ra" (2) — both below min=3
291        assert_eq!(t.tokenize("LoRA"), vec!["lora"]);
292    }
293
294    #[test]
295    fn identifier_punct_stripped() {
296        let got = IdentifierTokenizer::default().tokenize("LoRA,");
297        assert_eq!(got, vec!["lora", "lo", "ra"]);
298    }
299
300    // --- UnicodeWordTokenizer ---
301
302    #[cfg(feature = "unicode")]
303    #[test]
304    fn unicode_word_basic() {
305        let got = UnicodeWordTokenizer.tokenize("Hello, world!");
306        assert_eq!(got, vec!["Hello", "world"]);
307    }
308
309    #[cfg(feature = "unicode")]
310    #[test]
311    fn unicode_word_empty() {
312        assert!(UnicodeWordTokenizer.tokenize("").is_empty());
313    }
314}
khive_text/tokenizer.rs

khive_text/
tokenizer.rs
