mxbai-edge-colbert-v0-32m

pip install "sentence-transformers>=6.0.0"

from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder("mixedbread-ai/mxbai-edge-colbert-v0-32m")

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)
# (12, 64) (18, 64)

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

# uv
uv add pylate
# uv + pip
uv pip install pylate

# pip
pip install -U pylate

from pylate import indexes, models, retrieve

# Step 1: Load the model
model = models.ColBERT(
    model_name_or_path="mixedbread-ai/mxbai-edge-colbert-v0-32m",
)


# Step 2: Initialize an index (here, PLAID, for larger document collections)
index = indexes.PLAID(
    index_folder="pylate-index",
    index_name="index",
    override=True,  # This overwrites the existing index if any
)

# Step 3: Encode your documents
documents_ids = ["1", "2", "3"]
documents = ["document 1 text", "document 2 text", "document 3 text"]

documents_embeddings = model.encode(
    documents,
    batch_size=32,
    is_query=False,  # Ensure that it is set to False to indicate that these are documents, not queries
    show_progress_bar=True,
)

# Step 4: Add document embeddings to the index by providing embeddings and corresponding ids
index.add_documents(
    documents_ids=documents_ids,
    documents_embeddings=documents_embeddings,
)

# Step 5.1: Initialize the ColBERT retriever
retriever = retrieve.ColBERT(index=index)

# Step 2: Encode the queries
queries_embeddings = model.encode(
    ["query for document 3", "query for document 1"],
    batch_size=32,
    is_query=True,  #  # Ensure that it is set to False to indicate that these are queries
    show_progress_bar=True,
)

# Step 3: Retrieve top-k documents
scores = retriever.retrieve(
    queries_embeddings=queries_embeddings,
    k=10,  # Retrieve the top 10 matches for each query
)

from pylate import rank, models

# Load the model
model = models.ColBERT(
    model_name_or_path="mixedbread-ai/mxbai-edge-colbert-v0-32m",
)

# Define queries and documents
queries = [
    "query A",
    "query B",
]

documents = [
    ["document A", "document B"],
    ["document 1", "document C", "document B"],
]
documents_ids = [
    [1, 2],
    [1, 3, 2],
]

# Embed them
queries_embeddings = model.encode(
    queries,
    is_query=True,
)

documents_embeddings = model.encode(
    documents,
    is_query=False,
)

# Perform reranking
reranked_documents = rank.rerank(
    documents_ids=documents_ids,
    queries_embeddings=queries_embeddings,
    documents_embeddings=documents_embeddings,
)

@misc{takehi2025fantasticsmallretrieverstrain,
      title={Fantastic (small) Retrievers and How to Train Them: mxbai-edge-colbert-v0 Tech Report}, 
      author={Rikiya Takehi and Benjamin Clavié and Sean Lee and Aamir Shakir},
      year={2025},
      eprint={2510.14880},
      archivePrefix={arXiv},
      primaryClass={cs.IR},
      url={https://arxiv.org/abs/2510.14880}, 
}

@misc{clavie2025simpleprojectionvariantsimprove,
      title={Simple Projection Variants Improve ColBERT Performance}, 
      author={Benjamin Clavié and Sean Lee and Rikiya Takehi and Aamir Shakir and Makoto P. Kato},
      year={2025},
      eprint={2510.12327},
      archivePrefix={arXiv},
      primaryClass={cs.IR},
      url={https://arxiv.org/abs/2510.12327}, 
}

@misc{PyLate,
title={PyLate: Flexible Training and Retrieval for Late Interaction Models},
author={Chaffin, Antoine and Sourty, Raphaël},
url={https://github.com/lightonai/pylate},
year={2024}
}