Add Sentence Transformers usage
Hello!
As of Sentence Transformers v6.0.1, this checkpoint loads directly as a multi-vector (ColBERT-style late interaction) retriever through the MultiVectorEncoder, alongside its existing PyLate usage. This PR adds a Sentence Transformers usage section to the model card and the multi-vector tag. The weights and the existing files are untouched.
Sister PR of https://huggingface.co/lightonai/ColBERT-Zero/discussions/2. Worth calling out why these say v6.0.1 rather than v6.0.0: your checkpoints pair the [Q] /[D] markers with the search_query: /search_document: prompts, and v6.0.0 applied only the prompts, dropping the markers. PyLate applies both, so that was a bug on my side. It is fixed now, and encode_query renders [CLS][Q] search_query: ... again, matching PyLate token for token. On NanoBEIR the difference was worth 0.0256 nDCG@10 on ColBERT-Zero (0.6569 before, 0.6824 after), so the ColBERT-Zero row in my documentation moves up accordingly. Thank you for shipping configs that made this catchable.
Heads up, this PR was AI-generated and human-reviewed.
pip install "sentence-transformers>=6.0.1"
from sentence_transformers import MultiVectorEncoder
model = MultiVectorEncoder("lightonai/ColBERT-Zero-supervised")
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)
# torch.Size([16, 128]) torch.Size([19, 128])
# MaxSim late-interaction scoring (higher is more relevant)
scores = model.similarity(query_embeddings, document_embeddings)
print(scores)
# tensor([[ 9.7358, 11.2758, 10.7421, 10.7044]], device='cuda:0')
Verified against PyLate itself: per-token cosine 1.000000 on both the query and the document, with a maximum absolute deviation around 2e-07.
The config also gains a requirements block pinning sentence-transformers>=6.0.1. Sentence Transformers reads it at load time and fails with a clear message rather than letting an older install return the wrong embeddings silently, which is what happens today.
To try this before merging, pass revision="refs/pr/1" to MultiVectorEncoder.
Happy to tweak anything you'd like changed. Please let me know if you have any questions or feedback!
- Tom Aarsen