colbert-bert-tiny / README.md
tomaarsen's picture
tomaarsen HF Staff
Add Sentence Transformers usage
a91ed33 verified
|
Raw
History Blame
1.91 kB
---
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](https://github.com/stanford-futuredata/ColBERT) model finetuned from [google/bert_uncased_L-2_H-128_A-2](https://huggingface.co/google/bert_uncased_L-2_H-128_A-2) on the [msmarco-bm25](https://huggingface.co/datasets/sentence-transformers/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](https://www.sbert.net/) v6.0.0, this model loads directly as a multi-vector (ColBERT-style late interaction) retriever via the `MultiVectorEncoder`:
```bash
pip install "sentence-transformers>=6.0.0"
```
```python
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')
```