naver

/

splade-code-06B

like 5

Follow

NAVER LABS Europe 183

Hello!

Congratulations on the release! I believed it was possible to both simplify and integrate with Sentence Transformers, so here is my attempt.

Pull Request overview

• Simplify the bidirectional Qwen3 implementation heavily by relying on https://github.com/huggingface/transformers/pull/43705 (requires transformers v5.2.0+)
• Support SparseEncoder from Sentence Transformers, matches original implementation outputs

Details

The https://github.com/huggingface/transformers/pull/43705 pull request in transformers starts supporting the is_causal parameter on all architectures, meaning that turning a causal model bidirectional becomes as simple as setting "is_causal": falsein the config.json. This means that we don't even need the custom modeling code. However, sadly trust_remote_code=Trueis still required because we're loading a CausalLM model using AutoModelForMaskedLM. The implementation can be very simple now, e.g. with:

"""
This file exists solely to allow loading the Qwen3ForCausalLM via the AutoModelForMaskedLM class.
Compared to standard Qwen3, we're using bidirectional attention and not causal attention, but it's specified
with `is_causal=False` in the config.
"""
from transformers import Qwen3ForCausalLM

__all__ = ["Qwen3ForCausalLM"]

Although you can also extend this again to re-add the encode etc. methods.

Instead, I kept it simple for now and just used a SparseEncoder implementation instead. That just required some configuration files. I also copied the Qwen3-0.6B config.json as that's just loaded directly now instead of loading some custom code which itself loads a Qwen3 model. I also copied the Qwen3 tokenizer files.

The SparseEncoder usage is now:

from sentence_transformers import SparseEncoder

model = SparseEncoder("naver/splade-code-06B", trust_remote_code=True, revision="refs/pr/1")

queries = [
    "SELECT *\nFROM Student\nWHERE Age = (\nSELECT MAX(Age)\nFROM Student\nWHERE Group = 'specific_group'\n)\nAND Group = 'specific_group';"
]

query_embeddings = model.encode(queries)
print(query_embeddings.shape)
# torch.Size([1, 151936])

sparsity = model.sparsity(query_embeddings)
print(sparsity)
# {'active_dims': 1231.0, 'sparsity_ratio': 0.991897904380792}

decoded = model.decode(query_embeddings, top_k=10)
print(decoded)
# [[
#     ("Ġgroup", 2.34375),
#     ("Ġage", 2.34375),
#     ("ĠAge", 2.34375),
#     ("ĠStudent", 2.296875),
#     ("Ġspecific", 2.296875),
#     ("_group", 2.296875),
#     ("ĠMax", 2.21875),
#     ("Ġmax", 2.21875),
#     ("Ġstudent", 2.203125),
#     ("ĠGroup", 2.1875),
# ]]

And it works with transformers>5.2.0 and sentence-transformers>5.0.0. If you install kernels, you can set model_kwargs={"attn_implementation": "flash_attention_2"} and it will use a kernel from the Hub without having to actually install flash-attn (which is always annoying).

It also works with e.g. sdpa or eager, unlike the current implementation, and is likely to keep working with future transformers versions as it imports very little from transformers.

Note that the above script has revision="refs/pr/1" so you can test it directly from this PR branch without having to check anything out locally or merge it.

I believe the SparseEncoder is also integrated in MTEB, so you can evaluate this easily on MTEB for submission. In Sentence Transformers v5.4.0 I suspect that you'll also be able to avoid the trust_remote_code=True, but I'm not 100% sure on that yet.

Normally my integration PRs aren't this big, and the original code generally still works (but now it doesn't), so I totally understand if you're not as interested. Let me know your thoughts! Perhaps there's a nice middle-ground where both Sentence Transformers and regular transformers still works. I can also copy these changes to the 8B model if you're interested.

• Tom Aarsen