mixedbread-ai

/

mxbai-rerank-base-v2

like 63

Follow

Mixedbread 371

Hello!

Pull Request overview

• Integrate this model using a Sentence Transformers CrossEncoder

Details

This PR adds the configuration files needed to load this model directly as a CrossEncoder via Sentence Transformers. The model uses a text-generation Transformer with a LogitScore head that computes the logit difference between the "1" and "0" tokens (token IDs 16 and 15), i.e. the model's confidence that a document is relevant to a query.

A custom chat_template.jinja maps Sentence Transformers' structured messages (with "query" and "document" roles) to the model's expected Qwen2 chat format with the relevance judgment instruction.

Added files:

• modules.json: pipeline: Transformer & LogitScore
• sentence_bert_config.json: text-generation task, flat message format
• config_sentence_transformers.json: CrossEncoder model type
• chat_template.jinja: custom template for the reranker format
• 1_LogitScore/config.json: true/false token IDs

Changed files:

• tokenizer_config.json: removed old generic Qwen2 chat template (replaced by chat_template.jinja)
• README.md: fixed missing </b> tag

Once the Sentence Transformers v5.4 release is out, the model can be used immediately like so:

from sentence_transformers import CrossEncoder

model = CrossEncoder("mixedbread-ai/mxbai-rerank-base-v2", revision="refs/pr/9")

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.",
    "The 'Harry Potter' series, which consists of seven fantasy novels written by British author J.K. Rowling, is among the most popular and critically acclaimed books of the modern era.",
    "'The Great Gatsby', a novel written by American author F. Scott Fitzgerald, was published in 1925. The story is set in the Jazz Age and follows the life of millionaire Jay Gatsby and his pursuit of Daisy Buchanan.",
]

pairs = [(query, doc) for doc in documents]
scores = model.predict(pairs)
print(scores)
# [9.750735   1.3697281  8.080784   2.6611633  0.8729458  0.39884853]

rankings = model.rank(query, documents)
print(rankings)
# [{'corpus_id': 0, 'score': np.float32(9.750735)}, {'corpus_id': 2, 'score': np.float32(8.080784)}, {'corpus_id': 3, 'score': np.float32(2.6611633)}, {'corpus_id': 1, 'score': np.float32(1.3697281)}, {'corpus_id': 4, 'score': np.float32(0.8729458)}, {'corpus_id': 5, 'score': np.float32(0.39884853)}]

And after merging, the revision argument can be dropped.

Note that none of the old behaviour is affected/changed. It only adds an additional way to run this model in a familiar and common format.

• Tom Aarsen