Views
No views yet
| Model | Inference-free for Retrieval | Model Parameters | AVG NDCG@10 | AVG FLOPS |
|---|---|---|---|---|
| opensearch-neural-sparse-encoding-v1 | 133M | 0.524 | 11.4 | |
| opensearch-neural-sparse-encoding-v2-distill | 67M | 0.528 | 8.3 | |
| opensearch-neural-sparse-encoding-doc-v1 | ✔️ | 133M | 0.490 | 2.3 |
| opensearch-neural-sparse-encoding-doc-v2-distill | ✔️ | 67M | 0.504 | 1.8 |
| opensearch-neural-sparse-encoding-doc-v2-mini | ✔️ | 23M | 0.497 | 1.7 |
| opensearch-neural-sparse-encoding-doc-v3-distill | ✔️ | 67M | 0.517 | 1.8 |
| opensearch-neural-sparse-encoding-doc-v3-gte | ✔️ | 133M | 0.546 | 1.7 |
pip install -U sentence-transformers1from sentence_transformers.sparse_encoder import SparseEncoder
2
3# Download from the 🤗 Hub
4model = SparseEncoder("opensearch-project/opensearch-neural-sparse-encoding-v1")
5
6query = "What's the weather in ny now?"
7document = "Currently New York is rainy."
8
9query_embed = model.encode_query(query)
10document_embed = model.encode_document(document)
11
12sim = model.similarity(query_embed, document_embed)
13print(f"Similarity: {sim}")
14# Similarity: tensor([[22.3299]])
15
16decoded_query = model.decode(query_embed)
17decoded_document = model.decode(document_embed)
18
19for i in range(len(decoded_query)):
20 query_token, query_score = decoded_query[i]
21 doc_score = next((score for token, score in decoded_document if token == query_token), 0)
22 if doc_score != 0:
23 print(f"Token: {query_token}, Query score: {query_score:.4f}, Document score: {doc_score:.4f}")
24
25# Token: ny, Query score: 2.9262, Document score: 2.1335
26# Token: weather, Query score: 2.5206, Document score: 1.5277
27# Token: york, Query score: 2.0373, Document score: 2.3489
28# Token: cool, Query score: 1.5786, Document score: 0.8752
29# Token: current, Query score: 1.4636, Document score: 1.5132
30# Token: season, Query score: 0.7761, Document score: 0.8860
31# Token: 2020, Query score: 0.7560, Document score: 0.6726
32# Token: summer, Query score: 0.7222, Document score: 0.6292
33# Token: nina, Query score: 0.6888, Document score: 0.6419
34# Token: storm, Query score: 0.6451, Document score: 0.8200
35# Token: brooklyn, Query score: 0.4698, Document score: 0.7635
36# Token: julian, Query score: 0.4562, Document score: 0.1208
37# Token: wow, Query score: 0.3484, Document score: 0.3903
38# Token: usa, Query score: 0.3439, Document score: 0.4160
39# Token: manhattan, Query score: 0.2751, Document score: 0.8260
40# Token: fog, Query score: 0.2013, Document score: 0.7735
41# Token: mood, Query score: 0.1989, Document score: 0.2961
42# Token: climate, Query score: 0.1653, Document score: 0.3437
43# Token: nature, Query score: 0.1191, Document score: 0.1533
44# Token: temperature, Query score: 0.0665, Document score: 0.0599
45# Token: windy, Query score: 0.0552, Document score: 0.33961import itertools
2import torch
3from transformers import AutoModelForMaskedLM, AutoTokenizer
4
5
6# get sparse vector from dense vectors with shape batch_size * seq_len * vocab_size
7def get_sparse_vector(feature, output):
8 values, _ = torch.max(output*feature["attention_mask"].unsqueeze(-1), dim=1)
9 values = torch.log(1 + torch.relu(values))
10 values[:,special_token_ids] = 0
11 return values
12
13# transform the sparse vector to a dict of (token, weight)
14def transform_sparse_vector_to_dict(sparse_vector):
15 sample_indices,token_indices=torch.nonzero(sparse_vector,as_tuple=True)
16 non_zero_values = sparse_vector[(sample_indices,token_indices)].tolist()
17 number_of_tokens_for_each_sample = torch.bincount(sample_indices).cpu().tolist()
18 tokens = [transform_sparse_vector_to_dict.id_to_token[_id] for _id in token_indices.tolist()]
19
20 output = []
21 end_idxs = list(itertools.accumulate([0]+number_of_tokens_for_each_sample))
22 for i in range(len(end_idxs)-1):
23 token_strings = tokens[end_idxs[i]:end_idxs[i+1]]
24 weights = non_zero_values[end_idxs[i]:end_idxs[i+1]]
25 output.append(dict(zip(token_strings, weights)))
26 return output
27
28
29# load the model
30model = AutoModelForMaskedLM.from_pretrained("opensearch-project/opensearch-neural-sparse-encoding-v1")
31tokenizer = AutoTokenizer.from_pretrained("opensearch-project/opensearch-neural-sparse-encoding-v1")
32
33# set the special tokens and id_to_token transform for post-process
34special_token_ids = [tokenizer.vocab[token] for token in tokenizer.special_tokens_map.values()]
35get_sparse_vector.special_token_ids = special_token_ids
36id_to_token = ["" for i in range(tokenizer.vocab_size)]
37for token, _id in tokenizer.vocab.items():
38 id_to_token[_id] = token
39transform_sparse_vector_to_dict.id_to_token = id_to_token
40
41
42
43query = "What's the weather in ny now?"
44document = "Currently New York is rainy."
45
46# encode the query & document
47feature = tokenizer([query, document], padding=True, truncation=True, return_tensors='pt', return_token_type_ids=False)
48output = model(**feature)[0]
49sparse_vector = get_sparse_vector(feature, output)
50
51# get similarity score
52sim_score = torch.matmul(sparse_vector[0],sparse_vector[1])
53print(sim_score) # tensor(22.3299, grad_fn=<DotBackward0>)
54
55
56query_token_weight, document_query_token_weight = transform_sparse_vector_to_dict(sparse_vector)
57for token in sorted(query_token_weight, key=lambda x:query_token_weight[x], reverse=True):
58 if token in document_query_token_weight:
59 print("score in query: %.4f, score in document: %.4f, token: %s"%(query_token_weight[token],document_query_token_weight[token],token))
60
61
62
63# result:
64# score in query: 2.9262, score in document: 2.1335, token: ny
65# score in query: 2.5206, score in document: 1.5277, token: weather
66# score in query: 2.0373, score in document: 2.3489, token: york
67# score in query: 1.5786, score in document: 0.8752, token: cool
68# score in query: 1.4636, score in document: 1.5132, token: current
69# score in query: 0.7761, score in document: 0.8860, token: season
70# score in query: 0.7560, score in document: 0.6726, token: 2020
71# score in query: 0.7222, score in document: 0.6292, token: summer
72# score in query: 0.6888, score in document: 0.6419, token: nina
73# score in query: 0.6451, score in document: 0.8200, token: storm
74# score in query: 0.4698, score in document: 0.7635, token: brooklyn
75# score in query: 0.4562, score in document: 0.1208, token: julian
76# score in query: 0.3484, score in document: 0.3903, token: wow
77# score in query: 0.3439, score in document: 0.4160, token: usa
78# score in query: 0.2751, score in document: 0.8260, token: manhattan
79# score in query: 0.2013, score in document: 0.7735, token: fog
80# score in query: 0.1989, score in document: 0.2961, token: mood
81# score in query: 0.1653, score in document: 0.3437, token: climate
82# score in query: 0.1191, score in document: 0.1533, token: nature
83# score in query: 0.0665, score in document: 0.0600, token: temperature
84# score in query: 0.0552, score in document: 0.3396, token: windy| Model | Average | Trec Covid | NFCorpus | NQ | HotpotQA | FiQA | ArguAna | Touche | DBPedia | SCIDOCS | FEVER | Climate FEVER | SciFact | Quora |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| opensearch-neural-sparse-encoding-v1 | 0.524 | 0.771 | 0.360 | 0.553 | 0.697 | 0.376 | 0.508 | 0.278 | 0.447 | 0.164 | 0.821 | 0.263 | 0.723 | 0.856 |
| opensearch-neural-sparse-encoding-v2-distill | 0.528 | 0.775 | 0.347 | 0.561 | 0.685 | 0.374 | 0.551 | 0.278 | 0.435 | 0.173 | 0.849 | 0.249 | 0.722 | 0.863 |
| opensearch-neural-sparse-encoding-doc-v1 | 0.490 | 0.707 | 0.352 | 0.521 | 0.677 | 0.344 | 0.461 | 0.294 | 0.412 | 0.154 | 0.743 | 0.202 | 0.716 | 0.788 |
| opensearch-neural-sparse-encoding-doc-v2-distill | 0.504 | 0.690 | 0.343 | 0.528 | 0.675 | 0.357 | 0.496 | 0.287 | 0.418 | 0.166 | 0.818 | 0.224 | 0.715 | 0.841 |
| opensearch-neural-sparse-encoding-doc-v2-mini | 0.497 | 0.709 | 0.336 | 0.510 | 0.666 | 0.338 | 0.480 | 0.285 | 0.407 | 0.164 | 0.812 | 0.216 | 0.699 | 0.837 |
| opensearch-neural-sparse-encoding-doc-v3-distill | 0.517 | 0.724 | 0.345 | 0.544 | 0.694 | 0.356 | 0.520 | 0.294 | 0.424 | 0.163 | 0.845 | 0.239 | 0.708 | 0.863 |
| opensearch-neural-sparse-encoding-doc-v3-gte | 0.546 | 0.734 | 0.360 | 0.582 | 0.716 | 0.407 | 0.520 | 0.389 | 0.455 | 0.167 | 0.860 | 0.312 | 0.725 | 0.873 |