feature: support bm42 embeddings

Author: cdxker

Cargo.lock

@@ -17,6 +17,7 @@ tokio = { workspace = true }
 tracing = { workspace = true }
 
 [features]
+default = ["ort"]
 clap = ["dep:clap", "text-embeddings-backend-core/clap"]
 python = ["dep:text-embeddings-backend-python"]
 ort = ["dep:text-embeddings-backend-ort"]

backends/Cargo.toml

@@ -856,6 +856,7 @@ impl BertModel {
 
                     (outputs.sum(1)?.broadcast_div(&input_lengths))?
                 }
+                Pool::BM42 => unreachable!(),
                 Pool::Splade => {
                     // Unwrap is safe here
                     let splade_head = self.splade.as_ref().unwrap();
@@ -874,7 +875,7 @@ impl BertModel {
                     }
 
                     relu_log.max(1)?
-                }
+                },
             };
             Some(pooled_embeddings)
         } else {

backends/candle/src/models/bert.rs

@@ -587,6 +587,7 @@ impl DistilBertModel {
 
                     (outputs.sum(1)?.broadcast_div(&input_lengths))?
                 }
+                Pool::BM42 => unreachable!(),
                 Pool::Splade => {
                     // Unwrap is safe here
                     let splade_head = self.splade.as_ref().unwrap();

backends/candle/src/models/distilbert.rs

@@ -616,6 +616,7 @@ impl JinaBertModel {
 
                     (outputs.sum(1)?.broadcast_div(&input_lengths))?
                 }
+                Pool::BM42 => unreachable!(),
                 Pool::Splade => unreachable!(),
             };
             Some(pooled_embeddings)

backends/candle/src/models/jina.rs

@@ -604,7 +604,8 @@ impl JinaCodeBertModel {
                     }
 
                     (outputs.sum(1)?.broadcast_div(&input_lengths))?
-                }
+                },
+                Pool::BM42 => unreachable!(),
                 Pool::Splade => unreachable!(),
             };
             Some(pooled_embeddings)

backends/candle/src/models/jina_code.rs

@@ -631,6 +631,7 @@ impl NomicBertModel {
 
                     (outputs.sum(1)?.broadcast_div(&input_lengths))?
                 }
+                Pool::BM42 => unreachable!(),
                 Pool::Splade => unreachable!(),
             };
             Some(pooled_embeddings)

backends/candle/src/models/nomic.rs

@@ -63,16 +63,34 @@ pub enum Pool {
     /// This option is only available if the loaded model is a `ForMaskedLM` Transformer
     /// model.
     Splade,
+    /// Apply BM42 to the model embeddings.
+    /// This option is only availale if the loaded model is Qdrant/all_miniLM_L6_v2_with_attentions 
+    BM42,
     /// Select the last token as embedding
     LastToken,
 }
 
+#[derive(Debug, Clone)]
+pub struct Bm42Params {
+    pub invert_vocab: std::collections::HashMap<u32, String>,
+    pub stopwords: Vec<String>,
+    pub special_tokens: Vec<String>,
+}
+
+#[derive(Debug, Clone)]
+pub enum ModelParams {
+    Bm42(Bm42Params),
+    None
+}
+
+
 impl fmt::Display for Pool {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             Pool::Cls => write!(f, "cls"),
             Pool::Mean => write!(f, "mean"),
             Pool::Splade => write!(f, "splade"),
+            Pool::BM42  => write!(f, "bm42"),
             Pool::LastToken => write!(f, "last_token"),
         }
     }

backends/core/src/lib.rs

@@ -15,3 +15,5 @@ tracing = { workspace = true }
 thiserror = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
+rust-stemmers = "1.2.0"
+murmur3 = "0.5.2"