Datasets documentation
検索インデックス
検索インデックス
FAISSとElasticsearchは、データセット内の例を検索することを可能にします。これは、NLPタスクに関連する特定の例をデータセットから取得したい場合に便利です。例えば、オープンドメインの質問応答タスクに取り組んでいる場合、質問への回答に関連する例のみを返すことを望むかもしれません。
このガイドでは、データセットのインデックスを構築し、それを使って検索する方法を示します。
FAISS
FAISSは、ベクトル表現の類似性に基づいてドキュメントを検索します。この例では、DPRモデルを使用してベクトル表現を生成します。
- 🤗 TransformersからDPRモデルをダウンロードします:
>>> from transformers import DPRContextEncoder, DPRContextEncoderTokenizer
>>> import torch
>>> torch.set_grad_enabled(False)
>>> ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
>>> ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")- データセットを読み込み、ベクトル表現を計算します:
>>> from datasets import load_dataset
>>> ds = load_dataset('community-datasets/crime_and_punish', split='train[:100]')
>>> ds_with_embeddings = ds.map(lambda example: {'embeddings': ctx_encoder(**ctx_tokenizer(example["line"], return_tensors="pt"))[0][0].numpy()})- Dataset.add_faiss_index()でインデックスを作成します:
>>> ds_with_embeddings.add_faiss_index(column='embeddings')- これで、
embeddingsインデックスを使ってデータセットをクエリできます。DPR Question Encoderを読み込み、Dataset.get_nearest_examples()で質問を検索します:
>>> from transformers import DPRQuestionEncoder, DPRQuestionEncoderTokenizer
>>> q_encoder = DPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
>>> q_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
>>> question = "Is it serious ?"
>>> question_embedding = q_encoder(**q_tokenizer(question, return_tensors="pt"))[0][0].numpy()
>>> scores, retrieved_examples = ds_with_embeddings.get_nearest_examples('embeddings', question_embedding, k=10)
>>> retrieved_examples["line"][0]
'_that_ serious? It is not serious at all. It’s simply a fantasy to amuse\r\n'- Dataset.get_index() を使用してインデックスにアクセスし、
range_searchによるクエリなどの特殊な操作に利用できます。
>>> faiss_index = ds_with_embeddings.get_index('embeddings').faiss_index
>>> limits, distances, indices = faiss_index.range_search(x=question_embedding.reshape(1, -1), thresh=0.95)- クエリが完了したら、Dataset.save_faiss_index() を使用してインデックスをディスクに保存します:
>>> ds_with_embeddings.save_faiss_index('embeddings', 'my_index.faiss')- 後でDataset.load_faiss_index()を再読み込みします:
>>> ds = load_dataset('community-datasets/crime_and_punish', split='train[:100]')
>>> ds.load_faiss_index('embeddings', 'my_index.faiss')Elasticsearch
FAISSとは異なり、Elasticsearchは完全一致に基づいてドキュメントを検索します。
マシン上でElasticsearchを起動してください。まだインストールしていない場合は、Elasticsearchのインストールガイドを参照してください。
- インデックスを作成したいデータセットを読み込みます:
>>> from datasets import load_dataset
>>> squad = load_dataset('rajpurkar/squad', split='validation')- Dataset.add_elasticsearch_index()でインデックスを構築します:
>>> squad.add_elasticsearch_index("context", host="localhost", port="9200")- その後、Dataset.get_nearest_examples() を使用して
contextインデックスをクエリできます:
>>> query = "machine"
>>> scores, retrieved_examples = squad.get_nearest_examples("context", query, k=10)
>>> retrieved_examples["title"][0]
'Computational_complexity_theory'- インデックスを再利用する場合は、インデックスを構築する際に
es_index_nameパラメータを定義します。
>>> from datasets import load_dataset
>>> squad = load_dataset('rajpurkar/squad', split='validation')
>>> squad.add_elasticsearch_index("context", host="localhost", port="9200", es_index_name="hf_squad_val_context")
>>> squad.get_index("context").es_index_name
hf_squad_val_context- 後でインデックス名を指定して Dataset.load_elasticsearch_index() を呼び出すことで、再度読み込むことができます:
>>> from datasets import load_dataset
>>> squad = load_dataset('rajpurkar/squad', split='validation')
>>> squad.load_elasticsearch_index("context", host="localhost", port="9200", es_index_name="hf_squad_val_context")
>>> query = "machine"
>>> scores, retrieved_examples = squad.get_nearest_examples("context", query, k=10)より高度なElasticsearchの使用方法については、カスタム設定を含む独自の構成を指定できます。
>>> import elasticsearch as es
>>> import elasticsearch.helpers
>>> from elasticsearch import Elasticsearch
>>> es_client = Elasticsearch([{"host": "localhost", "port": "9200"}]) # default client
>>> es_config = {
... "settings": {
... "number_of_shards": 1,
... "analysis": {"analyzer": {"stop_standard": {"type": "standard", " stopwords": "_english_"}}},
... },
... "mappings": {"properties": {"text": {"type": "text", "analyzer": "standard", "similarity": "BM25"}}},
... } # default config
>>> es_index_name = "hf_squad_context" # name of the index in Elasticsearch
>>> squad.add_elasticsearch_index("context", es_client=es_client, es_config=es_config, es_index_name=es_index_name)