1#[cfg(any(feature = "native", feature = "wasm"))]
4mod hf_tokenizer;
5
6#[cfg(feature = "native")]
7mod idf_weights;
8
9#[cfg(any(feature = "native", feature = "wasm"))]
10pub use hf_tokenizer::{HfTokenizer, TokenizerSource};
11
12#[cfg(feature = "native")]
13pub use hf_tokenizer::{TokenizerCache, tokenizer_cache};
14
15#[cfg(feature = "native")]
16pub use idf_weights::{IdfWeights, IdfWeightsCache, idf_weights_cache};
17
18use std::collections::HashMap;
19use std::sync::Arc;
20
21use parking_lot::RwLock;
22use rust_stemmers::Algorithm;
23use serde::{Deserialize, Serialize};
24use stop_words::LANGUAGE;
25
26#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
28pub struct Token {
29 pub text: String,
31 pub position: u32,
33 pub offset_from: usize,
35 pub offset_to: usize,
37}
38
39impl Token {
40 pub fn new(text: String, position: u32, offset_from: usize, offset_to: usize) -> Self {
41 Self {
42 text,
43 position,
44 offset_from,
45 offset_to,
46 }
47 }
48}
49
50pub trait Tokenizer: Send + Sync + Clone + 'static {
52 fn tokenize(&self, text: &str) -> Vec<Token>;
54
55 fn tokenize_hinted(&self, text: &str, hint: Option<&str>) -> Vec<Token> {
62 let _ = hint;
63 self.tokenize(text)
64 }
65}
66
67#[derive(Debug, Clone, Default)]
71pub struct SimpleTokenizer;
72
73impl Tokenizer for SimpleTokenizer {
74 fn tokenize(&self, text: &str) -> Vec<Token> {
75 tokenize_and_clean(text, std::convert::identity)
76 }
77}
78
79#[derive(Debug, Clone, Default)]
83pub struct RawTokenizer;
84
85impl Tokenizer for RawTokenizer {
86 fn tokenize(&self, text: &str) -> Vec<Token> {
87 let trimmed = text.trim();
88 if trimmed.is_empty() {
89 return Vec::new();
90 }
91 let offset = text.as_ptr() as usize;
92 let trimmed_offset = trimmed.as_ptr() as usize - offset;
93 vec![Token::new(
94 trimmed.to_string(),
95 0,
96 trimmed_offset,
97 trimmed_offset + trimmed.len(),
98 )]
99 }
100}
101
102#[derive(Debug, Clone, Default)]
106pub struct RawCiTokenizer;
107
108impl Tokenizer for RawCiTokenizer {
109 fn tokenize(&self, text: &str) -> Vec<Token> {
110 let trimmed = text.trim();
111 if trimmed.is_empty() {
112 return Vec::new();
113 }
114 let offset = text.as_ptr() as usize;
115 let trimmed_offset = trimmed.as_ptr() as usize - offset;
116 vec![Token::new(
117 lowercase_word(trimmed),
118 0,
119 trimmed_offset,
120 trimmed_offset + trimmed.len(),
121 )]
122 }
123}
124
125#[inline]
129fn lowercase_word(word: &str) -> String {
130 if word.is_ascii() {
131 if word.bytes().all(|b| !b.is_ascii_uppercase()) {
132 return word.to_string();
133 }
134 let mut s = word.to_string();
135 s.make_ascii_lowercase();
136 s
137 } else {
138 word.chars().flat_map(|c| c.to_lowercase()).collect()
139 }
140}
141
142#[inline]
147fn clean_word(word: &str) -> String {
148 if word.is_ascii() {
149 let bytes = word.as_bytes();
150 if bytes
152 .iter()
153 .all(|&b| b.is_ascii_lowercase() || b.is_ascii_digit())
154 {
155 return word.to_string();
156 }
157 let mut result = String::with_capacity(bytes.len());
159 for &b in bytes {
160 if b.is_ascii_alphanumeric() {
161 result.push(b.to_ascii_lowercase() as char);
162 }
163 }
164 result
165 } else {
166 word.chars()
168 .filter(|c| c.is_alphanumeric())
169 .flat_map(|c| c.to_lowercase())
170 .collect()
171 }
172}
173
174fn tokenize_and_clean(text: &str, transform: impl Fn(String) -> String) -> Vec<Token> {
180 let mut tokens = Vec::with_capacity(text.len() / 5);
181 let mut position = 0u32;
182 for (offset, word) in split_whitespace_with_offsets(text) {
183 if !word.is_empty() {
184 let cleaned = clean_word(word);
185 if !cleaned.is_empty() {
186 tokens.push(Token::new(
187 transform(cleaned),
188 position,
189 offset,
190 offset + word.len(),
191 ));
192 position += 1;
193 }
194 }
195 }
196 tokens
197}
198
199fn split_whitespace_with_offsets(text: &str) -> impl Iterator<Item = (usize, &str)> {
204 let base = text.as_ptr() as usize;
205 text.split_whitespace()
206 .map(move |word| (word.as_ptr() as usize - base, word))
207}
208
209#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
211#[allow(missing_docs)]
212#[derive(Default)]
213pub enum Language {
214 Arabic,
215 Danish,
216 Dutch,
217 #[default]
218 English,
219 Finnish,
220 French,
221 German,
222 Greek,
223 Hungarian,
224 Italian,
225 Norwegian,
226 Portuguese,
227 Romanian,
228 Russian,
229 Spanish,
230 Swedish,
231 Tamil,
232 Turkish,
233}
234
235impl Language {
236 fn to_algorithm(self) -> Algorithm {
237 match self {
238 Language::Arabic => Algorithm::Arabic,
239 Language::Danish => Algorithm::Danish,
240 Language::Dutch => Algorithm::Dutch,
241 Language::English => Algorithm::English,
242 Language::Finnish => Algorithm::Finnish,
243 Language::French => Algorithm::French,
244 Language::German => Algorithm::German,
245 Language::Greek => Algorithm::Greek,
246 Language::Hungarian => Algorithm::Hungarian,
247 Language::Italian => Algorithm::Italian,
248 Language::Norwegian => Algorithm::Norwegian,
249 Language::Portuguese => Algorithm::Portuguese,
250 Language::Romanian => Algorithm::Romanian,
251 Language::Russian => Algorithm::Russian,
252 Language::Spanish => Algorithm::Spanish,
253 Language::Swedish => Algorithm::Swedish,
254 Language::Tamil => Algorithm::Tamil,
255 Language::Turkish => Algorithm::Turkish,
256 }
257 }
258
259 fn to_stop_words_language(self) -> LANGUAGE {
260 match self {
261 Language::Arabic => LANGUAGE::Arabic,
262 Language::Danish => LANGUAGE::Danish,
263 Language::Dutch => LANGUAGE::Dutch,
264 Language::English => LANGUAGE::English,
265 Language::Finnish => LANGUAGE::Finnish,
266 Language::French => LANGUAGE::French,
267 Language::German => LANGUAGE::German,
268 Language::Greek => LANGUAGE::Greek,
269 Language::Hungarian => LANGUAGE::Hungarian,
270 Language::Italian => LANGUAGE::Italian,
271 Language::Norwegian => LANGUAGE::Norwegian,
272 Language::Portuguese => LANGUAGE::Portuguese,
273 Language::Romanian => LANGUAGE::Romanian,
274 Language::Russian => LANGUAGE::Russian,
275 Language::Spanish => LANGUAGE::Spanish,
276 Language::Swedish => LANGUAGE::Swedish,
277 Language::Tamil => LANGUAGE::English, Language::Turkish => LANGUAGE::Turkish,
279 }
280 }
281}
282
283#[derive(Debug, Clone)]
287pub struct StopWordTokenizer<T: Tokenizer> {
288 inner: T,
289 stop_words: HashSet<String>,
290}
291
292use std::collections::HashSet;
293
294impl<T: Tokenizer> StopWordTokenizer<T> {
295 pub fn new(inner: T, language: Language) -> Self {
297 let stop_words: HashSet<String> = stop_words::get(language.to_stop_words_language())
298 .iter()
299 .map(|s| s.to_string())
300 .collect();
301 Self { inner, stop_words }
302 }
303
304 pub fn english(inner: T) -> Self {
306 Self::new(inner, Language::English)
307 }
308
309 pub fn with_custom_stop_words(inner: T, stop_words: HashSet<String>) -> Self {
311 Self { inner, stop_words }
312 }
313
314 pub fn is_stop_word(&self, word: &str) -> bool {
316 self.stop_words.contains(word)
317 }
318}
319
320impl<T: Tokenizer> Tokenizer for StopWordTokenizer<T> {
321 fn tokenize(&self, text: &str) -> Vec<Token> {
322 self.inner
323 .tokenize(text)
324 .into_iter()
325 .filter(|token| !self.stop_words.contains(token.text.as_str()))
326 .collect()
327 }
328}
329
330#[derive(Debug, Clone)]
335pub struct StemmerTokenizer {
336 language: Language,
337}
338
339impl StemmerTokenizer {
340 pub fn new(language: Language) -> Self {
342 Self { language }
343 }
344
345 pub fn english() -> Self {
347 Self::new(Language::English)
348 }
349}
350
351impl Default for StemmerTokenizer {
352 fn default() -> Self {
353 Self::english()
354 }
355}
356
357impl Tokenizer for StemmerTokenizer {
358 fn tokenize(&self, text: &str) -> Vec<Token> {
359 let stemmer = rust_stemmers::Stemmer::create(self.language.to_algorithm());
360 tokenize_and_clean(text, |s| stemmer.stem(&s).into_owned())
361 }
362}
363
364#[derive(Debug, Clone)]
369pub struct MultiLanguageStemmer {
370 default_language: Language,
371}
372
373impl MultiLanguageStemmer {
374 pub fn new(default_language: Language) -> Self {
376 Self { default_language }
377 }
378
379 pub fn tokenize_with_language(&self, text: &str, language: Language) -> Vec<Token> {
381 let stemmer = rust_stemmers::Stemmer::create(language.to_algorithm());
382 tokenize_and_clean(text, |s| stemmer.stem(&s).into_owned())
383 }
384
385 pub fn default_language(&self) -> Language {
387 self.default_language
388 }
389}
390
391impl Default for MultiLanguageStemmer {
392 fn default() -> Self {
393 Self::new(Language::English)
394 }
395}
396
397impl Tokenizer for MultiLanguageStemmer {
398 fn tokenize(&self, text: &str) -> Vec<Token> {
399 self.tokenize_with_language(text, self.default_language)
400 }
401}
402
403#[derive(Clone)]
408pub struct LanguageAwareTokenizer<F>
409where
410 F: Fn(&str) -> Language + Clone + Send + Sync + 'static,
411{
412 language_selector: F,
413 stemmer: MultiLanguageStemmer,
414}
415
416impl<F> LanguageAwareTokenizer<F>
417where
418 F: Fn(&str) -> Language + Clone + Send + Sync + 'static,
419{
420 pub fn new(language_selector: F) -> Self {
437 Self {
438 language_selector,
439 stemmer: MultiLanguageStemmer::default(),
440 }
441 }
442
443 pub fn tokenize_with_hint(&self, text: &str, language_hint: &str) -> Vec<Token> {
447 let language = (self.language_selector)(language_hint);
448 self.stemmer.tokenize_with_language(text, language)
449 }
450}
451
452impl<F> Tokenizer for LanguageAwareTokenizer<F>
453where
454 F: Fn(&str) -> Language + Clone + Send + Sync + 'static,
455{
456 fn tokenize(&self, text: &str) -> Vec<Token> {
457 self.stemmer.tokenize_with_language(text, Language::English)
459 }
460
461 fn tokenize_hinted(&self, text: &str, hint: Option<&str>) -> Vec<Token> {
462 match hint {
463 Some(hint) => self.tokenize_with_hint(text, hint),
464 None => Tokenizer::tokenize(self, text),
465 }
466 }
467}
468
469pub fn parse_language(s: &str) -> Language {
473 match s.to_lowercase().as_str() {
474 "ar" | "arabic" => Language::Arabic,
475 "da" | "danish" => Language::Danish,
476 "nl" | "dutch" => Language::Dutch,
477 "en" | "english" => Language::English,
478 "fi" | "finnish" => Language::Finnish,
479 "fr" | "french" => Language::French,
480 "de" | "german" => Language::German,
481 "el" | "greek" => Language::Greek,
482 "hu" | "hungarian" => Language::Hungarian,
483 "it" | "italian" => Language::Italian,
484 "no" | "norwegian" => Language::Norwegian,
485 "pt" | "portuguese" => Language::Portuguese,
486 "ro" | "romanian" => Language::Romanian,
487 "ru" | "russian" => Language::Russian,
488 "es" | "spanish" => Language::Spanish,
489 "sv" | "swedish" => Language::Swedish,
490 "ta" | "tamil" => Language::Tamil,
491 "tr" | "turkish" => Language::Turkish,
492 _ => Language::English, }
494}
495
496pub fn parse_language_opt(s: &str) -> Option<Language> {
501 Some(match s.trim().to_lowercase().as_str() {
502 "ar" | "arabic" => Language::Arabic,
503 "da" | "danish" => Language::Danish,
504 "nl" | "dutch" => Language::Dutch,
505 "en" | "english" => Language::English,
506 "fi" | "finnish" => Language::Finnish,
507 "fr" | "french" => Language::French,
508 "de" | "german" => Language::German,
509 "el" | "greek" => Language::Greek,
510 "hu" | "hungarian" => Language::Hungarian,
511 "it" | "italian" => Language::Italian,
512 "no" | "norwegian" => Language::Norwegian,
513 "pt" | "portuguese" => Language::Portuguese,
514 "ro" | "romanian" => Language::Romanian,
515 "ru" | "russian" => Language::Russian,
516 "es" | "spanish" => Language::Spanish,
517 "sv" | "swedish" => Language::Swedish,
518 "ta" | "tamil" => Language::Tamil,
519 "tr" | "turkish" => Language::Turkish,
520 _ => return None,
521 })
522}
523
524#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
526pub enum Script {
527 Latin,
528 Cyrillic,
529 Greek,
530 Arabic,
531 Tamil,
532 Other,
534}
535
536impl Script {
537 pub fn of_token(token: &str) -> Script {
539 let Some(c) = token.chars().find(|c| c.is_alphabetic()) else {
540 return Script::Other;
541 };
542 Script::of_char(c)
543 }
544
545 fn of_char(c: char) -> Script {
546 match c as u32 {
547 0x0041..=0x024F | 0x1E00..=0x1EFF => Script::Latin,
549 0x0370..=0x03FF | 0x1F00..=0x1FFF => Script::Greek,
550 0x0400..=0x052F => Script::Cyrillic,
551 0x0600..=0x06FF | 0x0750..=0x077F | 0x08A0..=0x08FF => Script::Arabic,
552 0x0B80..=0x0BFF => Script::Tamil,
553 _ => Script::Other,
554 }
555 }
556}
557
558impl Language {
559 pub fn script(self) -> Script {
561 match self {
562 Language::Russian => Script::Cyrillic,
563 Language::Greek => Script::Greek,
564 Language::Arabic => Script::Arabic,
565 Language::Tamil => Script::Tamil,
566 _ => Script::Latin,
567 }
568 }
569}
570
571thread_local! {
572 static STEMMER_CACHE: std::cell::RefCell<HashMap<Language, rust_stemmers::Stemmer>> =
573 std::cell::RefCell::new(HashMap::new());
574}
575
576#[inline]
580fn stem_owned(stemmer: &rust_stemmers::Stemmer, word: String) -> String {
581 match stemmer.stem(&word) {
582 std::borrow::Cow::Borrowed(_) => word,
583 std::borrow::Cow::Owned(stemmed) => stemmed,
584 }
585}
586
587fn with_stemmers<R>(languages: &[Language], f: impl FnOnce(&[&rust_stemmers::Stemmer]) -> R) -> R {
592 STEMMER_CACHE.with(|cache| {
593 let mut cache = cache.borrow_mut();
594 for language in languages {
595 cache
596 .entry(*language)
597 .or_insert_with(|| rust_stemmers::Stemmer::create(language.to_algorithm()));
598 }
599 let stemmers: Vec<&rust_stemmers::Stemmer> =
600 languages.iter().map(|language| &cache[language]).collect();
601 f(&stemmers)
602 })
603}
604
605#[derive(Debug, Clone, Default)]
619pub struct DynamicStemmer {
620 default: Option<Language>,
621}
622
623impl DynamicStemmer {
624 pub fn new(default: Option<Language>) -> Self {
626 Self { default }
627 }
628
629 pub fn default_language(&self) -> Option<Language> {
631 self.default
632 }
633
634 pub fn parse_hint(hint: &str) -> Vec<Language> {
636 let mut languages = Vec::new();
637 for part in hint.split(',') {
638 if let Some(language) = parse_language_opt(part)
639 && !languages.contains(&language)
640 {
641 languages.push(language);
642 }
643 }
644 languages
645 }
646
647 fn tokenize_with_languages(&self, text: &str, languages: &[Language]) -> Vec<Token> {
648 match languages {
649 [] => tokenize_and_clean(text, |s| s),
650 [single] => {
651 let script = single.script();
652 with_stemmers(languages, |stemmers| {
653 let stemmer = stemmers[0];
654 tokenize_and_clean(text, |s| {
655 if Script::of_token(&s) == script {
656 stem_owned(stemmer, s)
657 } else {
658 s
659 }
660 })
661 })
662 }
663 many => with_stemmers(many, |stemmers| {
664 tokenize_and_clean(text, |s| {
665 let script = Script::of_token(&s);
666 match many.iter().position(|language| language.script() == script) {
667 Some(index) => stem_owned(stemmers[index], s),
668 None => s,
669 }
670 })
671 }),
672 }
673 }
674}
675
676impl Tokenizer for DynamicStemmer {
677 fn tokenize(&self, text: &str) -> Vec<Token> {
678 match self.default {
679 Some(language) => self.tokenize_with_languages(text, &[language]),
680 None => tokenize_and_clean(text, |s| s),
681 }
682 }
683
684 fn tokenize_hinted(&self, text: &str, hint: Option<&str>) -> Vec<Token> {
685 match hint.map(str::trim).filter(|hint| !hint.is_empty()) {
686 Some(hint) => {
687 let languages = Self::parse_hint(hint);
688 if languages.is_empty() {
689 Tokenizer::tokenize(self, text)
691 } else {
692 self.tokenize_with_languages(text, &languages)
693 }
694 }
695 None => Tokenizer::tokenize(self, text),
696 }
697 }
698}
699
700#[derive(Debug, Clone, PartialEq, Eq)]
708pub enum TokenizerSpec {
709 Named(String),
711 DynamicStem {
713 by: String,
715 default: Option<Language>,
717 },
718}
719
720impl TokenizerSpec {
721 pub fn parse(spec: &str) -> Result<TokenizerSpec, String> {
723 let spec = spec.trim();
724 let Some(rest) = spec.strip_prefix("stem(") else {
725 if spec.is_empty() || spec.contains(['(', ')', ':', ',']) {
726 return Err(format!("invalid tokenizer spec '{spec}'"));
727 }
728 return Ok(TokenizerSpec::Named(spec.to_string()));
729 };
730 let Some(params) = rest.strip_suffix(')') else {
731 return Err(format!("tokenizer spec '{spec}' is missing ')'"));
732 };
733 let mut by = None;
734 let mut default = None;
735 for param in params.split(',') {
736 let param = param.trim();
737 if param.is_empty() {
738 continue;
739 }
740 let Some((key, value)) = param.split_once(':') else {
741 return Err(format!(
742 "tokenizer spec '{spec}': parameter '{param}' must be 'key: value'"
743 ));
744 };
745 let (key, value) = (key.trim(), value.trim());
746 match key {
747 "by" if !value.is_empty() => by = Some(value.to_string()),
748 "by" => return Err(format!("tokenizer spec '{spec}': 'by' needs a field name")),
749 "default" => {
750 default = match value {
751 "simple" | "none" => None,
752 other => Some(parse_language_opt(other).ok_or_else(|| {
753 format!("tokenizer spec '{spec}': unknown default language '{other}'")
754 })?),
755 };
756 }
757 other => {
758 return Err(format!(
759 "tokenizer spec '{spec}': unknown parameter '{other}'"
760 ));
761 }
762 }
763 }
764 let by = by.ok_or_else(|| format!("tokenizer spec '{spec}' requires 'by: <field>'"))?;
765 Ok(TokenizerSpec::DynamicStem { by, default })
766 }
767
768 pub fn hint_field(&self) -> Option<&str> {
770 match self {
771 TokenizerSpec::Named(_) => None,
772 TokenizerSpec::DynamicStem { by, .. } => Some(by),
773 }
774 }
775
776 pub fn dynamic_tokenizer(&self) -> Option<BoxedTokenizer> {
778 match self {
779 TokenizerSpec::Named(_) => None,
780 TokenizerSpec::DynamicStem { default, .. } => {
781 Some(Box::new(DynamicStemmer::new(*default)))
782 }
783 }
784 }
785}
786
787impl std::fmt::Display for TokenizerSpec {
788 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
789 match self {
790 TokenizerSpec::Named(name) => f.write_str(name),
791 TokenizerSpec::DynamicStem { by, default } => {
792 let default = match default {
793 None => "simple".to_string(),
794 Some(language) => language_code(*language).to_string(),
795 };
796 write!(f, "stem(by: {by}, default: {default})")
797 }
798 }
799 }
800}
801
802pub fn language_code(language: Language) -> &'static str {
804 match language {
805 Language::Arabic => "ar",
806 Language::Danish => "da",
807 Language::Dutch => "nl",
808 Language::English => "en",
809 Language::Finnish => "fi",
810 Language::French => "fr",
811 Language::German => "de",
812 Language::Greek => "el",
813 Language::Hungarian => "hu",
814 Language::Italian => "it",
815 Language::Norwegian => "no",
816 Language::Portuguese => "pt",
817 Language::Romanian => "ro",
818 Language::Russian => "ru",
819 Language::Spanish => "es",
820 Language::Swedish => "sv",
821 Language::Tamil => "ta",
822 Language::Turkish => "tr",
823 }
824}
825
826pub type BoxedTokenizer = Box<dyn TokenizerClone>;
828
829pub trait TokenizerClone: Send + Sync {
830 fn tokenize(&self, text: &str) -> Vec<Token>;
831 fn tokenize_hinted(&self, text: &str, hint: Option<&str>) -> Vec<Token>;
833 fn clone_box(&self) -> BoxedTokenizer;
834}
835
836impl<T: Tokenizer> TokenizerClone for T {
837 fn tokenize(&self, text: &str) -> Vec<Token> {
838 Tokenizer::tokenize(self, text)
839 }
840
841 fn tokenize_hinted(&self, text: &str, hint: Option<&str>) -> Vec<Token> {
842 Tokenizer::tokenize_hinted(self, text, hint)
843 }
844
845 fn clone_box(&self) -> BoxedTokenizer {
846 Box::new(self.clone())
847 }
848}
849
850impl Clone for BoxedTokenizer {
851 fn clone(&self) -> Self {
852 self.clone_box()
853 }
854}
855
856#[derive(Clone)]
861pub struct TokenizerRegistry {
862 tokenizers: Arc<RwLock<HashMap<String, BoxedTokenizer>>>,
863 dynamic: Arc<RwLock<HashMap<String, BoxedTokenizer>>>,
867}
868
869impl TokenizerRegistry {
870 pub fn new() -> Self {
872 let registry = Self {
873 tokenizers: Arc::new(RwLock::new(HashMap::new())),
874 dynamic: Arc::new(RwLock::new(HashMap::new())),
875 };
876 registry.register_defaults();
877 registry
878 }
879
880 fn register_defaults(&self) {
882 self.register("simple", SimpleTokenizer);
884 self.register("default", SimpleTokenizer);
885 self.register("raw", RawTokenizer);
886 self.register("raw_ci", RawCiTokenizer);
887
888 self.register("en_stem", StemmerTokenizer::new(Language::English));
890 self.register("english", StemmerTokenizer::new(Language::English));
891
892 self.register("ar_stem", StemmerTokenizer::new(Language::Arabic));
894 self.register("arabic", StemmerTokenizer::new(Language::Arabic));
895 self.register("da_stem", StemmerTokenizer::new(Language::Danish));
896 self.register("danish", StemmerTokenizer::new(Language::Danish));
897 self.register("nl_stem", StemmerTokenizer::new(Language::Dutch));
898 self.register("dutch", StemmerTokenizer::new(Language::Dutch));
899 self.register("fi_stem", StemmerTokenizer::new(Language::Finnish));
900 self.register("finnish", StemmerTokenizer::new(Language::Finnish));
901 self.register("fr_stem", StemmerTokenizer::new(Language::French));
902 self.register("french", StemmerTokenizer::new(Language::French));
903 self.register("de_stem", StemmerTokenizer::new(Language::German));
904 self.register("german", StemmerTokenizer::new(Language::German));
905 self.register("el_stem", StemmerTokenizer::new(Language::Greek));
906 self.register("greek", StemmerTokenizer::new(Language::Greek));
907 self.register("hu_stem", StemmerTokenizer::new(Language::Hungarian));
908 self.register("hungarian", StemmerTokenizer::new(Language::Hungarian));
909 self.register("it_stem", StemmerTokenizer::new(Language::Italian));
910 self.register("italian", StemmerTokenizer::new(Language::Italian));
911 self.register("no_stem", StemmerTokenizer::new(Language::Norwegian));
912 self.register("norwegian", StemmerTokenizer::new(Language::Norwegian));
913 self.register("pt_stem", StemmerTokenizer::new(Language::Portuguese));
914 self.register("portuguese", StemmerTokenizer::new(Language::Portuguese));
915 self.register("ro_stem", StemmerTokenizer::new(Language::Romanian));
916 self.register("romanian", StemmerTokenizer::new(Language::Romanian));
917 self.register("ru_stem", StemmerTokenizer::new(Language::Russian));
918 self.register("russian", StemmerTokenizer::new(Language::Russian));
919 self.register("es_stem", StemmerTokenizer::new(Language::Spanish));
920 self.register("spanish", StemmerTokenizer::new(Language::Spanish));
921 self.register("sv_stem", StemmerTokenizer::new(Language::Swedish));
922 self.register("swedish", StemmerTokenizer::new(Language::Swedish));
923 self.register("ta_stem", StemmerTokenizer::new(Language::Tamil));
924 self.register("tamil", StemmerTokenizer::new(Language::Tamil));
925 self.register("tr_stem", StemmerTokenizer::new(Language::Turkish));
926 self.register("turkish", StemmerTokenizer::new(Language::Turkish));
927
928 self.register(
930 "en_stop",
931 StopWordTokenizer::new(SimpleTokenizer, Language::English),
932 );
933 self.register(
934 "de_stop",
935 StopWordTokenizer::new(SimpleTokenizer, Language::German),
936 );
937 self.register(
938 "fr_stop",
939 StopWordTokenizer::new(SimpleTokenizer, Language::French),
940 );
941 self.register(
942 "ru_stop",
943 StopWordTokenizer::new(SimpleTokenizer, Language::Russian),
944 );
945 self.register(
946 "es_stop",
947 StopWordTokenizer::new(SimpleTokenizer, Language::Spanish),
948 );
949
950 self.register(
952 "en_stem_stop",
953 StopWordTokenizer::new(StemmerTokenizer::new(Language::English), Language::English),
954 );
955 self.register(
956 "de_stem_stop",
957 StopWordTokenizer::new(StemmerTokenizer::new(Language::German), Language::German),
958 );
959 self.register(
960 "fr_stem_stop",
961 StopWordTokenizer::new(StemmerTokenizer::new(Language::French), Language::French),
962 );
963 self.register(
964 "ru_stem_stop",
965 StopWordTokenizer::new(StemmerTokenizer::new(Language::Russian), Language::Russian),
966 );
967 self.register(
968 "es_stem_stop",
969 StopWordTokenizer::new(StemmerTokenizer::new(Language::Spanish), Language::Spanish),
970 );
971 }
972
973 pub fn register<T: Tokenizer>(&self, name: &str, tokenizer: T) {
975 let mut tokenizers = self.tokenizers.write();
976 tokenizers.insert(name.to_string(), Box::new(tokenizer));
977 }
978
979 pub fn get(&self, name: &str) -> Option<BoxedTokenizer> {
984 if name.starts_with("stem(") {
985 if let Some(tokenizer) = self.dynamic.read().get(name) {
986 return Some(tokenizer.clone());
987 }
988 let tokenizer = TokenizerSpec::parse(name)
989 .ok()
990 .and_then(|spec| spec.dynamic_tokenizer())?;
991 self.dynamic
992 .write()
993 .entry(name.to_string())
994 .or_insert_with(|| tokenizer.clone());
995 return Some(tokenizer);
996 }
997 let tokenizers = self.tokenizers.read();
998 tokenizers.get(name).cloned()
999 }
1000
1001 pub fn contains(&self, name: &str) -> bool {
1003 let tokenizers = self.tokenizers.read();
1004 tokenizers.contains_key(name)
1005 }
1006
1007 pub fn names(&self) -> Vec<String> {
1009 let tokenizers = self.tokenizers.read();
1010 tokenizers.keys().cloned().collect()
1011 }
1012}
1013
1014impl Default for TokenizerRegistry {
1015 fn default() -> Self {
1016 Self::new()
1017 }
1018}
1019
1020#[cfg(test)]
1021mod tests {
1022 use super::*;
1023
1024 #[test]
1025 fn test_simple_tokenizer() {
1026 let tokenizer = SimpleTokenizer;
1027 let tokens = Tokenizer::tokenize(&tokenizer, "Hello World");
1028
1029 assert_eq!(tokens.len(), 2);
1030 assert_eq!(tokens[0].text, "hello");
1031 assert_eq!(tokens[0].position, 0);
1032 assert_eq!(tokens[1].text, "world");
1033 assert_eq!(tokens[1].position, 1);
1034 }
1035
1036 #[test]
1037 fn test_raw_tokenizer() {
1038 let tokenizer = RawTokenizer;
1039 let tokens = Tokenizer::tokenize(&tokenizer, "Hello, World!");
1041 assert_eq!(tokens.len(), 1);
1042 assert_eq!(tokens[0].text, "Hello, World!");
1043 assert_eq!(tokens[0].position, 0);
1044 }
1045
1046 #[test]
1047 fn test_raw_tokenizer_trims() {
1048 let tokenizer = RawTokenizer;
1049 let tokens = Tokenizer::tokenize(&tokenizer, " spaced ");
1050 assert_eq!(tokens.len(), 1);
1051 assert_eq!(tokens[0].text, "spaced");
1052 assert_eq!(tokens[0].offset_from, 2);
1053 }
1054
1055 #[test]
1056 fn test_raw_tokenizer_empty() {
1057 let tokenizer = RawTokenizer;
1058 assert!(Tokenizer::tokenize(&tokenizer, "").is_empty());
1059 assert!(Tokenizer::tokenize(&tokenizer, " ").is_empty());
1060 }
1061
1062 #[test]
1063 fn test_raw_ci_tokenizer() {
1064 let tokenizer = RawCiTokenizer;
1065 let tokens = Tokenizer::tokenize(&tokenizer, "Hello, World!");
1067 assert_eq!(tokens.len(), 1);
1068 assert_eq!(tokens[0].text, "hello, world!");
1069 assert_eq!(tokens[0].position, 0);
1070 }
1071
1072 #[test]
1073 fn test_raw_ci_tokenizer_preserves_structure() {
1074 let tokenizer = RawCiTokenizer;
1075 let tokens = Tokenizer::tokenize(&tokenizer, "HTTPS://Example.COM/Page");
1076 assert_eq!(tokens.len(), 1);
1077 assert_eq!(tokens[0].text, "https://example.com/page");
1078 }
1079
1080 #[test]
1081 fn test_simple_tokenizer_strips_punctuation() {
1082 let tokenizer = SimpleTokenizer;
1083 let tokens = Tokenizer::tokenize(&tokenizer, "Hello, World!");
1084
1085 assert_eq!(tokens.len(), 2);
1086 assert_eq!(tokens[0].text, "hello");
1087 assert_eq!(tokens[1].text, "world");
1088 }
1089
1090 #[test]
1091 fn test_empty_text() {
1092 let tokenizer = SimpleTokenizer;
1093 let tokens = Tokenizer::tokenize(&tokenizer, "");
1094 assert!(tokens.is_empty());
1095 }
1096
1097 #[test]
1098 fn test_stemmer_tokenizer_english() {
1099 let tokenizer = StemmerTokenizer::english();
1100 let tokens = Tokenizer::tokenize(&tokenizer, "Dogs are running quickly");
1101
1102 assert_eq!(tokens.len(), 4);
1103 assert_eq!(tokens[0].text, "dog"); assert_eq!(tokens[1].text, "are"); assert_eq!(tokens[2].text, "run"); assert_eq!(tokens[3].text, "quick"); }
1108
1109 #[test]
1110 fn test_stemmer_tokenizer_preserves_offsets() {
1111 let tokenizer = StemmerTokenizer::english();
1112 let tokens = Tokenizer::tokenize(&tokenizer, "Running dogs");
1113
1114 assert_eq!(tokens.len(), 2);
1115 assert_eq!(tokens[0].text, "run");
1116 assert_eq!(tokens[0].offset_from, 0);
1117 assert_eq!(tokens[0].offset_to, 7); assert_eq!(tokens[1].text, "dog");
1119 assert_eq!(tokens[1].offset_from, 8);
1120 assert_eq!(tokens[1].offset_to, 12); }
1122
1123 #[test]
1124 fn test_stemmer_tokenizer_german() {
1125 let tokenizer = StemmerTokenizer::new(Language::German);
1126 let tokens = Tokenizer::tokenize(&tokenizer, "Häuser Bücher");
1127
1128 assert_eq!(tokens.len(), 2);
1129 assert_eq!(tokens[0].text, "haus"); assert_eq!(tokens[1].text, "buch"); }
1133
1134 #[test]
1135 fn test_stemmer_tokenizer_russian() {
1136 let tokenizer = StemmerTokenizer::new(Language::Russian);
1137 let tokens = Tokenizer::tokenize(&tokenizer, "бегущие собаки");
1138
1139 assert_eq!(tokens.len(), 2);
1140 assert_eq!(tokens[0].text, "бегущ"); assert_eq!(tokens[1].text, "собак"); }
1144
1145 #[test]
1146 fn test_multi_language_stemmer() {
1147 let stemmer = MultiLanguageStemmer::new(Language::English);
1148
1149 let tokens = stemmer.tokenize_with_language("running dogs", Language::English);
1151 assert_eq!(tokens[0].text, "run");
1152 assert_eq!(tokens[1].text, "dog");
1153
1154 let tokens = stemmer.tokenize_with_language("Häuser Bücher", Language::German);
1156 assert_eq!(tokens[0].text, "haus");
1157 assert_eq!(tokens[1].text, "buch");
1158
1159 let tokens = stemmer.tokenize_with_language("бегущие собаки", Language::Russian);
1161 assert_eq!(tokens[0].text, "бегущ");
1162 assert_eq!(tokens[1].text, "собак");
1163 }
1164
1165 #[test]
1166 fn test_language_aware_tokenizer() {
1167 let tokenizer = LanguageAwareTokenizer::new(parse_language);
1168
1169 let tokens = tokenizer.tokenize_with_hint("running dogs", "en");
1171 assert_eq!(tokens[0].text, "run");
1172 assert_eq!(tokens[1].text, "dog");
1173
1174 let tokens = tokenizer.tokenize_with_hint("Häuser Bücher", "de");
1176 assert_eq!(tokens[0].text, "haus");
1177 assert_eq!(tokens[1].text, "buch");
1178
1179 let tokens = tokenizer.tokenize_with_hint("бегущие собаки", "russian");
1181 assert_eq!(tokens[0].text, "бегущ");
1182 assert_eq!(tokens[1].text, "собак");
1183 }
1184
1185 #[test]
1186 fn test_parse_language() {
1187 assert_eq!(parse_language("en"), Language::English);
1188 assert_eq!(parse_language("english"), Language::English);
1189 assert_eq!(parse_language("English"), Language::English);
1190 assert_eq!(parse_language("de"), Language::German);
1191 assert_eq!(parse_language("german"), Language::German);
1192 assert_eq!(parse_language("ru"), Language::Russian);
1193 assert_eq!(parse_language("russian"), Language::Russian);
1194 assert_eq!(parse_language("unknown"), Language::English); }
1196
1197 #[test]
1198 fn test_tokenizer_registry_defaults() {
1199 let registry = TokenizerRegistry::new();
1200
1201 assert!(registry.contains("simple"));
1203 assert!(registry.contains("raw"));
1204 assert!(registry.contains("raw_ci"));
1205 assert!(registry.contains("raw"));
1206 assert!(registry.contains("raw_ci"));
1207 assert!(registry.contains("en_stem"));
1208 assert!(registry.contains("german"));
1209 assert!(registry.contains("russian"));
1210 }
1211
1212 #[test]
1213 fn test_tokenizer_registry_get() {
1214 let registry = TokenizerRegistry::new();
1215
1216 let tokenizer = registry.get("en_stem").unwrap();
1218 let tokens = tokenizer.tokenize("running dogs");
1219 assert_eq!(tokens[0].text, "run");
1220 assert_eq!(tokens[1].text, "dog");
1221
1222 let tokenizer = registry.get("german").unwrap();
1224 let tokens = tokenizer.tokenize("Häuser Bücher");
1225 assert_eq!(tokens[0].text, "haus");
1226 assert_eq!(tokens[1].text, "buch");
1227 }
1228
1229 #[test]
1230 fn test_tokenizer_registry_custom() {
1231 let registry = TokenizerRegistry::new();
1232
1233 registry.register("my_tokenizer", SimpleTokenizer);
1235
1236 assert!(registry.contains("my_tokenizer"));
1237 let tokenizer = registry.get("my_tokenizer").unwrap();
1238 let tokens = tokenizer.tokenize("Hello World");
1239 assert_eq!(tokens[0].text, "hello");
1240 assert_eq!(tokens[1].text, "world");
1241 }
1242
1243 #[test]
1244 fn test_tokenizer_registry_nonexistent() {
1245 let registry = TokenizerRegistry::new();
1246 assert!(registry.get("nonexistent").is_none());
1247 }
1248
1249 #[test]
1250 fn test_stop_word_tokenizer_english() {
1251 let tokenizer = StopWordTokenizer::english(SimpleTokenizer);
1252 let tokens = Tokenizer::tokenize(&tokenizer, "The quick brown fox jumps over the lazy dog");
1253
1254 let texts: Vec<&str> = tokens.iter().map(|t| t.text.as_str()).collect();
1256 assert!(!texts.contains(&"the"));
1257 assert!(!texts.contains(&"over"));
1258 assert!(texts.contains(&"quick"));
1259 assert!(texts.contains(&"brown"));
1260 assert!(texts.contains(&"fox"));
1261 assert!(texts.contains(&"jumps"));
1262 assert!(texts.contains(&"lazy"));
1263 assert!(texts.contains(&"dog"));
1264 }
1265
1266 #[test]
1267 fn test_stop_word_tokenizer_with_stemmer() {
1268 let tokenizer = StopWordTokenizer::new(StemmerTokenizer::english(), Language::English);
1272 let tokens = Tokenizer::tokenize(&tokenizer, "elephants galaxies quantum");
1273
1274 let texts: Vec<&str> = tokens.iter().map(|t| t.text.as_str()).collect();
1275 assert!(texts.contains(&"eleph")); assert!(texts.contains(&"galaxi")); assert!(texts.contains(&"quantum")); }
1280
1281 #[test]
1282 fn test_stop_word_tokenizer_german() {
1283 let tokenizer = StopWordTokenizer::new(SimpleTokenizer, Language::German);
1284 let tokens = Tokenizer::tokenize(&tokenizer, "Der Hund und die Katze");
1285
1286 let texts: Vec<&str> = tokens.iter().map(|t| t.text.as_str()).collect();
1288 assert!(!texts.contains(&"der"));
1289 assert!(!texts.contains(&"und"));
1290 assert!(!texts.contains(&"die"));
1291 assert!(texts.contains(&"hund"));
1292 assert!(texts.contains(&"katze"));
1293 }
1294
1295 #[test]
1296 fn test_stop_word_tokenizer_custom() {
1297 let custom_stops: HashSet<String> = ["foo", "bar"].iter().map(|s| s.to_string()).collect();
1298 let tokenizer = StopWordTokenizer::with_custom_stop_words(SimpleTokenizer, custom_stops);
1299 let tokens = Tokenizer::tokenize(&tokenizer, "foo baz bar qux");
1300
1301 let texts: Vec<&str> = tokens.iter().map(|t| t.text.as_str()).collect();
1302 assert!(!texts.contains(&"foo"));
1303 assert!(!texts.contains(&"bar"));
1304 assert!(texts.contains(&"baz"));
1305 assert!(texts.contains(&"qux"));
1306 }
1307
1308 #[test]
1309 fn test_stop_word_tokenizer_is_stop_word() {
1310 let tokenizer = StopWordTokenizer::english(SimpleTokenizer);
1311 assert!(tokenizer.is_stop_word("the"));
1312 assert!(tokenizer.is_stop_word("and"));
1313 assert!(tokenizer.is_stop_word("is"));
1314 assert!(!tokenizer.is_stop_word("elephant"));
1316 assert!(!tokenizer.is_stop_word("quantum"));
1317 }
1318
1319 #[test]
1320 fn test_tokenizer_registry_stop_word_tokenizers() {
1321 let registry = TokenizerRegistry::new();
1322
1323 assert!(registry.contains("en_stop"));
1325 assert!(registry.contains("en_stem_stop"));
1326 assert!(registry.contains("de_stop"));
1327 assert!(registry.contains("ru_stop"));
1328
1329 let tokenizer = registry.get("en_stop").unwrap();
1331 let tokens = tokenizer.tokenize("The quick fox");
1332 let texts: Vec<&str> = tokens.iter().map(|t| t.text.as_str()).collect();
1333 assert!(!texts.contains(&"the"));
1334 assert!(texts.contains(&"quick"));
1335 assert!(texts.contains(&"fox"));
1336
1337 let tokenizer = registry.get("en_stem_stop").unwrap();
1339 let tokens = tokenizer.tokenize("elephants galaxies");
1340 let texts: Vec<&str> = tokens.iter().map(|t| t.text.as_str()).collect();
1341 assert!(texts.contains(&"eleph")); assert!(texts.contains(&"galaxi")); }
1344
1345 fn hinted<T: Tokenizer>(tokenizer: &T, text: &str, hint: Option<&str>) -> Vec<Token> {
1346 Tokenizer::tokenize_hinted(tokenizer, text, hint)
1347 }
1348
1349 fn texts(tokens: &[Token]) -> Vec<&str> {
1350 tokens.iter().map(|t| t.text.as_str()).collect()
1351 }
1352
1353 #[test]
1354 fn dynamic_stemmer_selects_language_from_hint() {
1355 let stemmer = DynamicStemmer::new(None);
1356 assert_eq!(
1357 texts(&hinted(&stemmer, "Running Foxes", Some("en"))),
1358 vec!["run", "fox"]
1359 );
1360 assert_eq!(
1361 texts(&hinted(&stemmer, "бегущие собаки", Some("ru"))),
1362 vec!["бегущ", "собак"]
1363 );
1364 assert_eq!(
1366 texts(&hinted(&stemmer, "Running Foxes", Some("xx"))),
1367 vec!["running", "foxes"]
1368 );
1369 assert_eq!(
1370 texts(&hinted(&stemmer, "Running Foxes", None)),
1371 vec!["running", "foxes"]
1372 );
1373 assert_eq!(
1374 texts(&Tokenizer::tokenize(&stemmer, "Running, Foxes!")),
1375 vec!["running", "foxes"]
1376 );
1377 let english = DynamicStemmer::new(Some(Language::English));
1379 assert_eq!(
1380 texts(&hinted(&english, "Running Foxes", None)),
1381 vec!["run", "fox"]
1382 );
1383 }
1384
1385 #[test]
1386 fn dynamic_stemmer_routes_tokens_by_script() {
1387 let stemmer = DynamicStemmer::new(None);
1388 assert_eq!(
1390 texts(&hinted(&stemmer, "бегущие foxes", Some("ru,en"))),
1391 vec!["бегущ", "fox"]
1392 );
1393 assert_eq!(
1394 texts(&hinted(&stemmer, "бегущие foxes", Some("en, ru"))),
1395 vec!["бегущ", "fox"]
1396 );
1397 assert_eq!(
1399 texts(&hinted(&stemmer, "бегущие foxes", Some("ru"))),
1400 vec!["бегущ", "foxes"]
1401 );
1402 assert_eq!(
1403 texts(&hinted(&stemmer, "бегущие foxes", Some("en"))),
1404 vec!["бегущие", "fox"]
1405 );
1406 assert_eq!(
1408 texts(&hinted(&stemmer, "running", Some("de,en"))),
1409 vec!["running"]
1410 );
1411 assert_eq!(
1412 texts(&hinted(&stemmer, "running", Some("en,de"))),
1413 vec!["run"]
1414 );
1415 let tokens = hinted(&stemmer, "бегущие foxes run", Some("ru,en"));
1417 assert_eq!(
1418 tokens.iter().map(|t| t.position).collect::<Vec<_>>(),
1419 vec![0, 1, 2]
1420 );
1421 }
1422
1423 #[test]
1424 fn script_detection_covers_supported_stemmer_scripts() {
1425 assert_eq!(Script::of_token("hello"), Script::Latin);
1426 assert_eq!(Script::of_token("straße"), Script::Latin);
1427 assert_eq!(Script::of_token("собака"), Script::Cyrillic);
1428 assert_eq!(Script::of_token("γεια"), Script::Greek);
1429 assert_eq!(Script::of_token("مرحبا"), Script::Arabic);
1430 assert_eq!(Script::of_token("தமிழ்"), Script::Tamil);
1431 assert_eq!(Script::of_token("日本語"), Script::Other);
1432 assert_eq!(Script::of_token("2024"), Script::Other);
1433 assert_eq!(Language::Russian.script(), Script::Cyrillic);
1434 assert_eq!(Language::Turkish.script(), Script::Latin);
1435 }
1436
1437 #[test]
1438 fn tokenizer_spec_parses_and_renders_canonically() {
1439 assert_eq!(
1440 TokenizerSpec::parse("en_stem").unwrap(),
1441 TokenizerSpec::Named("en_stem".to_string())
1442 );
1443 let spec = TokenizerSpec::parse("stem(by:languages,default:simple)").unwrap();
1444 assert_eq!(
1445 spec,
1446 TokenizerSpec::DynamicStem {
1447 by: "languages".to_string(),
1448 default: None
1449 }
1450 );
1451 assert_eq!(spec.to_string(), "stem(by: languages, default: simple)");
1452 assert_eq!(spec.hint_field(), Some("languages"));
1453
1454 let spec = TokenizerSpec::parse("stem(by: lang, default: english)").unwrap();
1455 assert_eq!(spec.to_string(), "stem(by: lang, default: en)");
1456 assert_eq!(
1457 TokenizerSpec::parse("stem(by: lang)").unwrap(),
1458 TokenizerSpec::DynamicStem {
1459 by: "lang".to_string(),
1460 default: None
1461 }
1462 );
1463
1464 assert!(TokenizerSpec::parse("stem(default: en)").is_err());
1465 assert!(TokenizerSpec::parse("stem(by: lang, default: klingon)").is_err());
1466 assert!(TokenizerSpec::parse("stem(by: lang").is_err());
1467 assert!(TokenizerSpec::parse("stem(by: lang, color: red)").is_err());
1468 assert!(TokenizerSpec::parse("en_stem(foo)").is_err());
1469 assert!(TokenizerSpec::parse("").is_err());
1470 }
1471
1472 #[test]
1473 fn registry_builds_dynamic_stemmer_from_spec() {
1474 let registry = TokenizerRegistry::new();
1475 let tokenizer = registry
1476 .get("stem(by: languages, default: simple)")
1477 .expect("dynamic spec resolves without registration");
1478 assert_eq!(
1479 texts(&tokenizer.tokenize_hinted("Running Foxes", Some("en"))),
1480 vec!["run", "fox"]
1481 );
1482 assert_eq!(
1483 texts(&tokenizer.tokenize("Running Foxes")),
1484 vec!["running", "foxes"]
1485 );
1486 assert!(
1487 registry
1488 .get("stem(by: languages, default: klingon)")
1489 .is_none()
1490 );
1491 let simple = registry.get("en_stem").unwrap();
1493 assert_eq!(
1494 texts(&simple.tokenize_hinted("Running Foxes", Some("ru"))),
1495 vec!["run", "fox"]
1496 );
1497 }
1498
1499 #[test]
1500 fn parse_language_opt_rejects_unknown_codes() {
1501 assert_eq!(parse_language_opt(" RU "), Some(Language::Russian));
1502 assert_eq!(parse_language_opt("german"), Some(Language::German));
1503 assert_eq!(parse_language_opt("xx"), None);
1504 assert_eq!(parse_language_opt(""), None);
1505 for language in [Language::English, Language::Russian, Language::Tamil] {
1506 assert_eq!(parse_language_opt(language_code(language)), Some(language));
1507 }
1508 }
1509}