Add Sentence Transformers usage
Hello!
Starting with the next Sentence Transformers release (v6.0.0, planned for around the 18th), this checkpoint loads directly as a multi-vector (ColBERT-style late interaction) retriever through the new MultiVectorEncoder, alongside its existing PyLate usage. This PR adds a Sentence Transformers usage section to the model card and a multi-vector tag. The weights and the existing PyLate usage are untouched.
I'd love to feature this model in that release's blog post and documentation, especially once it loads without the revision pin (that is, once this PR is merged).
pip install "sentence-transformers @ git+https://github.com/huggingface/sentence-transformers.git"
from sentence_transformers import MultiVectorEncoder
model = MultiVectorEncoder("NeuML/biomedbert-base-colbert", revision="refs/pr/1")
query = "Which planet is known as the Red Planet?"
documents = [
"Venus is often called Earth's twin because of its similar size and proximity.",
"Mars, known for its reddish appearance, is often referred to as the Red Planet.",
"Jupiter, the largest planet in our solar system, has a prominent red spot.",
"Saturn, famous for its rings, is sometimes mistaken for the Red Planet.",
]
query_embeddings = model.encode_query(query)
document_embeddings = model.encode_document(documents)
print(query_embeddings.shape, document_embeddings[0].shape)
# (14, 128) (17, 128)
# MaxSim late-interaction scoring (higher is more relevant)
scores = model.similarity(query_embeddings, document_embeddings)
print(scores)
# tensor([[9.9712, 12.4123, 11.0444, 11.9934]])
- Tom Aarsen
Very exciting! Going to merge this in.
Have you considering adding in concepts like MUVERA to produce flat vectors from ColBERT vectors?
Hello!
I've been considering it indeed! I think you should be able to grab the sequential modules for a MultiVectorEncoder (Transformer, Dense, MultiVectorMask), then you can append a custom Muvera module (or related, I think there's also LEMUR or something, I'm a bit out of date but I saw it on txtai) to turn that into dense embeddings.
The only architectural downside (a pretty big one at that) is that Sentence Transformers fundamentally doesn't support multiple output types (i.e. dense, sparse, multi-vector), which prevents the most powerful use of MUVERA-like models: one forward to get 2 representations: one for retrieving and one for reranking the top retrieved inputs. Does that make sense?
- Tom Aarsen
Makes sense to me!