Views
No views yet
e5-base-unsupervised using an improved dataset that leverages modern hard-negative mining practices and includes three more high-quality retrieval datasets than the original E5 finetuning pipeline.| Model | BEIR Score (nDCG@10) | CLEF English (nDCG@10) |
|---|---|---|
| e5-base-v2 | 50.19 | 45.38 |
| arctic-e5-base | 54.70 | 52.77 |
| gte-base-en-v1.5 | 54.02 | 47.91 |
| arctic-embed-m-v1.0 | 54.89 | 47.62 |
| arctic-embed-m-v2.0 | 55.38 | 54.06 |
1from sentence_transformers import SentenceTransformer
2
3model = SentenceTransformer("Snowflake/snowflake-arctic-e5-base")
4
5queries = ['what is snowflake?', 'Where can I get the best tacos?']
6documents = ['The Data Cloud!', 'Mexico City of Course!']
7
8query_embeddings = model.encode(queries, prompt_name="query")
9document_embeddings = model.encode(documents)
10
11scores = query_embeddings @ document_embeddings.T
12for query, query_scores in zip(queries, scores):
13 doc_score_pairs = list(zip(documents, query_scores))
14 doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)
15 # Output passages & scores
16 print("Query:", query)
17 for document, score in doc_score_pairs:
18 print(score, document)Query: what is snowflake?
0.2747492 The Data Cloud!
0.19998045 Mexico City of Course!
Query: Where can I get the best tacos?
0.29974818 Mexico City of Course!
0.2344071 The Data Cloud!1import torch
2from transformers import AutoModel, AutoTokenizer
3
4tokenizer = AutoTokenizer.from_pretrained('Snowflake/snowflake-arctic-e5-base')
5model = AutoModel.from_pretrained('Snowflake/snowflake-arctic-e5-base')
6model.eval()
7
8query_prefix = 'query: '
9queries = ['what is snowflake?', 'Where can I get the best tacos?']
10queries_with_prefix = ["{}{}".format(query_prefix, q) for q in queries]
11query_tokens = tokenizer(queries_with_prefix, padding=True, truncation=True, return_tensors='pt', max_length=512)
12
13document_prefix = 'passage: '
14documents = ['The Data Cloud!', 'Mexico City of Course!']
15documents_with_prefix = ["{}{}".format(document_prefix, d) for d in documents]
16document_tokens = tokenizer(documents_with_prefix, padding=True, truncation=True, return_tensors='pt', max_length=512)
17
18# Compute token embeddings
19with torch.inference_mode():
20 query_embeddings = model(**query_tokens)[0][:, 0]
21 document_embeddings = model(**document_tokens)[0][:, 0]
22
23
24# normalize embeddings
25query_embeddings = torch.nn.functional.normalize(query_embeddings, p=2, dim=1)
26document_embeddings = torch.nn.functional.normalize(document_embeddings, p=2, dim=1)
27
28scores = torch.mm(query_embeddings, document_embeddings.transpose(0, 1))
29for query, query_scores in zip(queries, scores):
30 doc_score_pairs = list(zip(documents, query_scores))
31 doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)
32 #Output passages & scores
33 print("Query:", query)
34 for document, score in doc_score_pairs:
35 print(score, document)