Datasets

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.2.0).

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

検索インデックス

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)

Update on GitHub

←クラウドストレージ CLI→