mxbai-rerank-xsmall-v1

Actualmente, la mejor manera de usar nuestros modelos es con la versión más reciente de sentence-transformers.

pip install -U sentence-transformers

Vamos a suponer que tienes una consulta y quieres reordenar un conjunto de documentos. Puedes hacerlo con solo una línea de código:

from sentence_transformers import CrossEncoder
# Cargar el modelo
model = CrossEncoder("mixedbread-ai/mxbai-rerank-xsmall-v1")
# Ejemplo de consulta y documentos
query = "¿Quién escribió 'To Kill a Mockingbird'?"
documents = [
    "'To Kill a Mockingbird' es una novela de Harper Lee publicada en 1960. Fue inmediatamente exitosa, ganando el Premio Pulitzer, y se ha convertido en un clásico de la literatura estadounidense moderna.",
    "La novela 'Moby-Dick' fue escrita por Herman Melville y publicada por primera vez en 1851. Se considera una obra maestra de la literatura estadounidense y trata temas complejos de obsesión, venganza y el conflicto entre el bien y el mal.",
    "Harper Lee, una novelista estadounidense ampliamente conocida por su novela 'To Kill a Mockingbird', nació en 1926 en Monroeville, Alabama. Recibió el Premio Pulitzer de Ficción en 1961.",
    "Jane Austen fue una novelista inglesa conocida principalmente por sus seis novelas principales, que interpretan, critican y comentan la nobleza rural británica a finales del siglo XVIII.",
    "La serie 'Harry Potter', que consta de siete novelas de fantasía escritas por la autora británica J.K. Rowling, se encuentra entre los libros más populares y aclamados de la era moderna.",
    "'The Great Gatsby', una novela escrita por el autor estadounidense F. Scott Fitzgerald, fue publicada en 1925. La historia está ambientada en la Era del Jazz y sigue la vida del millonario Jay Gatsby y su búsqueda de Daisy Buchanan."
]
# Obtener los puntajes
results = model.rank(query, documents, return_documents=True, top_k=3)

Ejemplo en JavaScript

import { AutoTokenizer, AutoModelForSequenceClassification } from '@xenova/transformers';
const model_id = 'mixedbread-ai/mxbai-rerank-xsmall-v1';
const model = await AutoModelForSequenceClassification.from_pretrained(model_id);
const tokenizer = await AutoTokenizer.from_pretrained(model_id);
async function rank(query, documents, {
    top_k = undefined,
    return_documents = false,
} = {}) {
    const inputs = tokenizer(
        new Array(documents.length).fill(query),
        {
            text_pair: documents,
            padding: true,
            truncation: true,
        }
    );
    const { logits } = await model(inputs);
    return logits
        .sigmoid()
        .tolist()
        .map(([score], i) => ({
            corpus_id: i,
            score,
            ...(return_documents ? { text: documents[i] } : {})
        }))
        .sort((a, b) => b.score - a.score)
        .slice(0, top_k);
}
// Ejemplo de uso:
const query = "¿Quién escribió 'To Kill a Mockingbird'?"
const documents = [
    "'To Kill a Mockingbird' es una novela de Harper Lee publicada en 1960. Fue inmediatamente exitosa, ganando el Premio Pulitzer, y se ha convertido en un clásico de la literatura estadounidense moderna.",
    "La novela 'Moby-Dick' fue escrita por Herman Melville y publicada por primera vez en 1851. Se considera una obra maestra de la literatura estadounidense y trata temas complejos de obsesión, venganza y el conflicto entre el bien y el mal.",
    "Harper Lee, una novelista estadounidense ampliamente conocida por su novela 'To Kill a Mockingbird', nació en 1926 en Monroeville, Alabama. Recibió el Premio Pulitzer de Ficción en 1961.",
    "Jane Austen fue una novelista inglesa conocida principalmente por sus seis novelas principales, que interpretan, critican y comentan la nobleza rural británica a finales del siglo XVIII.",
    "La serie 'Harry Potter', que consta de siete novelas de fantasía escritas por la autora británica J.K. Rowling, se encuentra entre los libros más populares y aclamados de la era moderna.",
    "'The Great Gatsby', una novela escrita por el autor estadounidense F. Scott Fitzgerald, fue publicada en 1925. La historia está ambientada en la Era del Jazz y sigue la vida del millonario Jay Gatsby y su búsqueda de Daisy Buchanan."
]
const results = await rank(query, documents, { return_documents: true, top_k: 3 });
console.log(results);