metadata
license: apache-2.0
language:
- en
tags:
- ColBERT
- multi-vector
- sentence-transformers
- sentence-similarity
- feature-extraction
- generated_from_trainer
- dataset_size:497901
- loss:Contrastive
base_model: google/bert_uncased_L-2_H-128_A-2
datasets:
- sentence-transformers/msmarco-bm25
pipeline_tag: sentence-similarity
Model card for ColBERT v2 BERT Tiny
This is a ColBERT model finetuned from google/bert_uncased_L-2_H-128_A-2 on the msmarco-bm25 dataset. It maps sentences & paragraphs to sequences of 128-dimensional dense vectors and can be used for semantic textual similarity using the MaxSim operator.
This model is primarily designed for unit tests in limited compute environments such as GitHub Actions. But it does work to an extent for basic use cases.
Usage with Sentence Transformers
As of Sentence Transformers v6.0.0, this model loads directly as a multi-vector (ColBERT-style late interaction) retriever via the MultiVectorEncoder:
pip install "sentence-transformers>=6.0.0"
from sentence_transformers import MultiVectorEncoder
model = MultiVectorEncoder("NeuML/colbert-bert-tiny")
query = "What is the capital of France?"
documents = [
"Paris is the capital and largest city of France.",
"Berlin is the capital of Germany.",
]
query_embeddings = model.encode_query(query)
document_embeddings = model.encode_document(documents)
print(query_embeddings.shape, document_embeddings[0].shape)
# torch.Size([32, 128]) torch.Size([12, 128])
# MaxSim late-interaction scoring (higher is more relevant)
scores = model.similarity(query_embeddings, document_embeddings)
print(scores)
# tensor([[25.9327, 23.9168]], device='cuda:0')