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-doc-v3-gte", trust_remote_code=True, model_kwargs={"code_revision": "40ced75c3017eb27626c9d4ea981bde21a2662f4"})
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([[12.5747]])
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: 5.7729, Document score: 0.9703
26# Token: weather, Query score: 4.5684, Document score: 1.0387
27# Token: now, Query score: 3.5895, Document score: 0.5861
28# Token: in, Query score: 0.4989, Document score: 0.24941import json
2import itertools
3import torch
4
5from transformers import AutoModelForMaskedLM, AutoTokenizer
6
7
8# get sparse vector from dense vectors with shape batch_size * seq_len * vocab_size
9def get_sparse_vector(feature, output):
10 values, _ = torch.max(output*feature["attention_mask"].unsqueeze(-1), dim=1)
11 # note we update the activation for v3 model
12 values = torch.log(1 + torch.log(1 + torch.relu(values)))
13 values[:,special_token_ids] = 0
14 return values
15
16# transform the sparse vector to a dict of (token, weight)
17def transform_sparse_vector_to_dict(sparse_vector):
18 sample_indices,token_indices=torch.nonzero(sparse_vector,as_tuple=True)
19 non_zero_values = sparse_vector[(sample_indices,token_indices)].tolist()
20 number_of_tokens_for_each_sample = torch.bincount(sample_indices).cpu().tolist()
21 tokens = [transform_sparse_vector_to_dict.id_to_token[_id] for _id in token_indices.tolist()]
22
23 output = []
24 end_idxs = list(itertools.accumulate([0]+number_of_tokens_for_each_sample))
25 for i in range(len(end_idxs)-1):
26 token_strings = tokens[end_idxs[i]:end_idxs[i+1]]
27 weights = non_zero_values[end_idxs[i]:end_idxs[i+1]]
28 output.append(dict(zip(token_strings, weights)))
29 return output
30
31# download the idf file from model hub. idf is used to give weights for query tokens
32def get_tokenizer_idf(tokenizer):
33 from huggingface_hub import hf_hub_download
34 local_cached_path = hf_hub_download(repo_id="opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte", filename="idf.json")
35 with open(local_cached_path) as f:
36 idf = json.load(f)
37 idf_vector = [0]*tokenizer.vocab_size
38 for token,weight in idf.items():
39 _id = tokenizer._convert_token_to_id_with_added_voc(token)
40 idf_vector[_id]=weight
41 return torch.tensor(idf_vector)
42
43# load the model
44model = AutoModelForMaskedLM.from_pretrained("opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte", trust_remote_code=True, code_revision="40ced75c3017eb27626c9d4ea981bde21a2662f4")
45tokenizer = AutoTokenizer.from_pretrained("opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte")
46idf = get_tokenizer_idf(tokenizer)
47
48# set the special tokens and id_to_token transform for post-process
49special_token_ids = [tokenizer.vocab[token] for token in tokenizer.special_tokens_map.values()]
50get_sparse_vector.special_token_ids = special_token_ids
51id_to_token = ["" for i in range(tokenizer.vocab_size)]
52for token, _id in tokenizer.vocab.items():
53 id_to_token[_id] = token
54transform_sparse_vector_to_dict.id_to_token = id_to_token
55
56
57
58query = "What's the weather in ny now?"
59document = "Currently New York is rainy."
60
61# encode the query
62feature_query = tokenizer([query], padding=True, truncation=True, return_tensors='pt', return_token_type_ids=False)
63input_ids = feature_query["input_ids"]
64batch_size = input_ids.shape[0]
65query_vector = torch.zeros(batch_size, tokenizer.vocab_size)
66query_vector[torch.arange(batch_size).unsqueeze(-1), input_ids] = 1
67query_sparse_vector = query_vector*idf
68
69# encode the document
70feature_document = tokenizer([document], padding=True, truncation=True, return_tensors='pt', return_token_type_ids=False)
71output = model(**feature_document)[0]
72document_sparse_vector = get_sparse_vector(feature_document, output)
73
74
75# get similarity score
76sim_score = torch.matmul(query_sparse_vector[0],document_sparse_vector[0])
77print(sim_score) # tensor(12.5747, grad_fn=<DotBackward0>)
78
79
80query_token_weight = transform_sparse_vector_to_dict(query_sparse_vector)[0]
81document_query_token_weight = transform_sparse_vector_to_dict(document_sparse_vector)[0]
82for token in sorted(query_token_weight, key=lambda x:query_token_weight[x], reverse=True):
83 if token in document_query_token_weight:
84 print("score in query: %.4f, score in document: %.4f, token: %s"%(query_token_weight[token],document_query_token_weight[token],token))
85
86
87# result:
88# score in query: 5.7729, score in document: 0.9703, token: ny
89# score in query: 4.5684, score in document: 1.0387, token: weather
90# score in query: 3.5895, score in document: 0.5861, token: now
91# score in query: 0.4989, score in document: 0.2494, token: in| 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 |