multi-qa-distilbert-cos-v1

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

pip install -U sentence-transformers

from sentence_transformers import SentenceTransformer, util

query = "How many people live in London?"
docs = ["Around 9 Million people live in London", "London is known for its financial district"]

# Cargar el modelo
model = SentenceTransformer('sentence-transformers/multi-qa-distilbert-cos-v1')

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

# Calcular puntaje de punto entre la consulta y todos los embeddings 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)

Sin sentence-transformers, puedes usar el modelo así: Primero, pasas tu entrada a través del modelo transformador, luego tienes que aplicar la operación de pooling adecuada en la parte superior de los embeddings contextualizados.

from transformers import AutoTokenizer, AutoModel
import torch
import torch.nn.functional as F

# Pooling promedio: tomar el promedio de todos los tokens
def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output.last_hidden_state
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

# Codificar texto
def encode(texts):
    encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
    with torch.no_grad():
        model_output = model(**encoded_input, return_dict=True)
    embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
    embeddings = F.normalize(embeddings, p=2, dim=1)
    return embeddings

query = "How many people live in London?"
docs = ["Around 9 Million people live in London", "London is known for its financial district"]

# Cargar modelo desde HuggingFace Hub
tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/multi-qa-distilbert-cos-v1")
model = AutoModel.from_pretrained("sentence-transformers/multi-qa-distilbert-cos-v1")

query_emb = encode(query)
doc_emb = encode(docs)

# Calcular puntaje de punto entre la consulta y todos los embeddings 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)