1use hashbrown::HashSet;
8use std::time::Instant;
9use unicode_segmentation::{UnicodeSegmentation, UnicodeWords};
10use whatlang::{
11 Lang, detect as lang_detect_all, detect_lang as lang_detect, detect_script as script_detect,
12};
13
14#[cfg(feature = "tokenizer-chinese")]
15use std::vec::IntoIter;
16
17use super::stopwords::LexerStopWord;
18use crate::query::QueryGenericLang;
19use crate::store::identifiers::{StoreTermHash, StoreTermHashed};
20
21pub struct TokenLexerBuilder;
22
23pub struct TokenLexer<'a> {
24 mode: TokenLexerMode,
25 locale: Option<Lang>,
26 words: TokenLexerWords<'a>,
27 yields: HashSet<StoreTermHashed>,
28}
29
30#[derive(PartialEq)]
31pub enum TokenLexerMode {
32 NormalizeAndCleanup(Option<Lang>),
33 NormalizeOnly,
34}
35
36enum TokenLexerWords<'a> {
37 UAX29(UnicodeWords<'a>),
38
39 #[cfg(feature = "tokenizer-chinese")]
40 JieBa(IntoIter<&'a str>),
41
42 #[cfg(feature = "tokenizer-japanese")]
43 Lindera(IntoIter<lindera_tokenizer::token::Token<'a>>),
44}
45
46const TEXT_LANG_TRUNCATE_OVER_CHARS: usize = 200;
47const TEXT_LANG_DETECT_PROCEED_OVER_CHARS: usize = 20;
48const TEXT_LANG_DETECT_NGRAM_UNDER_CHARS: usize = 60;
49
50#[cfg(feature = "tokenizer-chinese")]
51lazy_static! {
52 static ref TOKENIZER_JIEBA: jieba_rs::Jieba = jieba_rs::Jieba::new();
53}
54
55#[cfg(feature = "tokenizer-japanese")]
56lazy_static! {
57 static ref TOKENIZER_LINDERA: lindera_tokenizer::tokenizer::Tokenizer =
58 lindera_tokenizer::tokenizer::Tokenizer::from_config(
59 lindera_tokenizer::tokenizer::TokenizerConfig {
60 dictionary: lindera_dictionary::DictionaryConfig {
61 kind: Some(lindera_dictionary::DictionaryKind::UniDic),
62 path: None
63 },
64 user_dictionary: None,
65 mode: lindera_core::mode::Mode::Normal,
66 }
67 )
68 .expect("unable to initialize japanese tokenizer");
69}
70
71impl TokenLexerBuilder {
72 pub fn from(mode: TokenLexerMode, text: &str) -> Result<TokenLexer<'_>, ()> {
73 let locale = match mode {
74 TokenLexerMode::NormalizeAndCleanup(None) => {
75 tracing::debug!("detecting locale from lexer text: {}", text);
77
78 Self::detect_lang(text)
79 }
80 TokenLexerMode::NormalizeAndCleanup(Some(lang)) => {
81 tracing::debug!("using hinted locale: {} from lexer text: {}", lang, text);
83
84 Some(lang)
85 }
86 TokenLexerMode::NormalizeOnly => {
87 tracing::debug!("not detecting locale from lexer text: {}", text);
88
89 None
91 }
92 };
93
94 Ok(TokenLexer::new(mode, text, locale))
96 }
97
98 fn detect_lang(text: &str) -> Option<Lang> {
99 if text.len() < TEXT_LANG_DETECT_PROCEED_OVER_CHARS {
102 return None;
103 }
104
105 let safe_text = if text.len() > TEXT_LANG_TRUNCATE_OVER_CHARS {
108 tracing::debug!(
109 "lexer text needs to be truncated, as it is too long ({}/{}): {}",
110 text.len(),
111 TEXT_LANG_TRUNCATE_OVER_CHARS,
112 text
113 );
114
115 text.char_indices()
125 .nth(TEXT_LANG_TRUNCATE_OVER_CHARS)
126 .map(|(end_index, _)| &text[0..end_index])
127 .unwrap_or(text)
128 } else {
129 text
130 };
131
132 tracing::debug!("will detect locale for lexer safe text: {}", safe_text);
133
134 if safe_text.len() < TEXT_LANG_DETECT_NGRAM_UNDER_CHARS {
144 tracing::debug!(
145 "lexer text is shorter than {} characters, using the slow method",
146 TEXT_LANG_DETECT_NGRAM_UNDER_CHARS
147 );
148
149 Self::detect_lang_slow(safe_text)
150 } else {
151 tracing::debug!(
152 "lexer text is equal or longer than {} characters, using the fast method",
153 TEXT_LANG_DETECT_NGRAM_UNDER_CHARS
154 );
155
156 Self::detect_lang_fast(safe_text)
157 }
158 }
159
160 fn detect_lang_slow(safe_text: &str) -> Option<Lang> {
161 let ngram_start = Instant::now();
162
163 match lang_detect_all(safe_text) {
164 Some(detector) => {
165 let ngram_took = ngram_start.elapsed();
166
167 let mut locale = detector.lang();
168
169 tracing::info!(
170 "[slow lexer] locale detected from text: {} ({} from {} at {}/1; {}s + {}ms)",
171 safe_text,
172 locale,
173 detector.script(),
174 detector.confidence(),
175 ngram_took.as_secs(),
176 ngram_took.subsec_millis()
177 );
178
179 if !detector.is_reliable() {
183 tracing::debug!("[slow lexer] trying to detect locale from stopwords instead");
184
185 if let Some(alternate_locale) =
187 LexerStopWord::guess_lang(safe_text, detector.script())
188 {
189 tracing::info!(
190 "[slow lexer] detected more accurate locale from stopwords: {}",
191 alternate_locale
192 );
193
194 locale = alternate_locale;
195 }
196 }
197
198 Some(locale)
199 }
200 None => {
201 tracing::info!(
202 "[slow lexer] no locale could be detected from text: {}",
203 safe_text
204 );
205
206 None
207 }
208 }
209 }
210
211 fn detect_lang_fast(safe_text: &str) -> Option<Lang> {
212 let stopwords_start = Instant::now();
213
214 match script_detect(safe_text) {
215 Some(script) => {
216 if let Some(locale) = LexerStopWord::guess_lang(safe_text, script) {
218 let stopwords_took = stopwords_start.elapsed();
219
220 tracing::info!(
221 "[fast lexer] locale detected from text: {} ({}; {}s + {}ms)",
222 safe_text,
223 locale,
224 stopwords_took.as_secs(),
225 stopwords_took.subsec_millis()
226 );
227
228 Some(locale)
229 } else {
230 tracing::debug!(
231 "[fast lexer] trying to detect locale from fallback ngram instead"
232 );
233
234 lang_detect(safe_text)
236 }
237 }
238 None => {
239 tracing::info!(
240 "[fast lexer] no script could be detected from text: {}",
241 safe_text
242 );
243
244 None
245 }
246 }
247 }
248}
249
250impl<'a> TokenLexer<'a> {
251 fn new(mode: TokenLexerMode, text: &'a str, locale: Option<Lang>) -> TokenLexer<'a> {
252 let words = match locale {
254 #[cfg(feature = "tokenizer-chinese")]
255 Some(Lang::Cmn) => TokenLexerWords::JieBa(TOKENIZER_JIEBA.cut(text, false).into_iter()),
256 #[cfg(feature = "tokenizer-japanese")]
257 Some(Lang::Jpn) => match TOKENIZER_LINDERA.tokenize(text) {
258 Ok(tokens) => TokenLexerWords::Lindera(tokens.into_iter()),
259 Err(err) => {
260 tracing::warn!("unable to tokenize japanese, falling back: {}", err);
261
262 TokenLexerWords::UAX29(text.unicode_words())
263 }
264 },
265 _ => TokenLexerWords::UAX29(text.unicode_words()),
266 };
267
268 TokenLexer {
269 mode,
270 locale,
271 words,
272 yields: HashSet::new(),
273 }
274 }
275}
276
277impl TokenLexerMode {
278 pub fn from_query_lang(lang: Option<QueryGenericLang>) -> TokenLexerMode {
279 match lang {
280 Some(QueryGenericLang::Enabled(lang)) => {
281 TokenLexerMode::NormalizeAndCleanup(Some(lang))
283 }
284 Some(QueryGenericLang::Disabled) => {
285 TokenLexerMode::NormalizeOnly
287 }
288 None => {
289 TokenLexerMode::NormalizeAndCleanup(None)
291 }
292 }
293 }
294}
295
296impl<'a> Iterator for TokenLexer<'a> {
297 type Item = (String, StoreTermHashed);
298
299 fn next(&mut self) -> Option<Self::Item> {
305 for word in &mut self.words {
306 let word = word.to_lowercase();
310
311 if self.mode == TokenLexerMode::NormalizeOnly || !LexerStopWord::is(&word, self.locale)
313 {
314 let term_hash = StoreTermHash::from(&word);
318
319 if !self.yields.contains(&term_hash) {
321 tracing::debug!("lexer yielded word: {}", word);
322
323 self.yields.insert(term_hash);
324
325 return Some((word, term_hash));
326 } else {
327 tracing::debug!(
328 "lexer did not yield word: {} because: word already yielded",
329 word
330 );
331 }
332 } else {
333 tracing::debug!(
334 "lexer did not yield word: {} because: word is a stop-word",
335 word
336 );
337 }
338 }
339
340 None
341 }
342}
343
344impl<'a> Iterator for TokenLexerWords<'a> {
345 type Item = &'a str;
346
347 fn next(&mut self) -> Option<Self::Item> {
348 match self {
349 TokenLexerWords::UAX29(token) => token.next(),
350
351 #[cfg(feature = "tokenizer-chinese")]
352 TokenLexerWords::JieBa(token) => token.next(),
353
354 #[cfg(feature = "tokenizer-japanese")]
355 TokenLexerWords::Lindera(token) => match token.next() {
356 Some(inner) => Some(inner.text),
357 None => None,
358 },
359 }
360 }
361}
362
363#[cfg(test)]
364mod tests {
365 use super::*;
366
367 #[test]
368 fn it_cleans_token_english() {
369 let mut token_cleaner = TokenLexerBuilder::from(
370 TokenLexerMode::NormalizeAndCleanup(None),
371 "The quick brown fox jumps over the lazy dog!",
372 )
373 .unwrap();
374
375 assert_eq!(token_cleaner.locale, Some(Lang::Eng));
376 assert_eq!(
377 token_cleaner.next(),
378 Some(("quick".to_string(), 4179131656))
379 );
380 assert_eq!(
381 token_cleaner.next(),
382 Some(("brown".to_string(), 1268820067))
383 );
384 assert_eq!(token_cleaner.next(), Some(("fox".to_string(), 667256324)));
385 assert_eq!(token_cleaner.next(), Some(("jumps".to_string(), 633865164)));
386 assert_eq!(token_cleaner.next(), Some(("lazy".to_string(), 4130433347)));
387 assert_eq!(token_cleaner.next(), Some(("dog".to_string(), 2044924251)));
388 assert_eq!(token_cleaner.next(), None);
389 }
390
391 #[test]
392 fn it_cleans_token_french() {
393 let mut token_cleaner = TokenLexerBuilder::from(
394 TokenLexerMode::NormalizeAndCleanup(None),
395 "Le vif renard brun saute par dessus le chien paresseux.",
396 )
397 .unwrap();
398
399 assert_eq!(token_cleaner.locale, Some(Lang::Fra));
400 assert_eq!(
401 token_cleaner.next(),
402 Some(("renard".to_string(), 1635186311))
403 );
404 assert_eq!(token_cleaner.next(), Some(("brun".to_string(), 2763604928)));
405 assert_eq!(
406 token_cleaner.next(),
407 Some(("saute".to_string(), 1918158211))
408 );
409 assert_eq!(
410 token_cleaner.next(),
411 Some(("chien".to_string(), 2177818351))
412 );
413 assert_eq!(
414 token_cleaner.next(),
415 Some(("paresseux".to_string(), 1678693110))
416 );
417 assert_eq!(token_cleaner.next(), None);
418 }
419
420 #[cfg(feature = "tokenizer-chinese")]
421 #[test]
422 fn it_cleans_token_chinese_jieba() {
423 let mut token_cleaner = TokenLexerBuilder::from(
424 TokenLexerMode::NormalizeAndCleanup(None),
425 "我们中出了一个叛徒",
426 )
427 .unwrap();
428
429 assert_eq!(token_cleaner.locale, Some(Lang::Cmn));
430 assert_eq!(token_cleaner.next(), Some(("出".to_string(), 241978070)));
431 assert_eq!(token_cleaner.next(), Some(("一个".to_string(), 2596274530)));
432 assert_eq!(token_cleaner.next(), Some(("叛徒".to_string(), 3244183759)));
433 assert_eq!(token_cleaner.next(), None);
434 }
435
436 #[cfg(not(feature = "tokenizer-chinese"))]
437 #[test]
438 fn it_cleans_token_chinese_naive() {
439 let mut token_cleaner = TokenLexerBuilder::from(
440 TokenLexerMode::NormalizeAndCleanup(None),
441 "快狐跨懒狗快狐跨懒狗",
442 )
443 .unwrap();
444
445 assert_eq!(token_cleaner.locale, Some(Lang::Cmn));
446 assert_eq!(token_cleaner.next(), Some(("快".to_string(), 126546256)));
447 assert_eq!(token_cleaner.next(), Some(("狐".to_string(), 2879689662)));
448 assert_eq!(token_cleaner.next(), Some(("跨".to_string(), 2913342670)));
449 assert_eq!(token_cleaner.next(), Some(("懒".to_string(), 3199935961)));
450 assert_eq!(token_cleaner.next(), Some(("狗".to_string(), 3360772096)));
451 assert_eq!(token_cleaner.next(), None);
452 }
453
454 #[cfg(feature = "tokenizer-japanese")]
455 #[test]
456 fn it_cleans_token_japanese_lindera_product() {
457 let mut token_cleaner = TokenLexerBuilder::from(
458 TokenLexerMode::NormalizeAndCleanup(None),
459 "関西国際空港限定トートバッグ",
460 )
461 .unwrap();
462
463 assert_eq!(token_cleaner.locale, Some(Lang::Jpn));
464 assert_eq!(token_cleaner.next(), Some(("関西".to_string(), 1283572620)));
465 assert_eq!(token_cleaner.next(), Some(("国際".to_string(), 2132457693)));
466 assert_eq!(token_cleaner.next(), Some(("空港".to_string(), 865668138)));
467 assert_eq!(token_cleaner.next(), Some(("限定".to_string(), 3708465176)));
468 assert_eq!(
469 token_cleaner.next(),
470 Some(("トート".to_string(), 881444746))
471 );
472 assert_eq!(
473 token_cleaner.next(),
474 Some(("バッグ".to_string(), 3515727814))
475 );
476 assert_eq!(token_cleaner.next(), None);
477 }
478
479 #[cfg(feature = "tokenizer-japanese")]
480 #[test]
481 fn it_cleans_token_japanese_lindera_food() {
482 let token_cleaner =
483 TokenLexerBuilder::from(TokenLexerMode::NormalizeAndCleanup(None), "𠮷野家").unwrap();
484
485 assert_eq!(token_cleaner.locale, None);
486
487 let token_cleaner =
488 TokenLexerBuilder::from(TokenLexerMode::NormalizeAndCleanup(None), "ヱビスビール")
489 .unwrap();
490
491 assert_eq!(token_cleaner.locale, None);
492 }
493
494 #[cfg(feature = "tokenizer-japanese")]
495 #[test]
496 fn it_cleans_token_japanese_lindera_sentence() {
497 let mut token_cleaner = TokenLexerBuilder::from(
498 TokenLexerMode::NormalizeAndCleanup(None),
499 "𠮷野家でヱビスビールを飲んだ",
500 )
501 .unwrap();
502
503 assert_eq!(token_cleaner.locale, Some(Lang::Jpn));
504 assert_eq!(token_cleaner.next(), Some(("𠮷".to_string(), 2866455824)));
505 assert_eq!(token_cleaner.next(), Some(("野家".to_string(), 1324395598)));
506 assert_eq!(
507 token_cleaner.next(),
508 Some(("ヱビス".to_string(), 1696836208))
509 );
510 assert_eq!(
511 token_cleaner.next(),
512 Some(("ビール".to_string(), 3421909800))
513 );
514 assert_eq!(token_cleaner.next(), Some(("飲ん".to_string(), 3196735184)));
515 assert_eq!(token_cleaner.next(), None);
516 }
517
518 #[test]
519 fn it_cleans_token_emojis() {
520 let mut token_cleaner =
521 TokenLexerBuilder::from(TokenLexerMode::NormalizeAndCleanup(None), "🚀 🙋♂️🙋♂️🙋♂️")
522 .unwrap();
523
524 assert_eq!(token_cleaner.locale, None);
525 assert_eq!(token_cleaner.next(), None);
526 }
527
528 #[test]
529 fn it_cleans_token_lang_hinted() {
530 let mut token_cleaner_right = TokenLexerBuilder::from(
531 TokenLexerMode::NormalizeAndCleanup(Some(Lang::Eng)),
532 "This will be cleaned properly, as English was hinted rightfully so.",
533 )
534 .unwrap();
535 let mut token_cleaner_wrong = TokenLexerBuilder::from(
536 TokenLexerMode::NormalizeAndCleanup(Some(Lang::Fra)),
537 "This will not be cleaned properly, as French was hinted but this is English.",
538 )
539 .unwrap();
540
541 assert_eq!(token_cleaner_right.locale, Some(Lang::Eng));
542 assert_eq!(token_cleaner_wrong.locale, Some(Lang::Fra));
543
544 assert_eq!(
545 token_cleaner_right.next(),
546 Some(("cleaned".to_string(), 3550382624))
547 );
548 assert_eq!(
549 token_cleaner_wrong.next(),
550 Some(("this".to_string(), 493303710))
551 );
552 }
553
554 #[test]
555 fn it_detects_lang_english_regular() {
556 assert_eq!(
557 TokenLexerBuilder::detect_lang("The quick brown fox jumps over the lazy dog!"),
558 Some(Lang::Eng)
559 );
560 }
561
562 #[test]
563 fn it_detects_lang_english_long() {
564 assert_eq!(
565 TokenLexerBuilder::detect_lang(
566 r#"Running an electrical current through water splits it into oxygen and hydrogen,
567 the latter of which can be used as a reliable, zero-emission fuel source. In the past,
568 the process of purifying water beforehand was too energy intensive for this process to
569 be useful — but now scientists have figured out how to skip the process altogether and
570 convert seawater into usable hydrogen"#
571 ),
572 Some(Lang::Eng)
573 );
574 }
575
576 #[test]
577 fn it_doesnt_detect_lang_english_tiny() {
578 assert_eq!(TokenLexerBuilder::detect_lang("The quick"), None);
579 }
580}
581
582#[cfg(all(feature = "benchmark", test))]
583mod benches {
584 extern crate test;
585
586 use super::*;
587 use test::Bencher;
588
589 #[bench]
590 fn bench_normalize_token_french_build(b: &mut Bencher) {
591 b.iter(|| {
592 TokenLexerBuilder::from(
593 TokenLexerMode::NormalizeOnly,
594 "Le vif renard brun saute par dessus le chien paresseux.",
595 )
596 });
597 }
598
599 #[bench]
600 fn bench_normalize_token_french_exhaust(b: &mut Bencher) {
601 b.iter(|| {
602 let token_cleaner = TokenLexerBuilder::from(
603 TokenLexerMode::NormalizeOnly,
604 "Le vif renard brun saute par dessus le chien paresseux.",
605 )
606 .unwrap();
607
608 token_cleaner.map(|value| value.1).collect::<Vec<u32>>()
609 });
610 }
611
612 #[bench]
613 fn bench_clean_token_english_regular_build(b: &mut Bencher) {
614 b.iter(|| {
615 TokenLexerBuilder::from(
616 TokenLexerMode::NormalizeAndCleanup(None),
617 "The quick brown fox jumps over the lazy dog!",
618 )
619 });
620 }
621
622 #[bench]
623 fn bench_clean_token_english_regular_exhaust(b: &mut Bencher) {
624 b.iter(|| {
625 let token_cleaner = TokenLexerBuilder::from(
626 TokenLexerMode::NormalizeAndCleanup(None),
627 "The quick brown fox jumps over the lazy dog!",
628 )
629 .unwrap();
630
631 token_cleaner.map(|value| value.1).collect::<Vec<u32>>()
632 });
633 }
634
635 #[bench]
636 fn bench_clean_token_english_long_exhaust(b: &mut Bencher) {
637 b.iter(|| {
638 let token_cleaner = TokenLexerBuilder::from(
639 TokenLexerMode::NormalizeAndCleanup(None),
640 r#"Running an electrical current through water splits it into oxygen and hydrogen,
641 the latter of which can be used as a reliable, zero-emission fuel source. In the
642 past, the process of purifying water beforehand was too energy intensive for this
643 process to be useful — but now scientists have figured out how to skip the process
644 altogether and convert seawater into usable hydrogen"#,
645 )
646 .unwrap();
647
648 token_cleaner.map(|value| value.1).collect::<Vec<u32>>()
649 });
650 }
651
652 #[bench]
653 fn bench_clean_token_english_hinted_build(b: &mut Bencher) {
654 b.iter(|| {
655 TokenLexerBuilder::from(
656 TokenLexerMode::NormalizeAndCleanup(Some(Lang::Eng)),
657 "The quick brown fox jumps over the lazy dog!",
658 )
659 });
660 }
661
662 #[bench]
663 fn bench_clean_token_english_hinted_exhaust(b: &mut Bencher) {
664 b.iter(|| {
665 let token_cleaner = TokenLexerBuilder::from(
666 TokenLexerMode::NormalizeAndCleanup(Some(Lang::Eng)),
667 "The quick brown fox jumps over the lazy dog!",
668 )
669 .unwrap();
670
671 token_cleaner.map(|value| value.1).collect::<Vec<u32>>()
672 });
673 }
674
675 #[bench]
676 fn bench_clean_token_chinese_build(b: &mut Bencher) {
677 b.iter(|| {
678 TokenLexerBuilder::from(
679 TokenLexerMode::NormalizeAndCleanup(None),
680 "我们中出了一个叛徒",
681 )
682 });
683 }
684
685 #[bench]
686 fn bench_clean_token_chinese_exhaust(b: &mut Bencher) {
687 b.iter(|| {
688 let token_cleaner = TokenLexerBuilder::from(
689 TokenLexerMode::NormalizeAndCleanup(None),
690 "我们中出了一个叛徒",
691 )
692 .unwrap();
693
694 token_cleaner.map(|value| value.1).collect::<Vec<u32>>()
695 });
696 }
697
698 #[bench]
699 fn bench_clean_token_japanese_build(b: &mut Bencher) {
700 b.iter(|| {
701 TokenLexerBuilder::from(
702 TokenLexerMode::NormalizeAndCleanup(None),
703 "関西国際空港限定トートバッグ",
704 )
705 });
706 }
707
708 #[bench]
709 fn bench_clean_token_japanese_exhaust(b: &mut Bencher) {
710 b.iter(|| {
711 let token_cleaner = TokenLexerBuilder::from(
712 TokenLexerMode::NormalizeAndCleanup(None),
713 "関西国際空港限定トートバッグ",
714 )
715 .unwrap();
716
717 token_cleaner.map(|value| value.1).collect::<Vec<u32>>()
718 });
719 }
720
721 #[bench]
722 fn bench_detect_lang_english_short(b: &mut Bencher) {
723 b.iter(|| TokenLexerBuilder::detect_lang("The quick brown fox."));
724 }
725
726 #[bench]
727 fn bench_detect_lang_english_regular(b: &mut Bencher) {
728 b.iter(|| TokenLexerBuilder::detect_lang("The quick brown fox jumps over the lazy dog!"));
729 }
730
731 #[bench]
732 fn bench_detect_lang_english_long(b: &mut Bencher) {
733 b.iter(|| {
734 TokenLexerBuilder::detect_lang(
735 r#"Running an electrical current through water splits it into oxygen and hydrogen,
736 the latter of which can be used as a reliable, zero-emission fuel source. In the past,
737 the process of purifying water beforehand was too energy intensive for this process to
738 be useful — but now scientists have figured out how to skip the process altogether and
739 convert seawater into usable hydrogen"#,
740 )
741 });
742 }
743
744 #[bench]
745 fn bench_dont_detect_lang_english_tiny(b: &mut Bencher) {
746 b.iter(|| TokenLexerBuilder::detect_lang("The quick"));
747 }
748}