From 34d568d5e229729031245aee488481c7f623125c Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Fri, 21 Jun 2024 07:22:43 +0800
Subject: refactor: improve rag (#622)

---
 src/rag/bm25.rs | 25 ++++++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

(limited to 'src/rag/bm25.rs')

diff --git a/src/rag/bm25.rs b/src/rag/bm25.rs
index 3dee357..c91dcf5 100644
--- a/src/rag/bm25.rs
+++ b/src/rag/bm25.rs
@@ -1,6 +1,7 @@
 use rayon::prelude::*;
 use std::collections::HashMap;
 use std::f64;
+use unicode_segmentation::UnicodeSegmentation;
 
 #[derive(Debug, Clone)]
 pub struct BM25Options {
@@ -147,13 +148,35 @@ impl<T: Clone> BM25<T> {
 }
 
 fn tokenize(text: &str) -> Vec<String> {
-    text.split(' ').map(|v| v.to_string()).collect()
+    text.unicode_words()
+        .filter_map(|v| {
+            if [
+                "a", "an", "and", "are", "as", "at", "be", "but", "by", "for", "if", "in", "into",
+                "is", "it", "no", "not", "of", "on", "or", "such", "that", "the", "their", "then",
+                "there", "these", "they", "this", "to", "was", "will", "with",
+            ]
+            .contains(&v)
+            {
+                None
+            } else {
+                Some(v.to_string())
+            }
+        })
+        .collect()
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
+    #[test]
+    fn test_tokenize() {
+        assert_eq!(
+            tokenize("a quick fox jumps over the lazy dog"),
+            vec!["quick", "fox", "jumps", "over", "lazy", "dog"]
+        );
+    }
+
     #[test]
     fn test_bm25() {
         let corpus = vec![
-- 
cgit v1.2.3