sentence-transformers/msmarco-distilbert-base-tas-b

Uso (Sentence-Transformers)

Usar este modelo es fácil cuando tienes sentence-transformers instalado:

pip install -U sentence-transformers

Luego puedes usar el modelo así:

from sentence_transformers import SentenceTransformer, util

query = "¿Cuántas personas viven en Londres?"
docs = ["Alrededor de 9 millones de personas viven en Londres", "Londres es conocido por su distrito financiero"]

# Cargar el modelo
model = SentenceTransformer('sentence-transformers/msmarco-distilbert-base-tas-b')

# Codificar consulta y documentos
query_emb = model.encode(query)
doc_emb = model.encode(docs)

# Calcular puntaje de puntos entre la consulta y todas las incrustaciones de documentos
ores = util.dot_score(query_emb, doc_emb)[0].cpu().tolist()

# Combinar documentos y puntajes
doc_score_pairs = list(zip(docs, scores))

# Ordenar en orden decreciente
doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)

# Poner de salida los pasajes y puntajes
for doc, score in doc_score_pairs:
    print(score, doc)

Uso (HuggingFace Transformers)

Sin sentence-transformers, puedes usar el modelo de esta manera: primero, pasas tu entrada a través del modelo transformer, luego debes aplicar la operación de agrupamiento correcta sobre las incrustaciones de palabras contextualizadas.

from transformers import AutoTokenizer, AutoModel
import torch

# CLS Pooling - Tomar salida del primer token
def cls_pooling(model_output):
    return model_output.last_hidden_state[:,0]

# Codificar texto
def encode(texts):
    # Tokenizar oraciones
    encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors='pt')

    # Calcular incrustaciones de tokens
    with torch.no_grad():
        model_output = model(**encoded_input, return_dict=True)

    # Realizar agrupamiento
    embeddings = cls_pooling(model_output)

    return embeddings

# Consultas para las cuales queremos incrustaciones de oraciones
query = "¿Cuántas personas viven en Londres?"
docs = ["Alrededor de 9 millones de personas viven en Londres", "Londres es conocido por su distrito financiero"]

# Cargar modelo del Hub de HuggingFace
tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/msmarco-distilbert-base-tas-b")
model = AutoModel.from_pretrained("sentence-transformers/msmarco-distilbert-base-tas-b")

# Codificar consulta y documentos
query_emb = encode(query)
doc_emb = encode(docs)

# Calcular puntaje de puntos entre la consulta y todas las incrustaciones de documentos
scores = torch.mm(query_emb, doc_emb.transpose(0, 1))[0].cpu().tolist()

# Combinar documentos y puntajes
doc_score_pairs = list(zip(docs, scores))

# Ordenar en orden decreciente
doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)

# Poner de salida los pasajes y puntajes
for doc, score in doc_score_pairs:
    print(score, doc)

Resultados de evaluación

Para una evaluación automatizada de este modelo, mira el Benchmark de Incrustaciones de Oraciones: https://seb.sbert.net