Support fallback to bm25

Author: EricLBuehler

Cargo.lock

@@ -95,6 +95,7 @@ parking_lot = "0.12.3"
 ahash = "0.8.12"
 num-traits = "0.2.19"
 libc = "0.2.172"
+bm25 = "2.2.1"
 
 [features]
 pyo3_macros = ["pyo3"]

mistralrs-core/Cargo.toml

@@ -33,7 +33,6 @@ impl Engine {
                     request.messages,
                     RequestMessage::Chat { .. } | RequestMessage::VisionChat { .. }
                 ) && request.web_search_options.is_some()
-                    && get_mut_arcmutex!(self.bert_pipeline).is_some()
                 {
                     search_request::search_request(self.clone(), *request).await;
                 } else {

mistralrs-core/src/engine/add_request.rs

@@ -1,5 +1,6 @@
 use std::{borrow::Cow, sync::Arc, time::Instant};
 
+use bm25::{Embedder, EmbedderBuilder, Language, ScoredDocument, Scorer};
 use either::Either;
 use indexmap::IndexMap;
 use tokenizers::InputSequence;
@@ -98,28 +99,76 @@ async fn do_search(
         results.sort_by_key(|(_, len)| *len);
 
         {
-            let device = get_mut_arcmutex!(this.pipeline).device();
+            // Determine ranking: use embedding model if available, otherwise fallback to BM25
+            let decreasing_indexes: Vec<usize> = if let Some(bert_pipeline) =
+                &mut *get_mut_arcmutex!(this.bert_pipeline)
+            {
+                // Semantic reranking with embeddings
+                let device = get_mut_arcmutex!(this.pipeline).device();
+                search::rag::compute_most_similar(
+                    &device,
+                    &tool_call_params.query,
+                    results.iter().map(|(res, _)| res).collect::<Vec<_>>(),
+                    bert_pipeline,
+                )
+                .unwrap()
+            } else {
+                tracing::warn!("No embedding model loaded; falling back to BM25 ranking for web search results.");
 
-            let Some(bert_pipeline) = &mut *get_mut_arcmutex!(this.bert_pipeline) else {
-                unreachable!()
-            };
+                // Build an Embedder over the corpus, fitting to the entire set of documents.
+                //    - Language::English is chosen here
+                //    - This computes an in‑memory sparse embedding for each document.
+
+                let docs: Vec<String> =
+                    results.iter().map(|(res, _)| res.content.clone()).collect();
+                let doc_refs: Vec<&str> = docs.iter().map(|s| s.as_str()).collect();
+
+                let embedder: Embedder =
+                    EmbedderBuilder::with_fit_to_corpus(Language::English, &doc_refs).build();
+
+                // Initialize a Scorer keyed by usize (document index type).
+                let mut scorer = Scorer::<usize>::new();
 
-            let decreasing_indexes = search::rag::compute_most_similar(
-                &device,
-                &tool_call_params.query,
-                results.iter().map(|(res, _)| res).collect::<Vec<_>>(),
-                bert_pipeline,
-            )
-            .unwrap();
-
-            // Rerank the results
-            let mut results_old = Vec::new();
-            std::mem::swap(&mut results_old, &mut results);
-            for &index in &decreasing_indexes {
-                let mut current_result: (SearchResult, usize) = Default::default();
-                std::mem::swap(&mut current_result, &mut results_old[index]);
-
-                results.push(current_result);
+                // For each document, compute its embedding and upsert into the scorer.
+                for (i, doc_text) in docs.iter().enumerate() {
+                    let doc_embedding = embedder.embed(doc_text);
+                    scorer.upsert(&i, doc_embedding);
+                }
+
+                // Embed the query string into the same sparse embedding space.
+                let query_embedding = embedder.embed(&tool_call_params.query);
+
+                // Score all documents individually
+                let mut scored_docs: Vec<ScoredDocument<usize>> = docs
+                    .iter()
+                    .enumerate()
+                    .filter_map(|(i, _)| {
+                        scorer
+                            .score(&i, &query_embedding)
+                            .map(|score| ScoredDocument { id: i, score })
+                    })
+                    .collect();
+
+                // Sort the scored documents by descending `score` (f32).
+                scored_docs.sort_by(|a, b| {
+                    b.score
+                        .partial_cmp(&a.score)
+                        .unwrap_or(std::cmp::Ordering::Equal)
+                });
+
+                // Extract only the document indices (usize) in ranked order.
+                let decreasing_indexes: Vec<usize> =
+                    scored_docs.into_iter().map(|d| d.id).collect();
+
+                decreasing_indexes
+            };
+            // Reorder results according to ranking
+            let mut old = Vec::new();
+            std::mem::swap(&mut old, &mut results);
+            for &idx in &decreasing_indexes {
+                let mut item: (SearchResult, usize) = Default::default();
+                std::mem::swap(&mut item, &mut old[idx]);
+                results.push(item);
             }
         }