This is a
sentence-transformers model: It maps sentences & paragraphs to a 384 dimensional dense vector space and was designed for
semantic search. It has been trained on 500k (query, answer) pairs from the
MS MARCO Passages dataset. For an introduction to semantic search, have a look at:
SBERT.net - Semantic Search
Using this model becomes easy when you have
sentence-transformers installed:
1from sentence_transformers import SentenceTransformer, util
2
3query = "How many people live in London?"
4docs = ["Around 9 Million people live in London", "London is known for its financial district"]
5
6#Load the model
7model = SentenceTransformer('sentence-transformers/msmarco-MiniLM-L6-cos-v5')
8
9#Encode query and documents
10query_emb = model.encode(query)
11doc_emb = model.encode(docs)
12
13#Compute dot score between query and all document embeddings
14scores = util.dot_score(query_emb, doc_emb)[0].cpu().tolist()
15
16#Combine docs & scores
17doc_score_pairs = list(zip(docs, scores))
18
19#Sort by decreasing score
20doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)
21
22#Output passages & scores
23for doc, score in doc_score_pairs:
24 print(score, doc)
Without
sentence-transformers, you can use the model like this: First, you pass your input through the transformer model, then you have to apply the correct pooling-operation on-top of the contextualized word embeddings.
1from transformers import AutoTokenizer, AutoModel
2import torch
3import torch.nn.functional as F
4
5#Mean Pooling - Take average of all tokens
6def mean_pooling(model_output, attention_mask):
7 token_embeddings = model_output.last_hidden_state #First element of model_output contains all token embeddings
8 input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
9 return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
10
11
12#Encode text
13def encode(texts):
14 # Tokenize sentences
15 encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
16
17 # Compute token embeddings
18 with torch.no_grad():
19 model_output = model(**encoded_input, return_dict=True)
20
21 # Perform pooling
22 embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
23
24 # Normalize embeddings
25 embeddings = F.normalize(embeddings, p=2, dim=1)
26
27 return embeddings
28
29
30# Sentences we want sentence embeddings for
31query = "How many people live in London?"
32docs = ["Around 9 Million people live in London", "London is known for its financial district"]
33
34# Load model from HuggingFace Hub
35tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/msmarco-MiniLM-L6-cos-v5")
36model = AutoModel.from_pretrained("sentence-transformers/msmarco-MiniLM-L6-cos-v5")
37
38#Encode query and docs
39query_emb = encode(query)
40doc_emb = encode(docs)
41
42#Compute dot score between query and all document embeddings
43scores = torch.mm(query_emb, doc_emb.transpose(0, 1))[0].cpu().tolist()
44
45#Combine docs & scores
46doc_score_pairs = list(zip(docs, scores))
47
48#Sort by decreasing score
49doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)
50
51#Output passages & scores
52for doc, score in doc_score_pairs:
53 print(score, doc)
This model was trained by
sentence-transformers.
If you find this model helpful, feel free to cite our publication
Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks:
1@inproceedings{reimers-2019-sentence-bert,
2 title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
3 author = "Reimers, Nils and Gurevych, Iryna",
4 booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
5 month = "11",
6 year = "2019",
7 publisher = "Association for Computational Linguistics",
8 url = "http://arxiv.org/abs/1908.10084",
9}