Datasets documentation
Índice de pesquisa
Í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.
- 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")- 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()})- Crie o índice com Dataset.add_faiss_index():
>>> ds_with_embeddings.add_faiss_index(column='embeddings')- 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'- 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)- Quando terminar de consultar, salve o índice no disco com Dataset.save_faiss_index():
>>> ds_with_embeddings.save_faiss_index('embeddings', 'my_index.faiss')- 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.
- Carregue o conjunto de dados que você deseja indexar:
>>> from datasets import load_dataset
>>> squad = load_dataset('rajpurkar/squad', split='validation')- Crie o índice com Dataset.add_elasticsearch_index():
>>> squad.add_elasticsearch_index("context", host="localhost", port="9200")- Em seguida, você pode consultar o índice
contextcom Dataset.get_nearest_examples():
>>> query = "machine"
>>> scores, retrieved_examples = squad.get_nearest_examples("context", query, k=10)
>>> retrieved_examples["title"][0]
'Computational_complexity_theory'- Se você deseja reutilizar o índice, defina o parâmetro
es_index_nameao 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- 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)