mixedbread-ai

/

mxbai-colbert-large-v1

like 53

Follow

Mixedbread 415

Hello!

The MultiVectorEncoder class ships in the next Sentence Transformers release, planned for around the 18th, so for now the model-card snippet installs from source. I would love to feature this model in that release's blog post and documentation, especially once it loads without a revision pin (that is, once this PR is merged).

Heads up, this PR was AI-generated and human-reviewed. Here's a summary of the changes as reported by my agent:

Pull Request overview

• Remove the stale modules.json so the checkpoint loads as the Stanford-NLP ColBERT model it is
• Add a Sentence Transformers (MultiVectorEncoder) usage section and library tags

Details

This repository is a Stanford-NLP ColBERT checkpoint. config.json declares the HF_ColBERT architecture, artifact.metadata carries the trained settings (dim: 128, doc_maxlen: 256, query_maxlen: 128, mask_punctuation: true, the [unused0] / [unused1] marker tokens), and the trained 128 x 1024 projection is stored inline as linear.weight. RAGatouille and PyLate read exactly those files.

The modules.json added later (commit d2f33f0, Sep 2024) was for PyLate, although PyLate can also read artifact.metadata to create its models. The crux is that modules.json is only compatible with PyLate and not with Sentence Transformers, which uses a different module configuration.

Removing modules.json fixes this and preserves the existing behaviour for Stanford-NLP, PyLate, and Sentence Transformers: I loaded the checkpoint with PyLate before and after the change and the multi-vector output is bit-identical (128-dim, L2-normalized, same MaxSim scores and ranking).

The Dense projection built from the checkpoint is bit-identical to the stored linear.weight.

I also bundled the model card change so this is a single PR: a "Using Sentence Transformers" section with a MultiVectorEncoder snippet, and the sentence-transformers / multi-vector tags.

pip install "sentence-transformers @ git+https://github.com/huggingface/sentence-transformers.git"

from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder("mixedbread-ai/mxbai-colbert-large-v1", revision="refs/pr/4")

query = "Who wrote 'To Kill a Mockingbird'?"
documents = [
    "'To Kill a Mockingbird' is a novel by Harper Lee published in 1960. It was immediately successful, winning the Pulitzer Prize, and has become a classic of modern American literature.",
    "The novel 'Moby-Dick' was written by Herman Melville and first published in 1851. It is considered a masterpiece of American literature and deals with complex themes of obsession, revenge, and the conflict between good and evil.",
    "Harper Lee, an American novelist widely known for her novel 'To Kill a Mockingbird', was born in 1926 in Monroeville, Alabama. She received the Pulitzer Prize for Fiction in 1961.",
    "Jane Austen was an English novelist known primarily for her six major novels, which interpret, critique and comment upon the British landed gentry at the end of the 18th century.",
]

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

scores = model.similarity(query_embeddings, document_embeddings)
print(scores)
# tensor([[119.3276,  95.9337, 114.1252,  76.5609]])

• Tom Aarsen