Implement BM25 Algorithm for better Document Ranking

Author: SamarthPyati

src/main.rs

@@ -268,7 +268,7 @@ fn entry() -> Result<(), ()> {
             })?.chars().collect::<Vec<char>>();
 
             let model = fetch_model(&index_path)?;
-            for (path, rank) in model.search_query(&prompt)?.iter().take(20) {
+            for (path, rank) in model.search_query(&prompt, &model)?.iter().take(20) {
                 println!("{path} - {rank}", path = path.display());
             } 
 
@@ -343,4 +343,3 @@ fn main() -> io::Result<()> {
 
 // TODO: Synonym terms
 // TODO: Add levenstein distance or cosine similarity
-// TODO: Add better document ranker specifically "Okapi BM-25"

src/model.rs

@@ -5,10 +5,11 @@ use std::{
 use std::collections::HashMap;
 use serde::{Deserialize, Serialize};
 use std::default::Default;
+
 use super::lexer::*;
 
 pub trait Model {
-    fn search_query(&self, query: &[char]) -> Result<Vec<(PathBuf, f32)>, ()>;
+    fn search_query(&self, query: &[char], model: &InMemoryModel) -> Result<Vec<(PathBuf, f32)>, ()>;
     fn add_document(&mut self, path: PathBuf, content: &[char], last_modified: SystemTime) -> Result<(), ()>;
     fn requires_reindexing(&mut self, path: &Path, last_modified: SystemTime) -> Result<bool, ()>;
 }
@@ -19,7 +20,7 @@ pub trait Model {
 pub type FreqTable = HashMap::<String, usize>;  
 
 /* Answers how frequently a term occurs in all documents. 
-   Map of term with frequency of occurence in all corpus of documents.*/
+   Map of term with frequency of occurence in all corpus of documents. */
 pub type GlobalTermFreq = HashMap::<String, usize>;
 
 #[derive(Serialize, Deserialize)]
@@ -37,16 +38,51 @@ pub struct InMemoryModel {
     pub docs: Docs, 
 }
 
-fn compute_tf(term: &str, doc: &Doc) -> f32 {
-    let n = doc.ft.get(term).cloned().unwrap_or(0) as f32;    // Number of times term occured in document
-    let d = doc.count as f32;                                          // Total number of terms present in document 
-    n / d   
+fn compute_avgdl(model: &InMemoryModel) -> f32 {
+    let total: usize = model.docs.values().map(|doc| doc.count).sum();
+    total as f32 / model.docs.len() as f32
 }
 
 fn compute_idf(term: &str, model: &InMemoryModel) -> f32 {
-    let n = model.docs.len() as f32;                                    // Number of documents in corpus
-    // NOTE: Can lead to division by 0 if term is not in Document Corpus, Set Denominator to 1 if turns to 0
-    let d  = model.gtf.get(term).cloned().unwrap_or(1) as f32;  // Number of times term occured in document corpus
+    // Number of documents in corpus
+    let total_docs = model.docs.len() as f32;                           
+
+    // Number of documents containing the unique 'term'
+    let doc_freq = model.docs.values().filter(|doc| doc.ft.contains_key(term)).count() as f32;
+
+    f32::ln(((total_docs - doc_freq + 0.5) + 1f32) / (doc_freq + 0.5))
+}
+
+// K and B are free parameters, usually chosen, in absence of an advanced optimization, as K = [1.2, 2.0] and B = 0.75
+const K: f32 = 2.0;
+const B: f32 = 0.75;
+fn bm25_score(query: &Vec<String>, doc: &Doc, model: &InMemoryModel) -> f32 {
+    // Ranking documents according to BM25 Algorithm: https://en.wikipedia.org/wiki/Okapi_BM25
+    let mut score= 0f32;
+    let avgdl = compute_avgdl(model);
+    let doc_length = doc.count as f32;
+
+    for term in query {
+        // Number of times a term occurs in the document
+        let tf = doc.ft.get(term).copied().unwrap_or(0) as f32;
+        let idf = compute_idf(&term, model);
+
+        let denom = tf + K * (1f32 - B + B * doc_length / avgdl);
+        score += idf * tf * (K + 1f32) / denom;
+    }
+    score
+}
+
+// For TF-IDF Ranking 
+fn tf(term: &str, doc: &Doc) -> f32 {
+    let n = doc.ft.get(term).cloned().unwrap_or(0) as f32;     // Number of times term occured in document
+    let d = doc.count as f32;                                           // Total number of terms present in document 
+    n / d   
+}
+// For TF-IDF Ranking
+fn idf(term: &str, model: &InMemoryModel) -> f32 {
+    let n = model.docs.len() as f32;
+    let d  = model.gtf.get(term).cloned().unwrap_or(1) as f32;
     f32::log10(n / d)
 }
 
@@ -57,30 +93,35 @@ impl InMemoryModel {
             for term in doc.ft.keys() {
                 // Update the GlobalTermFrequency table
                 if let Some(freq) = self.gtf.get_mut(term) {
-                    *freq -= 1;
+                    *freq -= freq.saturating_sub(1);
                 }
             }
         }
     }
 }
 
+static USE_BM25: bool = true;
 impl Model for InMemoryModel {
-    fn search_query(&self, query: &[char]) -> Result<Vec<(PathBuf, f32)>, ()> {
-        let mut results = Vec::new();
-        // Cache all the tokens and don't retokenize on each query 
+    fn search_query(&self, query: &[char], model: &InMemoryModel) -> Result<Vec<(PathBuf, f32)>, ()> {
         let tokens = Lexer::new(&query).collect::<Vec<_>>();
+        
+        let mut results = Vec::with_capacity(self.docs.len());
         for (path, doc) in &self.docs {
-            let mut rank = 0f32;   
-            for token in &tokens {
-                // Rank is value of tf-idf => tf * idf
-                rank += compute_tf(&token, &doc) * compute_idf(&token, &self);
-            }
+            let rank = if USE_BM25 {
+                // BM-25 Ranking
+                bm25_score(&tokens, doc, model)
+            } else {
+                // TF-IDF Ranking
+                tokens.iter()
+                    .map(|token| tf(token, doc) * idf(token, model))
+                    .sum()
+            };
+
             results.push((path.to_owned(), rank));
         }
 
         // Rank the files in desc order
-        results.sort_by(|(_, ra), (_, rb)| ra.partial_cmp(rb).expect("Compared with NaN values"));
-        results.reverse();
+        results.sort_by(|(_, ra), (_, rb)| rb.partial_cmp(ra).expect("Compared with NaN values"));
         Ok(results)
     }
 
@@ -92,7 +133,7 @@ impl Model for InMemoryModel {
         let tokens = Lexer::new(&content).collect::<Vec<_>>();
         let mut ft = FreqTable::new();
         for token in &tokens {
-            ft.entry(token.to_owned()).and_modify(|x| *x += 1).or_insert(1);
+            ft.entry(token.clone()).and_modify(|x| *x += 1).or_insert(1);
         }
 
         // Total count of terms in FreqTable 
@@ -116,4 +157,5 @@ impl Model for InMemoryModel {
     }
 }
 
+// TODO: BM25 is very slow, optimize it 
 // TODO: Implement a efficient sqlite Model with parellel processing support

src/server.rs

@@ -63,7 +63,7 @@ pub fn serve_api_search(mut request: Request, model: Arc<Mutex<InMemoryModel>>)
     println!("Recieved Query: \'{}\'", body.iter().collect::<String>().bright_blue());
 
     let model = model.lock().unwrap();
-    let results = match model.search_query(&body) {
+    let results = match model.search_query(&body, &model) {
         Ok(results) => results, 
         Err(()) => return serve_500(request)
     };