Opensearch Neural Sparse Encoding Doc V2 Distill

🚀 opensearch-neural-sparse-encoding-doc-v2-distill

このモデルは、検索の関連性、推論速度、および検索効率を考慮して設計された学習済みの疎行列検索モデルです。OpenSearchと連携して、高速かつ高精度な検索を実現します。

🚀 クイックスタート

モデルの選択には、検索の関連性、モデルの推論、および検索効率（FLOPS）を考慮する必要があります。我々は、BEIRベンチマークのサブセットであるTrecCovid、NFCorpus、NQ、HotpotQA、FiQA、ArguAna、Touche、DBPedia、SCIDOCS、FEVER、Climate FEVER、SciFact、Quoraで、モデルのゼロショット性能をベンチマークしています。

全体的に、v2シリーズのモデルは、v1シリーズよりも検索の関連性、効率、および推論速度が優れています。ただし、具体的な長所と短所は、データセットによって異なる場合があります。

✨ 主な機能

概要

• 論文: Towards Competitive Search Relevance For Inference-Free Learned Sparse Retrievers
• ファインチューニングのサンプル: opensearch-sparse-model-tuning-sample

これは学習済みの疎行列検索モデルです。ドキュメントを30522次元の疎ベクトルにエンコードします。クエリについては、トークナイザーと重みのルックアップテーブルを使用して疎ベクトルを生成します。非ゼロの次元インデックスは語彙内の対応するトークンを意味し、重みはトークンの重要度を意味します。そして、類似度スコアはクエリとドキュメントの疎ベクトルの内積です。

トレーニングデータセットには、MS MARCO、eli5_question_answer、squad_pairs、WikiAnswers、yahoo_answers_title_question、gooaq_pairs、stackexchange_duplicate_questions_body_body、wikihow、S2ORC_title_abstract、stackexchange_duplicate_questions_title-body_title-body、yahoo_answers_question_answer、searchQA_top5_snippets、stackexchange_duplicate_questions_title_title、yahoo_answers_title_answerが含まれています。

OpenSearchのニューラル疎行列機能は、Luceneの転置インデックスを使用した学習済みの疎行列検索をサポートしています。リンク: https://opensearch.org/docs/latest/query-dsl/specialized/neural-sparse/ 。インデックス作成と検索は、OpenSearchの高レベルAPIを使用して実行できます。

💻 使用例

基本的な使用法

import json
import itertools
import torch

from transformers import AutoModelForMaskedLM, AutoTokenizer


# get sparse vector from dense vectors with shape batch_size * seq_len * vocab_size
def get_sparse_vector(feature, output):
    values, _ = torch.max(output*feature["attention_mask"].unsqueeze(-1), dim=1)
    values = torch.log(1 + torch.relu(values))
    values[:,special_token_ids] = 0
    return values
    
# transform the sparse vector to a dict of (token, weight)
def transform_sparse_vector_to_dict(sparse_vector):
    sample_indices,token_indices=torch.nonzero(sparse_vector,as_tuple=True)
    non_zero_values = sparse_vector[(sample_indices,token_indices)].tolist()
    number_of_tokens_for_each_sample = torch.bincount(sample_indices).cpu().tolist()
    tokens = [transform_sparse_vector_to_dict.id_to_token[_id] for _id in token_indices.tolist()]

    output = []
    end_idxs = list(itertools.accumulate([0]+number_of_tokens_for_each_sample))
    for i in range(len(end_idxs)-1):
        token_strings = tokens[end_idxs[i]:end_idxs[i+1]]
        weights = non_zero_values[end_idxs[i]:end_idxs[i+1]]
        output.append(dict(zip(token_strings, weights)))
    return output
    
# download the idf file from model hub. idf is used to give weights for query tokens
def get_tokenizer_idf(tokenizer):
    from huggingface_hub import hf_hub_download
    local_cached_path = hf_hub_download(repo_id="opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill", filename="idf.json")
    with open(local_cached_path) as f:
        idf = json.load(f)
    idf_vector = [0]*tokenizer.vocab_size
    for token,weight in idf.items():
        _id = tokenizer._convert_token_to_id_with_added_voc(token)
        idf_vector[_id]=weight
    return torch.tensor(idf_vector)

# load the model
model = AutoModelForMaskedLM.from_pretrained("opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill")
tokenizer = AutoTokenizer.from_pretrained("opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill")
idf = get_tokenizer_idf(tokenizer)

# set the special tokens and id_to_token transform for post-process
special_token_ids = [tokenizer.vocab[token] for token in tokenizer.special_tokens_map.values()]
get_sparse_vector.special_token_ids = special_token_ids
id_to_token = ["" for i in range(tokenizer.vocab_size)]
for token, _id in tokenizer.vocab.items():
    id_to_token[_id] = token
transform_sparse_vector_to_dict.id_to_token = id_to_token



query = "What's the weather in ny now?"
document = "Currently New York is rainy."

# encode the query
feature_query = tokenizer([query], padding=True, truncation=True, return_tensors='pt')
input_ids = feature_query["input_ids"]
batch_size = input_ids.shape[0]
query_vector = torch.zeros(batch_size, tokenizer.vocab_size)
query_vector[torch.arange(batch_size).unsqueeze(-1), input_ids] = 1
query_sparse_vector = query_vector*idf

# encode the document
feature_document = tokenizer([document], padding=True, truncation=True, return_tensors='pt')
output = model(**feature_document)[0]
document_sparse_vector = get_sparse_vector(feature_document, output)


# get similarity score
sim_score = torch.matmul(query_sparse_vector[0],document_sparse_vector[0])
print(sim_score)   # tensor(17.5307, grad_fn=<DotBackward0>)


query_token_weight = transform_sparse_vector_to_dict(query_sparse_vector)[0]
document_query_token_weight = transform_sparse_vector_to_dict(document_sparse_vector)[0]
for token in sorted(query_token_weight, key=lambda x:query_token_weight[x], reverse=True):
    if token in document_query_token_weight:
        print("score in query: %.4f, score in document: %.4f, token: %s"%(query_token_weight[token],document_query_token_weight[token],token))
        

        
# result:
# score in query: 5.7729, score in document: 1.4109, token: ny
# score in query: 4.5684, score in document: 1.4673, token: weather
# score in query: 3.5895, score in document: 0.7473, token: now

上記のコードサンプルは、ニューラル疎行列検索の例を示しています。元のクエリとドキュメントに重複するトークンがないにもかかわらず、このモデルは良好なマッチングを行っています。

📚 ドキュメント

詳細な検索関連性

📄 ライセンス

このプロジェクトは、Apache v2.0 Licenseの下でライセンスされています。

著作権

Copyright OpenSearch Contributors. 詳細はNOTICEを参照してください。