Views
No views yet
MultiVectorEncoder:pip install "sentence-transformers>=6.0.0"1from sentence_transformers import MultiVectorEncoder
2
3model = MultiVectorEncoder("mixedbread-ai/mxbai-edge-colbert-v0-17m")
4
5query = "Which planet is known as the Red Planet?"
6documents = [
7 "Venus is often called Earth's twin because of its similar size and proximity.",
8 "Mars, known for its reddish appearance, is often referred to as the Red Planet.",
9 "Jupiter, the largest planet in our solar system, has a prominent red spot.",
10 "Saturn, famous for its rings, is sometimes mistaken for the Red Planet.",
11]
12
13query_embeddings = model.encode_query(query)
14document_embeddings = model.encode_document(documents)
15print(query_embeddings.shape, document_embeddings[0].shape)
16# (12, 48) (18, 48)
17
18# MaxSim late-interaction scoring (higher is more relevant)
19scores = model.similarity(query_embeddings, document_embeddings)
20print(scores)
21# tensor([[11.5693, 11.7584, 11.7099, 11.7229]])1# uv
2uv add pylate
3# uv + pip
4uv pip install pylate1# pip
2pip install -U pylate1from pylate import indexes, models, retrieve
2
3# Step 1: Load the model
4model = models.ColBERT(
5 model_name_or_path="mixedbread-ai/mxbai-edge-colbert-v0-17m",
6)
7
8
9# Step 2: Initialize an index (here, PLAID, for larger document collections)
10index = indexes.PLAID(
11 index_folder="pylate-index",
12 index_name="index",
13 override=True, # This overwrites the existing index if any
14)
15
16# Step 3: Encode your documents
17documents_ids = ["1", "2", "3"]
18documents = ["document 1 text", "document 2 text", "document 3 text"]
19
20documents_embeddings = model.encode(
21 documents,
22 batch_size=32,
23 is_query=False, # Ensure that it is set to False to indicate that these are documents, not queries
24 show_progress_bar=True,
25)
26
27# Step 4: Add document embeddings to the index by providing embeddings and corresponding ids
28index.add_documents(
29 documents_ids=documents_ids,
30 documents_embeddings=documents_embeddings,
31)1# Step 5.1: Initialize the ColBERT retriever
2retriever = retrieve.ColBERT(index=index)
3
4# Step 2: Encode the queries
5queries_embeddings = model.encode(
6 ["query for document 3", "query for document 1"],
7 batch_size=32,
8 is_query=True, # # Ensure that it is set to False to indicate that these are queries
9 show_progress_bar=True,
10)
11
12# Step 3: Retrieve top-k documents
13scores = retriever.retrieve(
14 queries_embeddings=queries_embeddings,
15 k=10, # Retrieve the top 10 matches for each query
16)1from pylate import rank, models
2
3# Load the model
4model = models.ColBERT(
5 model_name_or_path="mixedbread-ai/mxbai-edge-colbert-v0-17m",
6)
7
8# Define queries and documents
9queries = [
10 "query A",
11 "query B",
12]
13
14documents = [
15 ["document A", "document B"],
16 ["document 1", "document C", "document B"],
17]
18documents_ids = [
19 [1, 2],
20 [1, 3, 2],
21]
22
23# Embed them
24queries_embeddings = model.encode(
25 queries,
26 is_query=True,
27)
28
29documents_embeddings = model.encode(
30 documents,
31 is_query=False,
32)
33
34# Perform reranking
35reranked_documents = rank.rerank(
36 documents_ids=documents_ids,
37 queries_embeddings=queries_embeddings,
38 documents_embeddings=documents_embeddings,
39)| Model | AVG | MS MARCO | SciFact | Touche | FiQA | TREC-COVID | NQ | DBPedia |
|---|---|---|---|---|---|---|---|---|
| Large Models (>100M) | ||||||||
| GTE-ModernColBERT-v1 | 0.547 | 0.453 | 0.763 | 0.312 | 0.453 | 0.836 | 0.618 | 0.480 |
| ColBERTv2 | 0.488 | 0.456 | 0.693 | 0.263 | 0.356 | 0.733 | 0.562 | 0.446 |
| Medium Models (<35M) | ||||||||
| mxbai-edge-colbert-v0-32m | 0.521 | 0.450 | 0.740 | 0.313 | 0.390 | 0.775 | 0.600 | 0.455 |
| answerai-colbert-small-v1 | 0.534 | 0.434 | 0.740 | 0.250 | 0.410 | 0.831 | 0.594 | 0.464 |
| bge-small-en-v1.5 | 0.517 | 0.408 | 0.713 | 0.260 | 0.403 | 0.759 | 0.502 | 0.400 |
| snowflake-s | 0.519 | 0.402 | 0.722 | 0.235 | 0.407 | 0.801 | 0.509 | 0.410 |
| Small Models (<25M) | ||||||||
| mxbai-edge-colbert-v0-17m | 0.490 | 0.416 | 0.719 | 0.316 | 0.326 | 0.713 | 0.551 | 0.410 |
| colbert-muvera-micro | 0.394 | 0.364 | 0.662 | 0.251 | 0.254 | 0.561 | 0.386 | 0.332 |
| all-MiniLM-L6-v2 | 0.419 | 0.365 | 0.645 | 0.169 | 0.369 | 0.472 | 0.439 | 0.323 |
| Model | AVG |
|---|---|
| Large Models (>100M) | |
| GTE-ModernColBERT-v1 (32k) | 0.898 |
| GTE-ModernColBERT-v1 (4k) | 0.809 |
| granite-embedding-english-r2 | 0.656 |
| ColBERTv2 | 0.428 |
| Medium Models (<50M) | |
| mxbai-edge-colbert-v0-32m (32k) | 0.849 |
| mxbai-edge-colbert-v0-32m (4k) | 0.783 |
| granite-embedding-small-english-r2 | 0.637 |
| answerai-colbert-small-v1 | 0.441 |
| bge-small-en-v1.5 | 0.312 |
| snowflake-arctic-embed-s | 0.356 |
| Small Models (<25M) | |
| mxbai-edge-colbert-v0-17m (32k) | 0.847 |
| mxbai-edge-colbert-v0-17m (4k) | 0.776 |
| all-MiniLM-L6-v2 | 0.298 |
| colbert-muvera-micro | 0.405 |
1@misc{takehi2025fantasticsmallretrieverstrain,
2 title={Fantastic (small) Retrievers and How to Train Them: mxbai-edge-colbert-v0 Tech Report},
3 author={Rikiya Takehi and Benjamin Clavié and Sean Lee and Aamir Shakir},
4 year={2025},
5 eprint={2510.14880},
6 archivePrefix={arXiv},
7 primaryClass={cs.IR},
8 url={https://arxiv.org/abs/2510.14880},
9}1@misc{clavie2025simpleprojectionvariantsimprove,
2 title={Simple Projection Variants Improve ColBERT Performance},
3 author={Benjamin Clavié and Sean Lee and Rikiya Takehi and Aamir Shakir and Makoto P. Kato},
4 year={2025},
5 eprint={2510.12327},
6 archivePrefix={arXiv},
7 primaryClass={cs.IR},
8 url={https://arxiv.org/abs/2510.12327},
9}1@misc{PyLate,
2title={PyLate: Flexible Training and Retrieval for Late Interaction Models},
3author={Chaffin, Antoine and Sourty, Raphaël},
4url={https://github.com/lightonai/pylate},
5year={2024}
6}