Add Sentence Transformers usage

#22
by tomaarsen HF Staff - opened

Hello!

As of Sentence Transformers v6.0.0, this checkpoint loads directly as a multi-vector (ColBERT-style late interaction) retriever through the new MultiVectorEncoder. This PR adds a Sentence Transformers usage section to the model card and the multi-vector and sentence-transformers tags. The weights and the existing usage are untouched.

pip install "sentence-transformers>=6.0.0"
from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder("jinaai/jina-colbert-v2", trust_remote_code=True)

query = "What does ColBERT do?"
documents = [
    "ColBERT is a novel ranking model that adapts deep LMs for efficient retrieval.",
    "Jina-ColBERT is a ColBERT-style model but based on JinaBERT so it can support both 8k context length, fast and accurate retrieval.",
]

query_embeddings = model.encode_query(query)
document_embeddings = model.encode_document(documents)
print(query_embeddings.shape, document_embeddings[0].shape)
# (32, 128) (23, 128)

# MaxSim late-interaction scoring (higher is more relevant)
scores = model.similarity(query_embeddings, document_embeddings)
print(scores)
# tensor([[23.2578, 21.5039]])

Verified against the merged implementation (bd55a5ec): the snippet reproduces exactly, and in fp32 the output matches a transformers==4.43.2 reference to 2.4e-07 on the token embeddings with identical MaxSim scores.

  • Tom Aarsen
tomaarsen changed pull request status to open

Thanks @tomaarsen !

michael-guenther changed pull request status to merged

Sign up or log in to comment