multi-qa-MiniLM-L6-dot-v1

Uso (Sentence-Transformers)

Instalar sentence-transformers:

pip install -U sentence-transformers

Luego, puedes usar el modelo de esta manera:

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 conocida por su distrito financiero"]

# Cargar el modelo
model = SentenceTransformer('sentence-transformers/multi-qa-MiniLM-L6-dot-v1')

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

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

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

# Ordenar por puntaje decreciente
pairs = sorted(pairs, key=lambda x: x[1], reverse=True)

# Mostrar pasajes y puntajes
for doc, score in pairs:
    print(score, doc)

Uso (HuggingFace Transformers)

Sin sentence-transformers, puedes usar el modelo así:

from transformers import AutoTokenizer, AutoModel
import torch

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

# Codificar texto
def encode(texts):
    # Tokenizar frases
    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 pooling
    embeddings = cls_pooling(model_output)
    return embeddings

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

# Cargar modelo de HuggingFace Hub
tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/multi-qa-MiniLM-L6-dot-v1")
model = AutoModel.from_pretrained("sentence-transformers/multi-qa-MiniLM-L6-dot-v1")

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

# Calcular el puntaje punto 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
pairs = list(zip(docs, scores))

# Ordenar por puntaje decreciente
pairs = sorted(pairs, key=lambda x: x[1], reverse=True)

# Mostrar pasajes y puntajes
for doc, score in pairs:
    print(score, doc)