context_forge/analysis/
tokenizer.rs1use std::collections::HashSet;
2
3use stop_words::{get, LANGUAGE};
4
5#[derive(Debug, Default, Clone)]
7pub struct TokenizerConfig {
8 pub extra_stopwords: Vec<String>,
10 pub custom_only: bool,
12}
13
14#[derive(Debug, Clone)]
16pub struct Tokenizer {
17 stopwords: HashSet<String>,
18}
19
20impl Tokenizer {
21 #[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 #[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 fn is_normalized_stopword(&self, term: &str) -> bool {
72 self.stopwords.contains(term)
73 }
74
75 #[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}