# Índice de pesquisa

[FAISS](https://github.com/facebookresearch/faiss) e [Elasticsearch](https://www.elastic.co/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](https://huggingface.co/transformers/model_doc/dpr.html).

1. Baixe o modelo DPR do 🤗 Transformers:

```py
>>> 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")
```

2. Carregue seu conjunto de dados e calcule as representações vetoriais:

```py
>>> 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()})
```

3. Crie o índice com [Dataset.add_faiss_index()](/docs/datasets/main/pt/package_reference/main_classes#datasets.Dataset.add_faiss_index):

```py
>>> ds_with_embeddings.add_faiss_index(column='embeddings')
```

4. Agora você pode consultar seu dataset com o índice `embeddings`. Carregue o DPR Question Encoder e busque uma pergunta com [Dataset.get_nearest_examples()](/docs/datasets/main/pt/package_reference/main_classes#datasets.Dataset.get_nearest_examples):

```py
>>> 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'
```

5. Você pode acessar o índice com [Dataset.get_index()](/docs/datasets/main/pt/package_reference/main_classes#datasets.Dataset.get_index) e usá-lo para operações especiais, por exemplo, consultá-lo usando `range_search`:

```py
>>> 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)
```

6. Quando terminar de consultar, salve o índice no disco com [Dataset.save_faiss_index()](/docs/datasets/main/pt/package_reference/main_classes#datasets.Dataset.save_faiss_index):

```py
>>> ds_with_embeddings.save_faiss_index('embeddings', 'my_index.faiss')
```

7. Recarregue-o mais tarde com [Dataset.load_faiss_index()](/docs/datasets/main/pt/package_reference/main_classes#datasets.Dataset.load_faiss_index):

```py
>>> 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](https://www.elastic.co/guide/en/elasticsearch/reference/current/setup.html) caso ainda não o tenha instalado.

1. Carregue o conjunto de dados que você deseja indexar:

```py
>>> from datasets import load_dataset
>>> squad = load_dataset('rajpurkar/squad', split='validation')
```

2. Crie o índice com [Dataset.add_elasticsearch_index()](/docs/datasets/main/pt/package_reference/main_classes#datasets.Dataset.add_elasticsearch_index):

```py
>>> squad.add_elasticsearch_index("context", host="localhost", port="9200")
```

3. Em seguida, você pode consultar o índice `context` com [Dataset.get_nearest_examples()](/docs/datasets/main/pt/package_reference/main_classes#datasets.Dataset.get_nearest_examples):

```py
>>> query = "machine"
>>> scores, retrieved_examples = squad.get_nearest_examples("context", query, k=10)
>>> retrieved_examples["title"][0]
'Computational_complexity_theory'
```

4. Se você deseja reutilizar o índice, defina o parâmetro `es_index_name` ao construir o índice:

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

5. Recarregue-o posteriormente com o nome do índice ao chamar [Dataset.load_elasticsearch_index()](/docs/datasets/main/pt/package_reference/main_classes#datasets.Dataset.load_elasticsearch_index):

```py
>>> 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:

```py
>>> 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)
```