Skip to main content

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}