Datasets documentation

Индекс поиска

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v4.8.4).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Индекс поиска

FAISS и Elasticsearch позволяют искать примеры в наборе данных. Это может быть полезно, когда вам нужно извлечь из набора данных конкретные примеры, релевантные вашей задаче NLP. Например, если вы работаете над задачей Open Domain Question Answering, вам может потребоваться возвращать только те примеры, которые относятся к ответу на ваш вопрос.

В этом руководстве показано, как создать индекс для набора данных, чтобы обеспечить возможность поиска по нему.

FAISS

FAISS выполняет поиск документов на основе сходства их векторных представлений. В этом примере вы сгенерируете векторные представления с помощью модели DPR.

  1. Скачайте модель DPR из 🤗 Transformers:
>>> 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")
  1. Загрузите набор данных и вычислите векторные представления:
>>> 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()})
  1. Создайте индекс с помощью Dataset.add_faiss_index():
>>> ds_with_embeddings.add_faiss_index(column='embeddings')
  1. Теперь вы можете выполнять запросы к вашему набору данных с использованием индекса 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'
  1. Вы можете получить доступ к индексу с помощью 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)
  1. Когда вы закончите с запросами, сохраните индекс на диске с помощью Dataset.save_faiss_index():
>>> ds_with_embeddings.save_faiss_index('embeddings', 'my_index.faiss')
  1. Позже загрузите его с помощью 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, если вы ещё не установили его.

  1. Загрузите набор данных, который вы хотите проиндексировать:
>>> from datasets import load_dataset
>>> squad = load_dataset('rajpurkar/squad', split='validation')
  1. Создайте индекс с помощью Dataset.add_elasticsearch_index():
>>> squad.add_elasticsearch_index("context", host="localhost", port="9200")
  1. Затем можно выполнить запрос к индексу context с помощью Dataset.get_nearest_examples():
>>> query = "machine"
>>> scores, retrieved_examples = squad.get_nearest_examples("context", query, k=10)
>>> retrieved_examples["title"][0]
'Computational_complexity_theory'
  1. Если вы хотите повторно использовать индекс, определите параметр 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
  1. Позже загрузите его с именем индекса при вызове 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