uqa_analysis/
tokenizer.rs1use serde::{Deserialize, Serialize};
13use std::sync::OnceLock;
14use uqa_core::memory::{Budgeted, MemoryBudget};
15
16use crate::character_class::class;
17use crate::error::{AnalysisError, AnalysisResult};
18
19mod compiled;
20mod stream;
21pub(crate) use compiled::PreparedTokenizer;
22
23#[derive(Debug, Clone, Serialize, Deserialize)]
24#[serde(tag = "type", rename_all = "snake_case")]
25pub enum Tokenizer {
26 Whitespace,
27 Standard,
28 Letter,
29 NGram {
30 min_gram: usize,
31 max_gram: usize,
32 },
33 Pattern {
34 pattern: String,
35 },
36 Keyword,
37 #[cfg(feature = "nori")]
38 #[serde(rename = "nori_tokenizer")]
39 Nori(crate::nori::NoriTokenizerConfig),
40 #[cfg(feature = "kuromoji")]
41 #[serde(rename = "kuromoji_tokenizer")]
42 Kuromoji(crate::kuromoji::KuromojiTokenizerConfig),
43}
44
45impl Tokenizer {
46 pub fn validate(&self) -> AnalysisResult<()> {
50 match self {
51 #[cfg(feature = "nori")]
52 Tokenizer::Nori(_) => self.prepare().map(|_| ()),
53 #[cfg(feature = "kuromoji")]
54 Tokenizer::Kuromoji(_) => self.prepare().map(|_| ()),
55 Tokenizer::NGram { .. } | Tokenizer::Pattern { .. } => self.prepare().map(|_| ()),
56 _ => Ok(()),
57 }
58 }
59
60 pub fn tokenize(&self, text: &str) -> AnalysisResult<Vec<String>> {
61 self.tokenize_with_offsets(text)?.into_terms()
62 }
63
64 pub fn tokenize_with_offsets(&self, text: &str) -> AnalysisResult<crate::AnalyzedText> {
66 self.tokenize_mapped(&crate::FilteredText::new(text))
67 }
68
69 pub fn tokenize_with_offsets_budgeted(
73 &self,
74 text: &str,
75 budget: &MemoryBudget,
76 poll: impl FnMut() -> AnalysisResult<()>,
77 ) -> AnalysisResult<Budgeted<crate::AnalyzedText>> {
78 self.tokenize_mapped_budgeted(&crate::FilteredText::new(text), budget, poll)
79 }
80
81 pub fn tokenize_mapped_budgeted(
104 &self,
105 text: &crate::FilteredText<'_>,
106 budget: &MemoryBudget,
107 mut poll: impl FnMut() -> AnalysisResult<()>,
108 ) -> AnalysisResult<Budgeted<crate::AnalyzedText>> {
109 poll()?;
110 self.prepare()?
111 .tokenize_mapped_budgeted(text, budget, &mut poll)
112 }
113
114 pub(crate) fn tokenize_mapped(
115 &self,
116 text: &crate::FilteredText<'_>,
117 ) -> AnalysisResult<crate::AnalyzedText> {
118 self.prepare()?.tokenize_mapped(text)
119 }
120}
121
122fn validate_gram_bounds(
123 component: &'static str,
124 min_gram: usize,
125 max_gram: usize,
126) -> AnalysisResult<()> {
127 if min_gram == 0 || max_gram < min_gram {
128 return Err(AnalysisError::InvalidGramBounds {
129 component,
130 min_gram,
131 max_gram,
132 });
133 }
134 Ok(())
135}
136
137pub(super) fn standard_word_class() -> AnalysisResult<&'static regex_syntax::hir::ClassUnicode> {
138 static CLASS: OnceLock<Result<regex_syntax::hir::ClassUnicode, String>> = OnceLock::new();
139 CLASS
140 .get_or_init(|| class(r"\w"))
141 .as_ref()
142 .map_err(|message| AnalysisError::BuiltInRegex {
143 component: "standard tokenizer",
144 message: message.clone(),
145 })
146}
147
148#[cfg(test)]
149mod tests {
150 use super::*;
151
152 #[test]
153 fn whitespace_splits_on_whitespace() {
154 let t = Tokenizer::Whitespace;
155 assert_eq!(
156 t.tokenize("hello world\n rust").unwrap(),
157 vec!["hello", "world", "rust"]
158 );
159 }
160
161 #[test]
162 fn standard_extracts_unicode_words() {
163 let t = Tokenizer::Standard;
164 assert_eq!(
165 t.tokenize("Rust 2024! Carácter.").unwrap(),
166 vec!["Rust", "2024", "Carácter"]
167 );
168 }
169
170 #[test]
171 fn letter_extracts_ascii_letters_only() {
172 let t = Tokenizer::Letter;
173 assert_eq!(t.tokenize("abc123 xyz").unwrap(), vec!["abc", "xyz"]);
174 }
175
176 #[test]
177 fn ngram_emits_substrings_per_word() {
178 let t = Tokenizer::NGram {
179 min_gram: 2,
180 max_gram: 3,
181 };
182 assert_eq!(t.tokenize("ab abc").unwrap(), vec!["ab", "ab", "bc", "abc"]);
185 }
186
187 #[test]
188 fn pattern_splits_on_regex() {
189 let t = Tokenizer::Pattern {
190 pattern: r"\W+".to_string(),
191 };
192 assert_eq!(t.tokenize("hello, world!").unwrap(), vec!["hello", "world"]);
193 }
194
195 #[test]
196 fn keyword_emits_whole_input() {
197 let t = Tokenizer::Keyword;
198 assert_eq!(t.tokenize("a b c").unwrap(), vec!["a b c"]);
199 assert!(t.tokenize("").unwrap().is_empty());
200 }
201
202 #[test]
203 fn round_trips_via_serde_json() {
204 let t = Tokenizer::NGram {
205 min_gram: 2,
206 max_gram: 4,
207 };
208 let s = serde_json::to_string(&t).unwrap();
209 let back: Tokenizer = serde_json::from_str(&s).unwrap();
210 assert_eq!(
211 back.tokenize("foobar").unwrap(),
212 t.tokenize("foobar").unwrap()
213 );
214 }
215}