uqa_analysis/
tokenizer.rs1use std::sync::OnceLock;
13
14use regex::Regex;
15use serde::{Deserialize, Serialize};
16
17use crate::error::{AnalysisError, AnalysisResult};
18
19#[derive(Debug, Clone, Serialize, Deserialize)]
20#[serde(tag = "type", rename_all = "snake_case")]
21pub enum Tokenizer {
22 Whitespace,
23 Standard,
24 Letter,
25 NGram { min_gram: usize, max_gram: usize },
26 Pattern { pattern: String },
27 Keyword,
28}
29
30impl Tokenizer {
31 pub fn validate(&self) -> AnalysisResult<()> {
35 match self {
36 Tokenizer::NGram { min_gram, max_gram } => {
37 validate_gram_bounds("n-gram tokenizer", *min_gram, *max_gram)
38 }
39 Tokenizer::Pattern { pattern } => {
40 Regex::new(pattern)
41 .map(|_| ())
42 .map_err(|source| AnalysisError::InvalidRegex {
43 component: "pattern tokenizer",
44 pattern: pattern.clone(),
45 source,
46 })
47 }
48 _ => Ok(()),
49 }
50 }
51
52 pub fn tokenize(&self, text: &str) -> AnalysisResult<Vec<String>> {
53 let tokens = match self {
54 Tokenizer::Whitespace => text.split_whitespace().map(str::to_owned).collect(),
55 Tokenizer::Standard => standard_word_re()?
56 .find_iter(text)
57 .map(|m| m.as_str().to_owned())
58 .collect(),
59 Tokenizer::Letter => letter_re()?
60 .find_iter(text)
61 .map(|m| m.as_str().to_owned())
62 .collect(),
63 Tokenizer::NGram { min_gram, max_gram } => {
64 validate_gram_bounds("n-gram tokenizer", *min_gram, *max_gram)?;
65 let mut out = Vec::new();
66 for word in text.split_whitespace() {
67 let chars: Vec<char> = word.chars().collect();
68 for n in *min_gram..=*max_gram {
69 if chars.len() < n {
70 continue;
71 }
72 for i in 0..=(chars.len() - n) {
73 out.push(chars[i..i + n].iter().collect());
74 }
75 }
76 }
77 out
78 }
79 Tokenizer::Pattern { pattern } => {
80 let re = Regex::new(pattern).map_err(|source| AnalysisError::InvalidRegex {
81 component: "pattern tokenizer",
82 pattern: pattern.clone(),
83 source,
84 })?;
85 re.split(text)
86 .filter(|s| !s.is_empty())
87 .map(str::to_owned)
88 .collect()
89 }
90 Tokenizer::Keyword => {
91 if text.is_empty() {
92 Vec::new()
93 } else {
94 vec![text.to_owned()]
95 }
96 }
97 };
98 Ok(tokens)
99 }
100}
101
102fn validate_gram_bounds(
103 component: &'static str,
104 min_gram: usize,
105 max_gram: usize,
106) -> AnalysisResult<()> {
107 if min_gram == 0 || max_gram < min_gram {
108 return Err(AnalysisError::InvalidGramBounds {
109 component,
110 min_gram,
111 max_gram,
112 });
113 }
114 Ok(())
115}
116
117fn standard_word_re() -> AnalysisResult<&'static Regex> {
118 static RE: OnceLock<Result<Regex, String>> = OnceLock::new();
119 RE.get_or_init(|| Regex::new(r"\w+").map_err(|error| error.to_string()))
120 .as_ref()
121 .map_err(|message| AnalysisError::BuiltInRegex {
122 component: "standard tokenizer",
123 message: message.clone(),
124 })
125}
126
127fn letter_re() -> AnalysisResult<&'static Regex> {
128 static RE: OnceLock<Result<Regex, String>> = OnceLock::new();
129 RE.get_or_init(|| Regex::new(r"[a-zA-Z]+").map_err(|error| error.to_string()))
130 .as_ref()
131 .map_err(|message| AnalysisError::BuiltInRegex {
132 component: "letter tokenizer",
133 message: message.clone(),
134 })
135}
136
137#[cfg(test)]
138mod tests {
139 use super::*;
140
141 #[test]
142 fn whitespace_splits_on_whitespace() {
143 let t = Tokenizer::Whitespace;
144 assert_eq!(
145 t.tokenize("hello world\n rust").unwrap(),
146 vec!["hello", "world", "rust"]
147 );
148 }
149
150 #[test]
151 fn standard_extracts_unicode_words() {
152 let t = Tokenizer::Standard;
153 assert_eq!(
154 t.tokenize("Rust 2024! Carácter.").unwrap(),
155 vec!["Rust", "2024", "Carácter"]
156 );
157 }
158
159 #[test]
160 fn letter_extracts_ascii_letters_only() {
161 let t = Tokenizer::Letter;
162 assert_eq!(t.tokenize("abc123 xyz").unwrap(), vec!["abc", "xyz"]);
163 }
164
165 #[test]
166 fn ngram_emits_substrings_per_word() {
167 let t = Tokenizer::NGram {
168 min_gram: 2,
169 max_gram: 3,
170 };
171 assert_eq!(t.tokenize("ab abc").unwrap(), vec!["ab", "ab", "bc", "abc"]);
174 }
175
176 #[test]
177 fn pattern_splits_on_regex() {
178 let t = Tokenizer::Pattern {
179 pattern: r"\W+".to_string(),
180 };
181 assert_eq!(t.tokenize("hello, world!").unwrap(), vec!["hello", "world"]);
182 }
183
184 #[test]
185 fn keyword_emits_whole_input() {
186 let t = Tokenizer::Keyword;
187 assert_eq!(t.tokenize("a b c").unwrap(), vec!["a b c"]);
188 assert!(t.tokenize("").unwrap().is_empty());
189 }
190
191 #[test]
192 fn round_trips_via_serde_json() {
193 let t = Tokenizer::NGram {
194 min_gram: 2,
195 max_gram: 4,
196 };
197 let s = serde_json::to_string(&t).unwrap();
198 let back: Tokenizer = serde_json::from_str(&s).unwrap();
199 assert_eq!(
200 back.tokenize("foobar").unwrap(),
201 t.tokenize("foobar").unwrap()
202 );
203 }
204}