Integrate with Sentence Transformers

#2
by tomaarsen HF Staff - opened

Hello!

Preface

Congrats on the very strong results on turning SPLADE multimodal! I'm quite surprised at how well it does on ViDoRe etc., looks very impressive.
Nowadays, Sentence Transformers has multimodality support, even for Sparse models, even if they're inference-free. I've worked with an agent of mine to integrate this model with Sentence Transformers, which should simplify usage a lot. Please let me know what you think!

I've double-checked the outputs of the work compared to your original implementation, and it looks to match. Here's the PR description as produced by the agent. Note that I've updated the snippet to use revision. This means that you can already use it out of the box for testing once you have pip install -U sentence-transformers[image], without having to e.g. check out this PR branch locally or something.

Pull Request overview

  • Integrate this model with Sentence Transformers (v5.6.0 or newer): SparseEncoder("naver/v-splade-quality", trust_remote_code=True) runs text queries through the inference-free Li-LSR lookup and page images (or plain text) through the VBERT document encoder.

Details

The integration is a Router with two routes, mirroring the asymmetric design:

  • query: VSPLADEStaticEmbedding (modeling_st_vsplade.py), a SparseStaticEmbedding subclass over the precomputed lookup table softplus(projection(embedding)) (float32, ~200 KB, special tokens zeroed; the source tensors in model.safetensors are untouched). It reproduces InferenceFreeQueryEncoder.encode_with_lookup exactly: repeated query tokens accumulate via scatter-add, and the 40 added vision-token ids contribute nothing.
  • document: Transformer(transformer_task="fill-mask") -> SpladePooling(max), resolving through a new auto_map entry to VSPLADEForMaskedLM (modeling_vsplade.py): the native ModernVBERT backbone (transformers>=5.3.0) plus the V-SPLADE MLM head, with a module tree that mirrors the checkpoint layout so model.safetensors loads without any key remapping. Its forward returns the SPLADE term logits (scaled by hidden_size ** -0.25, special tokens zeroed), exactly matching UnifiedRetriever._apply_sparse_head. As a side effect, plain AutoModelForMaskedLM.from_pretrained(..., trust_remote_code=True) document encoding now works too. The "" module id in router_config.json makes this module load from the repository root, so no weights are duplicated into a subfolder.

Chat templates: the legacy chat_template.json became the equivalent modern chat_template.jinja (transformers raises when the legacy file coexists with an additional_chat_templates/ directory), and the new additional_chat_templates/sentence_transformers.jinja renders image inputs to the vsplade_inference.py format (User:<image><end_of_utterance>\nAssistant:) and text-only inputs as raw text (the plain RLHN passage format). The reference examples/quickstart.py flow was re-run after the conversion and is unchanged.

Verification against the reference implementation (the pinned transformers==4.57.6 + colpali-engine stack, float32 on CPU): query embeddings are bit-identical, document embeddings and scores match within 8e-06 and 2e-06, the README quickstart output is reproduced (nnz=552, top tokens dog, dogs, puppy, Records, Bennett), batches mixing images with different tile counts match their solo encodings, and a config-only load on the released sentence-transformers==5.6.0 + transformers==5.13.0 passes as well. Nothing the reference loaders read was modified, so the github.com/naver/v-splade inference and training code paths are unaffected.

Added files:

  • modules.json, router_config.json, sentence_bert_config.json, config_sentence_transformers.json: the Sentence Transformers pipeline configuration.
  • modeling_vsplade.py and modeling_st_vsplade.py: the document encoder and query module described above.
  • query_0_VSPLADEStaticEmbedding/: the lookup table, module config, and tokenizer copy.
  • document_1_SpladePooling/config.json: SPLADE max-pooling config.
  • chat_template.jinja and additional_chat_templates/sentence_transformers.jinja.

Modified files:

  • config.json: added the auto_map entry.
  • README.md: added the sentence-transformers tag and a "Using Sentence Transformers" section; the existing quickstart moved under a "Using the reference implementation" subsection, otherwise unchanged.

Removed files:

  • chat_template.json: replaced by chat_template.jinja using the exact same chat template.
  • modeling_modernvbert.py, configuration_modernvbert.py: reference copies that were never loadable from the repository (package-relative imports, no auto_map pointed at them, and the github.com/naver/v-splade code uses colpali-engine's vendored implementation). With the architecture native in transformers>=5.3.0, they only risked confusion; easy to restore if you would rather keep them.

Usage, as added to the README:

from sentence_transformers import SparseEncoder

model = SparseEncoder("naver/v-splade-quality", trust_remote_code=True, revision="refs/pr/2")

queries = ["send signed forms", "records office"]
documents = ["https://raw.githubusercontent.com/naver/v-splade/main/examples/sample_page.png"]

query_embeddings = model.encode_query(queries)
document_embeddings = model.encode_document(documents)
print(query_embeddings.shape, document_embeddings.shape)
# torch.Size([2, 50368]) torch.Size([1, 50368])

similarities = model.similarity(query_embeddings, document_embeddings)
print(similarities)
# tensor([[0.9843],
#         [0.5935]], device='cuda:0')

# Inspect the top activated tokens of the page image
decoded = model.decode(document_embeddings[0], top_k=5)
print([(token, round(weight, 3)) for token, weight in decoded])
# [('Ġdog', 1.836), ('Ġdogs', 1.664), ('Ġpuppy', 1.586), ('ĠRecords', 1.578), ('ĠBennett', 1.508)]

One optional follow-up to consider: switching library_name in the README metadata from transformers to sentence-transformers, so the Hub widget shows the retrieval-ready snippet by default. Both libraries can load the model either way, so this PR leaves it as is.

Happy to tweak anything you'd like changed, including porting the same integration to naver/v-splade-efficient. Please let me know if you have any questions or feedback!

  • Tom Aarsen
gyuhwung-cho changed pull request status to open
gyuhwung-cho changed pull request status to merged

Dear @tomaarsen

Thank you so much for porting our model to Sentence Transformers .

We really appreciate that this makes V-SPLADE much easier to use through the Sentence Transformers interface.

We added a few minor updates and have merged the PR for both naver/v-splade-quality and naver/v-splade-efficient.

Thanks again for your great contribution!

Best regards,
Gyu-Hwung Cho

Sign up or log in to comment