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

[KShivendu]

Adds SparseModel::OpenSearchNeuralSparseDocV3Gte, an inference-free SPLADE model, ported from the Python fastembed implementation in qdrant/fastembed#652 so the two libraries produce identical vectors.

The model

opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte, mirrored for fastembed at Qdrant/opensearch-neural-sparse-encoding-doc-v3-gte (~0.55 GB, single-file ONNX, apache-2.0). It is an MLM head over a GTE backbone: input_ids + attention_mask in (type_vocab_size is 0, so there is no token_type_ids input), one logits tensor of (batch, seq_len, 30522) out.

It is asymmetric, which is the whole point of it:

• Documents run through the ONNX encoder, once, at index time.
• Queries are a tokenizer pass plus an IDF lookup. No inference at all.

That makes query latency essentially free while keeping the term-expansion quality of SPLADE on the document side.

What changed

• src/models/sparse.rs — new variant with model_code: "Qdrant/opensearch-neural-sparse-encoding-doc-v3-gte", model_file: "model.onnx", additional_files: ["idf.json"].
• src/sparse_text_embedding/init.rs — two new fields on SparseTextEmbedding: token_id_to_idf: Option<HashMap<usize, f32>> and special_token_ids: HashSet<usize>. Both are inert for SPLADE++ and BGE-M3.
• src/sparse_text_embedding/impl.rs
  • try_new already downloaded additional_files but never read them back. It now keeps the path of the downloaded idf.json, and any model that declares one gets its IDF table loaded eagerly next to the tokenizer — no per-variant special-casing in the loader. idf.json is keyed by token string, so it is resolved through the tokenizer vocab, exactly like Python's _load_idf.
  • Special-token ids come from the tokenizer's added-tokens table (get_added_tokens_decoder(), filtered on special), which after load_tokenizer matches what Python derives from special_tokens_map.json: {0, 100, 101, 102, 103} here.
  • New post_process_if_splade arm.
  • New public query_embed.
• tests/if_splade.rs + a CI matrix entry, README model list and a short usage section.

Document post-processing — this is not post_process_splade

Reusing the existing SPLADE++ path would silently produce plausible-but-wrong vectors, in two ways.

Activation. SPLADE++ uses a single ln(1 + relu(x)). The v3 opensearch-neural-sparse models use a double one, ln(1 + ln(1 + relu(x))), to push document embeddings sparser:

values.push((1.0 + (1.0 + score).ln()).ln());

Order of operations. post_process_splade does relu → log → mask → max-pool from NEG_INFINITY. Python here masks → max-pools → relus:

pooled = np.max(output.model_output * np.expand_dims(output.attention_mask, axis=-1), axis=1)
scores = np.log1p(np.log1p(np.maximum(pooled, 0.0)))

The new arm initialises each vocabulary accumulator to 0.0 and maxes only over unmasked positions. That 0.0 encodes both the ReLU floor and the zero contribution of padding at once, so it agrees with Python whether or not the batch is padded — and it avoids materialising a second batch × seq_len × 30522 tensor for the mask multiply.

Special-token ids are then dropped from the document vector too (Python zeroes them after pooling); otherwise they would match every query.

query_embed

pub fn query_embed<S: AsRef<str> + Send + Sync>(
    &self,
    texts: impl AsRef<[S]>,
) -> Result<Vec<SparseEmbedding>>

Tokenize, drop special tokens, dedupe, sort ascending by token id, look up the IDF weight, skip ids absent from the table.

It takes &self, not &mut self like embed — a genuine ergonomic win, since a SparseTextEmbedding can now serve queries from behind a shared reference (an Arc, a request handler) without a lock. It is also a compile-time proof of the inference-free property: Session::run needs &mut self, so a &self method cannot possibly run the model. Python asserts the same thing at runtime with assert not hasattr(model.model, "model").

For SPLADE++ and BGE-M3, which have no separate query representation, it returns Error::InvalidArgument rather than silently falling back to embed. tests/text-embeddings.rs asserts that.

Parity

Verified against the real weights on the same input as the Python test (docs = ["Hello World"]), ONNX Runtime 1.24.4.

Query — identical to the Python values to all 8 published digits:

index	token	Rust	Python
2088	world	3.42086864	3.42086864
7592	hello	6.93775654	6.93775654

Document — 65 non-zero dimensions, of which the leading 15 are the ones the Python test pins:

index	Rust	Python
999	0.16544974	0.16544909
1010	0.00529340	0.00529129
1011	0.03921197	0.0392109
1024	0.12337571	0.12337475
1028	0.09640644	0.09640586
1029	0.05325796	0.05325737
1045	0.09611876	0.09611791
1074	0.03159864	0.03159865
1993	0.01349987	0.01349991
2017	0.09392501	0.09392473
2033	0.01928882	0.01928805
2054	0.05238366	0.05238346
2073	0.05515388	0.05515401
2080	0.03156827	0.03156782
2088	0.98263091	0.98263124

Largest absolute difference is 2.1e-6 (index 1010), i.e. ONNX Runtime float noise, three orders of magnitude inside the abs=0.001 tolerance the Python test uses. tests/if_splade.rs checks both sides against these goldens with a 1e-3 epsilon.

Caveat worth knowing about: batch size

This model emits one score per vocabulary entry per token. At the crate's DEFAULT_BATCH_SIZE of 256 and the default max length of 512 that is 256 × 512 × 30522 × 4 bytes ≈ 16 GB in a single output tensor. embed(docs, None) is therefore a bad idea here; the README section and the test both pass a small explicit batch size, and the test constant carries a comment saying why. I left DEFAULT_BATCH_SIZE alone since it is shared with the other sparse models — happy to add a per-model default if you would rather the footgun not exist.

One other intentional difference from Python: this crate's sparse DEFAULT_MAX_LENGTH is 512, while Python honours the model's model_max_length of 8192. Documents beyond 512 tokens will diverge unless you pass .with_max_length(8192). Left as-is to match the crate's existing behaviour for the other sparse models.

Checks

cargo fmt --all -- --check, cargo clippy (default features and --no-default-features --features image-models,ort-download-binaries-native-tls), and the test suite pass locally.

🤖 Generated with Claude Code

[Anush008]

Thanks for taking the time to contribute @KShivendu.

I've got some nit picks for you to consider.

[github-actions]

🎉 This PR is included in version 6.1.0 🎉

The release is available on:

• GitHub release
• v6.1.0

Your semantic-release bot 📦🚀