Views
No views yet
| Model | Inference-free for Retrieval | Model Parameters | AVG NDCG@10 | AVG FLOPS | AVG EMB SIZE |
|---|---|---|---|---|---|
| opensearch-neural-sparse-encoding-multilingual-v1 | ✔️ | 160M | 0.629 | 1.3 | 138 |
| opensearch-neural-sparse-encoding-multilingual-v1; prune_ratio 0.1 | ✔️ | 160M | 0.626 | 0.8 | 75 |
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-multilingual-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([[7.7400]])
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: weather, Query score: 3.0699, Document score: 1.2821
26# Token: now, Query score: 1.6406, Document score: 0.9018
27# Token: ?, Query score: 1.6108, Document score: 0.3141
28# Token: ny, Query score: 1.2721, Document score: 1.3446
29# Token: in, Query score: 0.6005, Document score: 0.18041import 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, prune_ratio=0.1):
10 values, _ = torch.max(output*feature["attention_mask"].unsqueeze(-1), dim=1)
11 values = torch.log(1 + torch.relu(values))
12 values[:,special_token_ids] = 0
13 max_values = values.max(dim=-1)[0].unsqueeze(1) * prune_ratio
14 return values * (values > max_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-multilingual-v1", 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-multilingual-v1")
45tokenizer = AutoTokenizer.from_pretrained("opensearch-project/opensearch-neural-sparse-encoding-multilingual-v1")
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(7.6317, 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
88# result:
89# score in query: 3.0699, score in document: 1.2821, token: weather
90# score in query: 1.6406, score in document: 0.9018, token: now
91# score in query: 1.6108, score in document: 0.3141, token: ?
92# score in query: 1.2721, score in document: 1.3446, token: ny| Model | Average | bn | te | es | fr | id | hi | ru | ar | zh | fa | ja | fi | sw | ko | en |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| BM25 | 0.305 | 0.482 | 0.383 | 0.077 | 0.115 | 0.297 | 0.350 | 0.256 | 0.395 | 0.175 | 0.287 | 0.312 | 0.458 | 0.351 | 0.371 | 0.267 |
| opensearch-neural-sparse-encoding-multilingual-v1 | 0.629 | 0.670 | 0.740 | 0.542 | 0.558 | 0.582 | 0.486 | 0.658 | 0.740 | 0.562 | 0.514 | 0.669 | 0.767 | 0.768 | 0.607 | 0.575 |
| opensearch-neural-sparse-encoding-multilingual-v1; prune_ratio 0.1 | 0.626 | 0.667 | 0.740 | 0.537 | 0.555 | 0.576 | 0.481 | 0.655 | 0.737 | 0.558 | 0.511 | 0.664 | 0.761 | 0.766 | 0.604 | 0.572 |