Integrate with Sentence Transformers via SparseEncoder

#1
by tomaarsen HF Staff - opened

Hello @brutusxu and team!

I hadn't heard of this model release yet until today, but today's https://huggingface.co/Linkup-Platform/linkup-sparseup-embed-v1 benchmarked against LACONIC-1b, so I went to check it out. I think the inference/usage of this model can be simplified a lot with the Sentence Transformers SparseEncoder class, so I'm pitching its integration here.

Pull Request overview

  • Add Sentence Transformers support through SparseEncoder
  • Include tokenizer files and a model card with usage examples

Details

I've added the configuration, tokenizer, and a minimal LlamaForMaskedLM shim to load LACONIC with bidirectional attention and SPLADE pooling. The example's scores match the original implementation within 2e-5 in float32.

You can try this PR with sentence-transformers>=5.4.0 and transformers>=5.2.0:

from sentence_transformers import SparseEncoder

model = SparseEncoder("utahnlp/laconic-1b", trust_remote_code=True, revision="refs/pr/1")

queries = ["What is the capital of France?", "How do plants make food?"]
documents = [
    "Paris is the capital and largest city of France.",
    "Plants use sunlight to turn carbon dioxide and water into sugars through photosynthesis.",
    "The piano is a musical instrument with a keyboard.",
]

query_embeddings = model.encode_query(queries, max_active_dims=512)
document_embeddings = model.encode_document(documents, max_active_dims=512)
print(query_embeddings.shape, document_embeddings.shape)
# torch.Size([2, 128256]) torch.Size([3, 128256])

print(model.sparsity(query_embeddings))
# {'active_dims': 83.5, 'sparsity_ratio': 0.9993489583333334}
print(model.sparsity(document_embeddings))
# {'active_dims': 512.0, 'sparsity_ratio': 0.9960079840319361}

scores = model.similarity(query_embeddings, document_embeddings)
print(scores.cpu())
# tensor([[17.1619,  0.8950,  3.1142],
#         [ 1.8311, 18.3808,  1.9908]])

My changes here are largely additive, so the existing usage via your own GitHub is not affected. I also updated the Model Card to make it clearer what this model is and does.

  • Tom Aarsen
tomaarsen changed pull request status to open
Ready to merge
This branch is ready to get merged automatically.

Sign up or log in to comment