Views
No views yet
pip install -U sentence-transformers1from 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)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)| Setting | Value |
|---|---|
| Dimensions | 384 |
| Produces normalized embeddings | Yes |
| Pooling-Method | Mean pooling |
| Suitable score functions | dot-product (util.dot_score), cosine-similarity (util.cos_sim), or euclidean distance |
sentence-transformers, this model produces normalized embeddings with length 1. In that case, dot-product and cosine-similarity are equivalent. dot-product is preferred as it is faster. Euclidean distance is proportional to dot-product and can also be used.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}