feat: inference-free SPLADE (opensearch-neural-sparse-encoding-doc-v3-gte)

.github/workflows/test.yml

@@ -51,6 +51,7 @@ jobs:
       matrix:
         suite:
           - { name: bgem3, filter: "--test bgem3 --test bgem3_comparison" }
+          - { name: if-splade, filter: "--test if_splade" }
           - { name: text,  filter: "--test text-embeddings" }
           - { name: image, filter: "--test image-embeddings" }
 

README.md

@@ -67,6 +67,7 @@ Quantized versions are also available for several models above (append `Q` to th
 
 - [**prithivida/Splade_PP_en_v1**](https://huggingface.co/prithivida/Splade_PP_en_v1) - Default
 - [**BAAI/bge-m3**](https://huggingface.co/BAAI/bge-m3)
+- [**opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte**](https://huggingface.co/Qdrant/opensearch-neural-sparse-encoding-doc-v3-gte) - Inference-free, see [Inference-free Sparse Embeddings](#inference-free-sparse-embeddings)
 
 </details>
 
@@ -164,6 +165,27 @@ let documents = vec![
 let embeddings: Vec<SparseEmbedding> = model.embed(documents, None)?;
 ```
 
+### Inference-free Sparse Embeddings
+
+`SparseModel::OpenSearchNeuralSparseDocV3Gte` is asymmetric: documents are expanded by the ONNX
+encoder, while queries are embedded by `query_embed` from the tokenizer and a precomputed IDF
+table alone, without any inference. Both sides are compared with a dot product.
+
+```rust
+use fastembed::{SparseInitOptions, SparseModel, SparseTextEmbedding};
+
+let mut model = SparseTextEmbedding::try_new(
+    SparseInitOptions::new(SparseModel::OpenSearchNeuralSparseDocV3Gte).with_max_length(8192),
+)?;
+
+//  This model emits one score per vocabulary entry per token.
+// So keep the batch size small.
+let documents = model.embed(vec!["Hello World"], Some(4))?;
+
+let queries = model.query_embed(vec!["Hello World"])?;
+```
+representation.
+
 ### Image Embeddings
 
 ```rust

src/models/sparse.rs

@@ -2,13 +2,18 @@ use std::{fmt::Display, str::FromStr};
 
 use crate::ModelInfo;
 
+/// Sidecar file with the per-token IDF weights used to embed queries without inference.
+pub(crate) const IDF_FILE: &str = "idf.json";
+
 #[derive(Default, Debug, Clone, PartialEq, Eq)]
 pub enum SparseModel {
     /// prithivida/Splade_PP_en_v1
     #[default]
     SPLADEPPV1,
     /// BAAI/bge-m3
     BGEM3,
+    /// opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte
+    OpenSearchNeuralSparseDocV3Gte,
 }
 
 pub fn models_list() -> Vec<ModelInfo<SparseModel>> {
@@ -36,6 +41,18 @@ pub fn models_list() -> Vec<ModelInfo<SparseModel>> {
             ],
             output_key: None,
         },
+        ModelInfo {
+            model: SparseModel::OpenSearchNeuralSparseDocV3Gte,
+            dim: 0,
+            description: String::from(
+                "Inference-free SPLADE model. Documents are expanded with an ONNX encoder, \
+                 queries are encoded with a tokenizer and an IDF lookup table only",
+            ),
+            model_code: String::from("Qdrant/opensearch-neural-sparse-encoding-doc-v3-gte"),
+            model_file: String::from("model.onnx"),
+            additional_files: vec![IDF_FILE.to_string()],
+            output_key: None,
+        },
     ]
 }
 
@@ -75,9 +92,14 @@ pub(crate) fn all_variants() -> Vec<SparseModel> {
         match m {
             SparseModel::SPLADEPPV1 => (),
             SparseModel::BGEM3 => (),
+            SparseModel::OpenSearchNeuralSparseDocV3Gte => (),
         }
     }
-    vec![SparseModel::SPLADEPPV1, SparseModel::BGEM3]
+    vec![
+        SparseModel::SPLADEPPV1,
+        SparseModel::BGEM3,
+        SparseModel::OpenSearchNeuralSparseDocV3Gte,
+    ]
 }
 
 #[cfg(test)]

src/sparse_text_embedding/impl.rs

@@ -1,5 +1,7 @@
 #[cfg(feature = "hf-hub")]
 use crate::common::{init_session_builder, load_tokenizer_hf_hub};
+#[cfg(feature = "hf-hub")]
+use crate::models::sparse::IDF_FILE;
 use crate::{
     common::{Error, Result},
     models::sparse::{models_list, SparseModel},
@@ -9,10 +11,10 @@ use crate::{
 use hf_hub::api::sync::ApiRepo;
 use ndarray::{Array, ArrayViewD, Axis, CowArray, Dim};
 use ort::{session::Session, value::Value};
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 #[cfg_attr(not(feature = "hf-hub"), allow(unused_imports))]
 #[cfg(feature = "hf-hub")]
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use tokenizers::Tokenizer;
 
 #[cfg(feature = "hf-hub")]
@@ -55,36 +57,76 @@ impl SparseTextEmbedding {
                 })?;
 
         // Download additional files if needed (e.g., model.onnx.data for large models)
+        let mut idf_file_reference: Option<PathBuf> = None;
         if !model_info.additional_files.is_empty() {
             for file in &model_info.additional_files {
-                model_repo.get(file).map_err(|e| Error::ModelRetrieval {
+                let reference = model_repo.get(file).map_err(|e| Error::ModelRetrieval {
                     file: file.clone(),
                     source: Box::new(e),
                 })?;
+                if file == IDF_FILE {
+                    idf_file_reference = Some(reference);
+                }
             }
         }
 
         let session = init_session_builder(execution_providers, intra_threads)?
             .commit_from_file(model_file_reference)?;
 
         let tokenizer = load_tokenizer_hf_hub(model_repo, max_length)?;
-        Ok(Self::new(tokenizer, session, model_name))
+        // Models declaring an `idf.json` embed queries from a lookup table instead of
+        // running inference, so the table is loaded up front alongside the tokenizer.
+        let token_id_to_idf = idf_file_reference
+            .map(|reference| Self::load_idf(&reference, &tokenizer))
+            .transpose()?;
+
+        Ok(Self::new(tokenizer, session, model_name, token_id_to_idf))
     }
 
     /// Private method to return an instance
     #[cfg_attr(not(feature = "hf-hub"), allow(dead_code))]
-    fn new(tokenizer: Tokenizer, session: Session, model: SparseModel) -> Self {
+    fn new(
+        tokenizer: Tokenizer,
+        session: Session,
+        model: SparseModel,
+        token_id_to_idf: Option<HashMap<usize, f32>>,
+    ) -> Self {
         let need_token_type_ids = session
             .inputs()
             .iter()
             .any(|input| input.name() == "token_type_ids");
+        let special_token_ids = tokenizer
+            .get_added_tokens_decoder()
+            .iter()
+            .filter(|(_, token)| token.special)
+            .map(|(id, _)| *id as usize)
+            .collect();
         Self {
             tokenizer,
             session,
             need_token_type_ids,
             model,
+            special_token_ids,
+            token_id_to_idf,
         }
     }
+
+    /// Read the `idf.json` sidecar, resolving its token strings to token ids via the
+    /// tokenizer's vocabulary. Tokens the tokenizer does not know about are dropped.
+    #[cfg(feature = "hf-hub")]
+    fn load_idf(idf_file: &Path, tokenizer: &Tokenizer) -> Result<HashMap<usize, f32>> {
+        let token_to_idf: HashMap<String, f32> = serde_json::from_slice(&std::fs::read(idf_file)?)
+            .map_err(|e| {
+                Error::Other(format!("Failed to parse the {IDF_FILE} of the model: {e}"))
+            })?;
+
+        let vocab = tokenizer.get_vocab(true);
+        Ok(token_to_idf
+            .into_iter()
+            .filter_map(|(token, idf)| vocab.get(&token).map(|&id| (id as usize, idf)))
+            .collect())
+    }
+
     /// Return the SparseTextEmbedding model's directory from cache or remote retrieval
     #[cfg(feature = "hf-hub")]
     fn retrieve_model(
@@ -235,6 +277,30 @@ impl SparseTextEmbedding {
                             &attention_mask_array,
                         )
                     }
+                    SparseModel::OpenSearchNeuralSparseDocV3Gte => {
+                        let logits_key = match outputs.len() {
+                            1 => outputs
+                                .keys()
+                                .next()
+                                .ok_or_else(|| Error::OutputKeyMissing {
+                                    key: "<only output>".into(),
+                                })?,
+                            _ => "logits",
+                        };
+
+                        let (shape, data) = outputs[logits_key]
+                            .try_extract_tensor::<f32>()
+                            .map_err(|e| Error::TensorExtraction(e.to_string()))?;
+                        let shape: Vec<usize> = shape.iter().map(|&d| d as usize).collect();
+                        let logits = ndarray::ArrayViewD::from_shape(shape.as_slice(), data)
+                            .map_err(|e| Error::InvalidShape(e.to_string()))?;
+
+                        Self::post_process_if_splade(
+                            &logits,
+                            &attention_mask_array,
+                            &self.special_token_ids,
+                        )
+                    }
                 };
 
                 Ok(embeddings)
@@ -247,6 +313,57 @@ impl SparseTextEmbedding {
         Ok(output)
     }
 
+    /// Method to generate sparse query embeddings without running any model inference.
+    ///
+    /// Only available for the inference-free (asymmetric) models, such as
+    /// [`SparseModel::OpenSearchNeuralSparseDocV3Gte`].
+    ///
+    /// Accepts anything that can be referenced as a slice of elements implementing
+    /// [`AsRef<str>`], such as `Vec<String>`, `Vec<&str>`, `&[String]`, or `&[&str]`.
+    pub fn query_embed<S: AsRef<str> + Send + Sync>(
+        &self,
+        texts: impl AsRef<[S]>,
+    ) -> Result<Vec<SparseEmbedding>> {
+        let token_id_to_idf = self.token_id_to_idf.as_ref().ok_or_else(|| {
+            Error::InvalidArgument(format!(
+                "{} has no IDF table and no separate query representation, use `embed` instead",
+                self.model
+            ))
+        })?;
+
+        texts
+            .as_ref()
+            .iter()
+            .map(|text| {
+                let encoding = self
+                    .tokenizer
+                    .encode(text.as_ref(), true)
+                    .map_err(|e| Error::Tokenization(format!("Failed to encode the query: {e}")))?;
+
+                // Every unique token contributes its IDF weight exactly once, ordered by token id
+                let mut token_ids: Vec<usize> = encoding
+                    .get_ids()
+                    .iter()
+                    .map(|&id| id as usize)
+                    .filter(|id| !self.special_token_ids.contains(id))
+                    .collect();
+                token_ids.sort_unstable();
+                token_ids.dedup();
+
+                let mut indices = Vec::with_capacity(token_ids.len());
+                let mut values = Vec::with_capacity(token_ids.len());
+                for token_id in token_ids {
+                    if let Some(&idf) = token_id_to_idf.get(&token_id) {
+                        indices.push(token_id);
+                        values.push(idf);
+                    }
+                }
+
+                Ok(SparseEmbedding { values, indices })
+            })
+            .collect()
+    }
+
     fn post_process_splade(
         model_output: &ArrayViewD<f32>,
         attention_mask: &CowArray<i64, Dim<[usize; 2]>>,
@@ -327,4 +444,54 @@ impl SparseTextEmbedding {
             })
             .collect()
     }
+
+    /// Post-processing for the inference-free SPLADE document encoder.
+    ///
+    /// The token logits are max-pooled over the unmasked positions and squashed with a double
+    /// log activation, `log(1 + log(1 + relu(x)))`, which the v3 models of the
+    /// opensearch-neural-sparse family use to make document embeddings sparser than the single
+    /// `log(1 + relu(x))` of SPLADE++.
+    fn post_process_if_splade(
+        logits: &ArrayViewD<f32>,
+        attention_mask: &Array<i64, Dim<[usize; 2]>>,
+        special_token_ids: &HashSet<usize>,
+    ) -> Vec<SparseEmbedding> {
+        let batch_size = attention_mask.shape()[0];
+        let seq_len = attention_mask.shape()[1];
+        let vocab_size = logits.shape()[2];
+
+        (0..batch_size)
+            .map(|batch_idx| {
+                // Starting the accumulator at `0.0` encodes both the ReLU floor and the zero
+                // contribution of the padded positions, which are skipped altogether.
+                let mut pooled = vec![0.0f32; vocab_size];
+
+                for seq_idx in 0..seq_len {
+                    if attention_mask[[batch_idx, seq_idx]] == 0 {
+                        continue;
+                    }
+
+                    let token_logits = logits.slice(ndarray::s![batch_idx, seq_idx, ..]);
+                    for (score, &logit) in pooled.iter_mut().zip(token_logits.iter()) {
+                        *score = score.max(logit);
+                    }
+                }
+
+                let mut values: Vec<f32> = Vec::new();
+                let mut indices: Vec<usize> = Vec::new();
+
+                for (token_id, &score) in pooled.iter().enumerate() {
+                    // Special tokens are dropped from the document side as well, otherwise they
+                    // would match every query
+                    if score <= 0.0 || special_token_ids.contains(&token_id) {
+                        continue;
+                    }
+                    values.push((1.0 + (1.0 + score).ln()).ln());
+                    indices.push(token_id);
+                }
+
+                SparseEmbedding { values, indices }
+            })
+            .collect()
+    }
 }

src/sparse_text_embedding/init.rs

@@ -1,4 +1,5 @@
 use ort::session::Session;
+use std::collections::{HashMap, HashSet};
 use tokenizers::Tokenizer;
 
 use crate::{
@@ -41,4 +42,10 @@ pub struct SparseTextEmbedding {
     pub(crate) session: Session,
     pub(crate) need_token_type_ids: bool,
     pub(crate) model: SparseModel,
+    /// Ids of the tokenizer's special tokens, excluded from every embedding produced by
+    /// the inference-free models.
+    pub(crate) special_token_ids: HashSet<usize>,
+    /// Token id to IDF weight, read from the `idf.json` shipped with inference-free models.
+    /// [`None`] for models which do not have a separate query representation.
+    pub(crate) token_id_to_idf: Option<HashMap<usize, f32>>,
 }