plugmem_core/tokenizer/mod.rs
1//! The core tokenizer, v2.
2//!
3//! The tokenizer is deliberately split into four stages:
4//!
5//! 1. NFKC normalization of the input;
6//! 2. UAX #29 word segmentation and CJK adjacency tracking;
7//! 3. policy-driven Unicode folding;
8//! 4. canonical, byte-budgeted token emission.
9//!
10//! The stages share caller-owned scratch buffers, so the ordinary Latin,
11//! Cyrillic, and generic-script paths remain allocation-free after warm-up and
12//! behave identically on native and WASM. ICU4X's dictionary/LSTM path for
13//! complex scripts may allocate inside its iterator, but uses the same token
14//! policy and canonical emission rules.
15//! Emitted tokens are canonical fixed points of the tokenizer.
16
17use alloc::string::String;
18
19mod emit;
20mod fold;
21mod normalize;
22mod policy;
23mod segment;
24mod tables;
25mod unicode;
26
27pub use self::emit::MAX_TOKEN_BYTES;
28pub use self::policy::TokenizerPolicy;
29use self::segment::CjkRun;
30use self::unicode::UnicodeBackend;
31
32/// Streaming tokenizer with reusable scratch buffers.
33///
34/// One instance should be reused by an engine or by one thread of a wrapper.
35/// After warm-up, [`Tokenizer::tokenize`] performs no heap allocation on the
36/// generic Unicode path. Complex scripts may use ICU4X dictionary/LSTM
37/// scratch allocations to obtain language-aware word boundaries.
38#[derive(Debug, Default, Clone)]
39pub struct Tokenizer {
40 /// Folding policy. This is copied into the hot loop and has no heap cost.
41 policy: TokenizerPolicy,
42 /// NFKC-normalized copy of the input.
43 norm: String,
44 /// The token being assembled (folded word or CJK bigram).
45 token: String,
46 /// Lowercase copy of the current ICU word segment.
47 lower: String,
48 /// Reused scratch for the rare post-fold NFKC pass.
49 canonical: String,
50 /// Compiled Unicode data and segmentation rules.
51 unicode: UnicodeBackend,
52}
53
54impl Tokenizer {
55 /// Creates a tokenizer with empty scratch buffers and the search policy.
56 pub fn new() -> Self {
57 Self::default()
58 }
59
60 /// Creates a tokenizer with an explicit folding policy.
61 pub fn with_policy(policy: TokenizerPolicy) -> Self {
62 Self {
63 policy,
64 ..Self::default()
65 }
66 }
67
68 /// Returns the policy used by this tokenizer.
69 pub const fn policy(&self) -> TokenizerPolicy {
70 self.policy
71 }
72
73 /// Splits `text` into normalized tokens, calling `sink` for each one.
74 ///
75 /// The emitted `&str` is only valid for the duration of one `sink` call.
76 ///
77 /// ```
78 /// use plugmem_core::tokenizer::Tokenizer;
79 ///
80 /// let mut tokenizer = Tokenizer::new();
81 /// let mut tokens = Vec::new();
82 /// tokenizer.tokenize("Hello, МИР-42! 東京タワー", &mut |token| {
83 /// tokens.push(token.to_owned())
84 /// });
85 /// assert_eq!(tokens, ["hello", "мир", "42", "東京", "タワー"]);
86 /// ```
87 pub fn tokenize(&mut self, text: &str, sink: &mut dyn FnMut(&str)) {
88 normalize::normalize_into(&self.unicode, text, &mut self.norm);
89
90 let policy = self.policy;
91 let token = &mut self.token;
92 let lower = &mut self.lower;
93 let canonical = &mut self.canonical;
94 let unicode = &self.unicode;
95 let mut cjk_run = CjkRun::default();
96 let mut processor =
97 segment::SegmentProcessor::new(policy, unicode, &mut cjk_run, token, lower, canonical);
98 let mut start = 0usize;
99
100 for end in unicode.word_boundaries(&self.norm) {
101 if end == start {
102 continue;
103 }
104 let segment = &self.norm[start..end];
105 if segment.chars().any(char::is_alphanumeric) {
106 processor.process(segment, sink);
107 } else {
108 processor.flush(sink);
109 }
110 start = end;
111 }
112 processor.flush(sink);
113 }
114}