From 7bbf5bd2cbba8346447427d2b34cf21b1672a22e Mon Sep 17 00:00:00 2001 From: KShivendu Date: Sat, 12 Sep 2026 07:31:14 +0530 Subject: [PATCH 1/2] feat: add inference-free SPLADE sparse model Adds `SparseModel::OpenSearchNeuralSparseDocV3Gte`, a port of the inference-free SPLADE support that landed in the Python fastembed (qdrant/fastembed#652), producing numerically identical vectors. The model is asymmetric: documents are expanded by the ONNX encoder, while queries are embedded from the tokenizer and the `idf.json` table shipped with the model, without touching the session. `try_new` now reads back a downloaded `idf.json` into an IDF lookup, and the new `query_embed` takes `&self` since it runs no inference. Document post-processing uses the double log activation of the v3 opensearch-neural-sparse family, `ln(1 + ln(1 + relu(x)))`, rather than the single `ln(1 + relu(x))` of SPLADE++, and max-pools before the ReLU to match the Python ordering. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/test.yml | 1 + README.md | 25 ++++ src/models/sparse.rs | 24 +++- src/sparse_text_embedding/impl.rs | 182 +++++++++++++++++++++++++++++- src/sparse_text_embedding/init.rs | 7 ++ tests/if_splade.rs | 76 +++++++++++++ tests/text-embeddings.rs | 11 ++ 7 files changed, 320 insertions(+), 6 deletions(-) create mode 100644 tests/if_splade.rs diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 6965e69..8c39a29 100644 --- a/.github/workflows/test.yml +++ b/.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" } diff --git a/README.md b/README.md index 94f47fd..25c4abd 100644 --- a/README.md +++ b/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) @@ -164,6 +165,30 @@ let documents = vec![ let embeddings: Vec = 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), +)?; + +// Documents run through the encoder. Keep the batch size small: this model emits one score per +// vocabulary entry per token, so the default batch size of 256 would allocate ~16 GB at once. +let documents = model.embed(vec!["Hello World"], Some(4))?; + +// Queries only need a shared reference, no session is touched +let queries = model.query_embed(vec!["Hello World"])?; +``` + +`query_embed` returns an error for the symmetric models, which have no separate query +representation. + ### Image Embeddings ```rust diff --git a/src/models/sparse.rs b/src/models/sparse.rs index bcabcb6..f4d5638 100644 --- a/src/models/sparse.rs +++ b/src/models/sparse.rs @@ -2,6 +2,9 @@ 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 @@ -9,6 +12,8 @@ pub enum SparseModel { SPLADEPPV1, /// BAAI/bge-m3 BGEM3, + /// opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte + OpenSearchNeuralSparseDocV3Gte, } pub fn models_list() -> Vec> { @@ -36,6 +41,18 @@ pub fn models_list() -> Vec> { ], 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 { match m { SparseModel::SPLADEPPV1 => (), SparseModel::BGEM3 => (), + SparseModel::OpenSearchNeuralSparseDocV3Gte => (), } } - vec![SparseModel::SPLADEPPV1, SparseModel::BGEM3] + vec![ + SparseModel::SPLADEPPV1, + SparseModel::BGEM3, + SparseModel::OpenSearchNeuralSparseDocV3Gte, + ] } #[cfg(test)] diff --git a/src/sparse_text_embedding/impl.rs b/src/sparse_text_embedding/impl.rs index 000c615..66ee10c 100644 --- a/src/sparse_text_embedding/impl.rs +++ b/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,12 +57,16 @@ impl SparseTextEmbedding { })?; // Download additional files if needed (e.g., model.onnx.data for large models) + let mut idf_file_reference: Option = 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); + } } } @@ -68,23 +74,59 @@ impl SparseTextEmbedding { .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>, + ) -> 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> { + let token_to_idf: HashMap = 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: "".into(), + })?, + _ => "logits", + }; + + let (shape, data) = outputs[logits_key] + .try_extract_tensor::() + .map_err(|e| Error::TensorExtraction(e.to_string()))?; + let shape: Vec = 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,62 @@ 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`]: a query is tokenized, and each of its + /// unique non-special tokens is assigned the IDF weight shipped with the model. Documents + /// still have to go through [`SparseTextEmbedding::embed`], and the two are compared with a + /// dot product. + /// + /// Takes `&self` rather than `&mut self` because the ONNX session is never touched. + /// + /// Accepts anything that can be referenced as a slice of elements implementing + /// [`AsRef`], such as `Vec`, `Vec<&str>`, `&[String]`, or `&[&str]`. + pub fn query_embed + Send + Sync>( + &self, + texts: impl AsRef<[S]>, + ) -> Result> { + 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 = 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, attention_mask: &CowArray>, @@ -327,4 +449,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, + attention_mask: &Array>, + special_token_ids: &HashSet, + ) -> Vec { + 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 = Vec::new(); + let mut indices: Vec = 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() + } } diff --git a/src/sparse_text_embedding/init.rs b/src/sparse_text_embedding/init.rs index 2992954..3db4f49 100644 --- a/src/sparse_text_embedding/init.rs +++ b/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, + /// 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>, } diff --git a/tests/if_splade.rs b/tests/if_splade.rs new file mode 100644 index 0000000..3b8f07c --- /dev/null +++ b/tests/if_splade.rs @@ -0,0 +1,76 @@ +#![cfg(feature = "hf-hub")] + +use fastembed::{SparseInitOptions, SparseModel, SparseTextEmbedding}; + +const EPS: f32 = 1e-3; + +/// The MLM head emits a `vocab_size` wide tensor per token, so batches have to stay small: +/// the default batch size of 256 would allocate `256 * 512 * 30522 * 4` bytes at once. +const BATCH_SIZE: usize = 4; + +#[test] +fn test_if_splade_embeddings_match_python() { + let mut model = SparseTextEmbedding::try_new(SparseInitOptions::new( + SparseModel::OpenSearchNeuralSparseDocV3Gte, + )) + .expect("Failed to initialize the inference-free SPLADE model"); + + let documents = vec!["Hello World"]; + + // Expected values from Python + // from fastembed import SparseTextEmbedding + // model = SparseTextEmbedding("opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte") + // The document embedding has many more non-zero dimensions, these are the leading ones. + let expected_document_indices = [ + 999, 1010, 1011, 1024, 1028, 1029, 1045, 1074, 1993, 2017, 2033, 2054, 2073, 2080, 2088, + ]; + let expected_document_values = [ + 0.16544909, 0.00529129, 0.0392109, 0.12337475, 0.09640586, 0.05325737, 0.09611791, + 0.03159865, 0.01349991, 0.09392473, 0.01928805, 0.05238346, 0.05515401, 0.03156782, + 0.98263124, + ]; + let expected_query_indices = [2088, 7592]; + let expected_query_values = [3.42086864, 6.93775654]; + + let embeddings = model + .embed(documents.clone(), Some(BATCH_SIZE)) + .expect("Embedding failed"); + + assert_eq!(embeddings.len(), documents.len()); + let document = &embeddings[0]; + assert_eq!(document.indices.len(), document.values.len()); + assert!(document.indices.len() > expected_document_indices.len()); + assert_eq!( + document.indices[..expected_document_indices.len()], + expected_document_indices + ); + for (i, expected) in expected_document_values.iter().enumerate() { + assert!( + (document.values[i] - expected).abs() < EPS, + "dimension {} is {}, expected {expected}", + document.indices[i], + document.values[i], + ); + } + + // Queries are embedded from the tokenizer and the IDF table alone. `query_embed` takes + // `&self` while running the ONNX session needs `&mut self`, so a shared borrow of the model + // is enough to prove that no inference happens on the query side. + let model: &SparseTextEmbedding = &model; + let query_embeddings = model + .query_embed(documents) + .expect("Query embedding failed"); + + assert_eq!(query_embeddings.len(), 1); + let query = &query_embeddings[0]; + assert_eq!(query.indices, expected_query_indices); + assert_eq!(query.values.len(), expected_query_values.len()); + for (i, expected) in expected_query_values.iter().enumerate() { + assert!( + (query.values[i] - expected).abs() < EPS, + "dimension {} is {}, expected {expected}", + query.indices[i], + query.values[i], + ); + } +} diff --git a/tests/text-embeddings.rs b/tests/text-embeddings.rs index b8c42c5..edac5ae 100644 --- a/tests/text-embeddings.rs +++ b/tests/text-embeddings.rs @@ -180,6 +180,14 @@ create_embeddings_test!( fn test_sparse_embeddings() { SparseTextEmbedding::list_supported_models() .iter() + // The inference-free models produce far denser documents and have a query path of their + // own, they are covered by the `if_splade` suite instead. + .filter(|supported_model| { + supported_model + .additional_files + .iter() + .all(|f| f != "idf.json") + }) .for_each(|supported_model| { let mut model: SparseTextEmbedding = SparseTextEmbedding::try_new(SparseInitOptions::new(supported_model.model.clone())) @@ -202,6 +210,9 @@ fn test_sparse_embeddings() { assert_eq!(embedding.indices.len(), embedding.values.len()); }); + // Symmetric models have no separate query representation + assert!(model.query_embed(documents).is_err()); + // Clear the model cache to avoid running out of space on GitHub Actions. if std::env::var("CI").is_ok() { clean_cache(supported_model.model_code.clone()) From d2360db3353ed93d07252b35da6535557da9f21c Mon Sep 17 00:00:00 2001 From: Anush Date: Sat, 12 Sep 2026 22:58:34 +0530 Subject: [PATCH 2/2] Apply batched suggestions from code review Co-authored-by: Anush --- README.md | 9 +++------ src/sparse_text_embedding/impl.rs | 7 +------ 2 files changed, 4 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 25c4abd..944697f 100644 --- a/README.md +++ b/README.md @@ -175,18 +175,15 @@ table alone, without any inference. Both sides are compared with a dot product. use fastembed::{SparseInitOptions, SparseModel, SparseTextEmbedding}; let mut model = SparseTextEmbedding::try_new( - SparseInitOptions::new(SparseModel::OpenSearchNeuralSparseDocV3Gte), + SparseInitOptions::new(SparseModel::OpenSearchNeuralSparseDocV3Gte).with_max_length(8192), )?; -// Documents run through the encoder. Keep the batch size small: this model emits one score per -// vocabulary entry per token, so the default batch size of 256 would allocate ~16 GB at once. +// 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))?; -// Queries only need a shared reference, no session is touched let queries = model.query_embed(vec!["Hello World"])?; ``` - -`query_embed` returns an error for the symmetric models, which have no separate query representation. ### Image Embeddings diff --git a/src/sparse_text_embedding/impl.rs b/src/sparse_text_embedding/impl.rs index 66ee10c..ffd2fd8 100644 --- a/src/sparse_text_embedding/impl.rs +++ b/src/sparse_text_embedding/impl.rs @@ -316,12 +316,7 @@ impl SparseTextEmbedding { /// Method to generate sparse query embeddings without running any model inference. /// /// Only available for the inference-free (asymmetric) models, such as - /// [`SparseModel::OpenSearchNeuralSparseDocV3Gte`]: a query is tokenized, and each of its - /// unique non-special tokens is assigned the IDF weight shipped with the model. Documents - /// still have to go through [`SparseTextEmbedding::embed`], and the two are compared with a - /// dot product. - /// - /// Takes `&self` rather than `&mut self` because the ONNX session is never touched. + /// [`SparseModel::OpenSearchNeuralSparseDocV3Gte`]. /// /// Accepts anything that can be referenced as a slice of elements implementing /// [`AsRef`], such as `Vec`, `Vec<&str>`, `&[String]`, or `&[&str]`.