Datasets documentation

Índice de pesquisa

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

Índice de pesquisa

FAISS e Elasticsearch permitem buscar exemplos em um conjunto de dados. Isso pode ser útil quando você deseja recuperar exemplos específicos de um conjunto de dados que sejam relevantes para sua tarefa de NLP. Por exemplo, se você estiver trabalhando em uma tarefa de Resposta a Perguntas em Domínio Aberto, pode querer retornar apenas exemplos que sejam relevantes para responder à sua pergunta.

Este guia mostrará como construir um índice para o seu conjunto de dados que permitirá pesquisá-lo.

FAISS

O FAISS recupera documentos com base na similaridade de suas representações vetoriais. Neste exemplo, você gerará as representações vetoriais com o modelo DPR.

  1. Baixe o modelo DPR do 🤗 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. Carregue seu conjunto de dados e calcule as representações vetoriais:
>>> 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. Crie o índice com Dataset.add_faiss_index():
>>> ds_with_embeddings.add_faiss_index(column='embeddings')
  1. Agora você pode consultar seu dataset com o índice embeddings. Carregue o DPR Question Encoder e busque uma pergunta com 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. Você pode acessar o índice com Dataset.get_index() e usá-lo para operações especiais, por exemplo, consultá-lo usando 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. Quando terminar de consultar, salve o índice no disco com Dataset.save_faiss_index():
>>> ds_with_embeddings.save_faiss_index('embeddings', 'my_index.faiss')
  1. Recarregue-o mais tarde com Dataset.load_faiss_index():
>>> ds = load_dataset('community-datasets/crime_and_punish', split='train[:100]')
>>> ds.load_faiss_index('embeddings', 'my_index.faiss')

Elasticsearch

Diferentemente do FAISS, o Elasticsearch recupera documentos com base em correspondências exatas.

Inicie o Elasticsearch em sua máquina ou consulte o guia de instalação do Elasticsearch caso ainda não o tenha instalado.

  1. Carregue o conjunto de dados que você deseja indexar:
>>> from datasets import load_dataset
>>> squad = load_dataset('rajpurkar/squad', split='validation')
  1. Crie o índice com Dataset.add_elasticsearch_index():
>>> squad.add_elasticsearch_index("context", host="localhost", port="9200")
  1. Em seguida, você pode consultar o índice context com 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. Se você deseja reutilizar o índice, defina o parâmetro es_index_name ao construir o índice:
>>> 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. Recarregue-o posteriormente com o nome do índice ao chamar 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)

Para um uso mais avançado do Elasticsearch, você pode especificar sua própria configuração com configurações personalizadas:

>>> 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