From 2eab71a641827e503b14952373aec82661192ba2 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Thu, 20 Jun 2024 11:26:45 +0800
Subject: feat: rag hybrid search (#618)

---
 src/rag/bm25.rs | 172 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 172 insertions(+)
 create mode 100644 src/rag/bm25.rs

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

diff --git a/src/rag/bm25.rs b/src/rag/bm25.rs
new file mode 100644
index 0000000..3dee357
--- /dev/null
+++ b/src/rag/bm25.rs
@@ -0,0 +1,172 @@
+use rayon::prelude::*;
+use std::collections::HashMap;
+use std::f64;
+
+#[derive(Debug, Clone)]
+pub struct BM25Options {
+    k1: f64,
+    b: f64,
+    epsilon: f64,
+}
+
+impl Default for BM25Options {
+    fn default() -> Self {
+        Self {
+            k1: 1.5,
+            b: 0.75,
+            epsilon: 0.25,
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct BM25<T> {
+    options: BM25Options,
+    corpus_size: usize,
+    avgdl: f64,
+    doc_freqs: Vec<HashMap<String, u32>>,
+    doc_ids: Vec<T>,
+    idf: HashMap<String, f64>,
+    doc_len: Vec<usize>,
+}
+
+impl<T: Clone> BM25<T> {
+    pub fn new(corpus: Vec<(T, String)>, options: BM25Options) -> Self {
+        let mut doc_ids = vec![];
+        let mut docs = vec![];
+        for (id, value) in corpus {
+            doc_ids.push(id);
+            docs.push(value);
+        }
+        let tokenized_docs = docs.into_par_iter().map(|text| tokenize(&text)).collect();
+
+        let mut bm25 = BM25 {
+            options,
+            corpus_size: 0,
+            avgdl: 0.0,
+            doc_freqs: Vec::new(),
+            doc_ids,
+            idf: HashMap::new(),
+            doc_len: Vec::new(),
+        };
+
+        let map = bm25.initialize(tokenized_docs);
+        bm25.calc_idf(map);
+
+        bm25
+    }
+
+    pub fn search(&self, query: &str, top_k: usize, min_score: Option<f64>) -> Vec<T> {
+        let scores = self.get_scores(query);
+        let mut indexed_scores: Vec<(T, f64)> = scores
+            .into_iter()
+            .enumerate()
+            .filter_map(|(i, v)| match min_score {
+                Some(minimum_score) => {
+                    if v < minimum_score {
+                        None
+                    } else {
+                        Some((self.doc_ids[i].clone(), v))
+                    }
+                }
+                None => Some((self.doc_ids[i].clone(), v)),
+            })
+            .collect();
+        indexed_scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
+        indexed_scores
+            .into_iter()
+            .take(top_k)
+            .map(|(id, _)| id)
+            .collect()
+    }
+
+    pub fn get_scores(&self, query: &str) -> Vec<f64> {
+        let mut score = vec![0.0; self.corpus_size];
+
+        for q in tokenize(query) {
+            if let Some(idf) = self.idf.get(&q) {
+                for (i, doc) in self.doc_freqs.iter().enumerate() {
+                    let q_freq = doc.get(&q).unwrap_or(&0);
+                    score[i] += *idf
+                        * (*q_freq as f64 * (self.options.k1 + 1.0)
+                            / (*q_freq as f64
+                                + self.options.k1
+                                    * (1.0 - self.options.b
+                                        + self.options.b * self.doc_len[i] as f64 / self.avgdl)));
+                }
+            }
+        }
+
+        score
+    }
+
+    fn initialize(&mut self, corpus: Vec<Vec<String>>) -> HashMap<String, usize> {
+        let mut map = HashMap::new();
+        let mut num_doc = 0;
+
+        for document in corpus {
+            self.doc_len.push(document.len());
+            num_doc += document.len();
+
+            let mut frequencies = HashMap::new();
+            for word in document {
+                *frequencies.entry(word).or_insert(0) += 1;
+            }
+            self.doc_freqs.push(frequencies);
+
+            for word in self.doc_freqs[self.doc_freqs.len() - 1].keys() {
+                *map.entry(word.clone()).or_insert(0) += 1;
+            }
+
+            self.corpus_size += 1;
+        }
+
+        self.avgdl = num_doc as f64 / self.corpus_size as f64;
+        map
+    }
+
+    fn calc_idf(&mut self, map: HashMap<String, usize>) {
+        let mut idf_sum = 0.0;
+        let mut negative_idfs = Vec::new();
+
+        for (word, freq) in map {
+            let idf = (self.corpus_size as f64 - freq as f64 + 0.5).ln() - (freq as f64 + 0.5).ln();
+            self.idf.insert(word.clone(), idf);
+            idf_sum += idf;
+            if idf < 0.0 {
+                negative_idfs.push(word);
+            }
+        }
+
+        let average_idf = idf_sum / self.idf.len() as f64;
+
+        for word in negative_idfs {
+            self.idf.insert(word, self.options.epsilon * average_idf);
+        }
+    }
+}
+
+fn tokenize(text: &str) -> Vec<String> {
+    text.split(' ').map(|v| v.to_string()).collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_bm25() {
+        let corpus = vec![
+            (0, "Hello there good man!".into()),
+            (1, "It is quite windy in London".into()),
+            (2, "How is the weather today?".into()),
+        ];
+        let bm25 = BM25::new(corpus, BM25Options::default());
+
+        let scores = bm25.get_scores("windy London");
+        assert_eq!(scores, [0.0, 0.9372947225064051, 0.0]);
+
+        let top_n = bm25.search("windy London", 3, None);
+        assert_eq!(top_n, vec![1, 0, 2])
+    }
+}
-- 
cgit v1.2.3