Skip to main content

context_forge/analysis/
tokenizer.rs

1use std::collections::HashSet;
2
3use stop_words::{get, LANGUAGE};
4
5/// Configuration for the tokenizer.
6#[derive(Debug, Default, Clone)]
7pub struct TokenizerConfig {
8    /// Custom stopwords to add to the default set.
9    pub extra_stopwords: Vec<String>,
10    /// If true, use only `extra_stopwords` and skip the default English set.
11    pub custom_only: bool,
12}
13
14/// A configured tokenizer that normalizes and filters text into terms.
15#[derive(Debug, Clone)]
16pub struct Tokenizer {
17    stopwords: HashSet<String>,
18}
19
20impl Tokenizer {
21    /// Create a new tokenizer with the given configuration.
22    ///
23    /// Loads the default English stopword list from `stop-words` crate
24    /// unless `config.custom_only` is true. Merges `extra_stopwords`.
25    #[must_use]
26    pub fn new(config: &TokenizerConfig) -> Self {
27        let mut stopwords: HashSet<String> = if config.custom_only {
28            HashSet::new()
29        } else {
30            get(LANGUAGE::English)
31                .iter()
32                .map(|word| word.to_ascii_lowercase())
33                .collect()
34        };
35
36        stopwords.extend(
37            config
38                .extra_stopwords
39                .iter()
40                .map(|word| word.to_ascii_lowercase()),
41        );
42
43        Self { stopwords }
44    }
45
46    /// Tokenize text into a list of normalized terms.
47    ///
48    /// 1. Lowercase the entire input (ASCII only)
49    /// 2. Split on ASCII whitespace
50    /// 3. Strip non-alphanumeric characters from each token (keep only a-z, 0-9)
51    /// 4. Remove empty tokens
52    /// 5. Remove stopwords
53    ///
54    /// **Note:** Input is treated as ASCII; non-ASCII characters are silently dropped.
55    #[must_use]
56    pub fn tokenize(&self, text: &str) -> Vec<String> {
57        text.to_ascii_lowercase()
58            .split_ascii_whitespace()
59            .map(|token| {
60                token
61                    .chars()
62                    .filter(char::is_ascii_alphanumeric)
63                    .collect::<String>()
64            })
65            .filter(|token| !token.is_empty())
66            .filter(|token| !self.is_normalized_stopword(token))
67            .collect()
68    }
69
70    /// Return true if the given normalized term is a stopword.
71    fn is_normalized_stopword(&self, term: &str) -> bool {
72        self.stopwords.contains(term)
73    }
74
75    /// Return true if the given term is a stopword.
76    #[must_use]
77    pub fn is_stopword(&self, term: &str) -> bool {
78        self.is_normalized_stopword(&term.to_ascii_lowercase())
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use super::{Tokenizer, TokenizerConfig};
85
86    #[test]
87    fn test_basic_tokenization() {
88        let config = TokenizerConfig {
89            extra_stopwords: Vec::new(),
90            custom_only: true,
91        };
92        let tokenizer = Tokenizer::new(&config);
93        let tokens = tokenizer.tokenize("Hello World");
94
95        assert_eq!(tokens, vec!["hello".to_string(), "world".to_string()]);
96    }
97
98    #[test]
99    fn test_stopword_removal() {
100        let config = TokenizerConfig {
101            extra_stopwords: vec!["the".to_string(), "is".to_string()],
102            custom_only: true,
103        };
104        let tokenizer = Tokenizer::new(&config);
105        let tokens = tokenizer.tokenize("the quick brown fox is running");
106
107        assert_eq!(
108            tokens,
109            vec![
110                "quick".to_string(),
111                "brown".to_string(),
112                "fox".to_string(),
113                "running".to_string()
114            ]
115        );
116    }
117
118    #[test]
119    fn test_punctuation_stripping() {
120        let config = TokenizerConfig {
121            extra_stopwords: Vec::new(),
122            custom_only: true,
123        };
124        let tokenizer = Tokenizer::new(&config);
125        let tokens = tokenizer.tokenize("don't, can't; hello!");
126
127        assert_eq!(
128            tokens,
129            vec!["dont".to_string(), "cant".to_string(), "hello".to_string()]
130        );
131    }
132
133    #[test]
134    fn test_empty_input() {
135        let tokenizer = Tokenizer::new(&TokenizerConfig::default());
136        let tokens = tokenizer.tokenize("");
137
138        assert!(tokens.is_empty());
139    }
140
141    #[test]
142    fn test_all_stopwords() {
143        let tokenizer = Tokenizer::new(&TokenizerConfig::default());
144        let tokens = tokenizer.tokenize("the is a an");
145
146        assert!(tokens.is_empty());
147    }
148
149    #[test]
150    fn test_custom_stopwords() {
151        let config = TokenizerConfig {
152            extra_stopwords: vec!["rust".to_string(), "tokio".to_string()],
153            custom_only: true,
154        };
155        let tokenizer = Tokenizer::new(&config);
156        let tokens = tokenizer.tokenize("Rust and Tokio are useful");
157
158        assert_eq!(
159            tokens,
160            vec!["and".to_string(), "are".to_string(), "useful".to_string()]
161        );
162    }
163
164    #[test]
165    fn test_is_stopword() {
166        let config = TokenizerConfig {
167            extra_stopwords: vec!["custom".to_string()],
168            custom_only: false,
169        };
170        let tokenizer = Tokenizer::new(&config);
171
172        assert!(tokenizer.is_stopword("the"));
173        assert!(tokenizer.is_stopword("custom"));
174        assert!(!tokenizer.is_stopword("contextforge"));
175    }
176}