Add searchable datasets
State: closed | Author: lhoestq (MEMBER)

# Better support for Numpy format + Add Indexed Datasets

I was working on adding Indexed Datasets but in the meantime I had to also add more support for Numpy arrays in the lib.

## Better support for Numpy format

New features:
- New fast method to convert Numpy arrays from Arrow structure (up to x100 speed up) using Pandas.
- Allow to output Numpy arrays in batched `.map`, which was the only missing part to fully support Numpy arrays.

Pandas offers fast zero-copy Numpy arrays conversion from Arrow structures.
Using it we can speed up the reading of memory-mapped Numpy array stored in Arrow format.

With these changes you can easily compute embeddings of texts using `.map()`. For example:
```python
def embed(text):
    tokenized_example = tokenizer.encode(text, return_tensors="pt")
    embeddings = bert_encoder(tokenized_examples).numpy()
    return embeddings
dset_with_embeddings = dset.map(lambda example: {"embeddings": embed(example["text])})
```
And then reading the embeddings from the arrow format is be very fast.

PS1: Note that right now only 1d arrays are supported.
PS2: It seems possible to do without pandas but it will require more _trickery_.
PS3: I did a simple benchmark with google colab that you can view here:
https://colab.research.google.com/drive/1QlLTR6LRwYOKGJ-hTHmHyolE3wJzvfFg?usp=sharing

## Add Indexed Datasets

For many retrieval tasks it is convenient to index a dataset to be able to run fast queries.
For example for models like DPR, REALM, RAG etc. that are models for Open Domain QA, the retrieval step is very important.

Therefore I added two ways to add an index to a column of a dataset:
1) You can index it using a Dense Index like Faiss. It is used to index vectors.
    Faiss is a library for efficient similarity search and clustering of dense vectors.
    It contains algorithms that search in sets of vectors of any size, up to ones that possibly do not fit in RAM.
2) You can index it using a Sparse Index like Elasticsearch. It is used to index text and run queries based on BM25 similarity.

Example of usage:

```python
ds = nlp.load_dataset('crime_and_punish', split='train')
ds_with_embeddings = ds.map(lambda example: {'embeddings': embed(example['line']}))  # `embed` outputs a `np.array`
ds_with_embeddings.add_vector_index(column='embeddings')
scores, retrieved_examples = ds_with_embeddings.get_nearest(column='embeddings', query=embed('my new query'), k=10)
```

```python
ds = nlp.load_dataset('crime_and_punish', split='train')
es_client = elasticsearch.Elasticsearch()
ds.add_text_index(column='line', es_client=es_client, index_name="my_es_index")
scores, retrieved_examples = ds.get_nearest(column='line', query='my new query', k=10)
```

PS4: Faiss allows to specify many options for the [index](https://github.com/facebookresearch/faiss/wiki/The-index-factory) and for [GPU settings](https://github.com/facebookresearch/faiss/wiki/Faiss-on-the-GPU). I made sure that the user has full control over those settings.

## Tests

I added tests for Faiss, Elasticsearch and indexed datasets.
I had to edit the CI config because all the test scripts were not being run by CircleCI.

------------------

I'd be really happy to have some feedbacks :)

Comments: Looks very cool! Only looked at it superficially though

Alright I think I've checked all your comments, thanks :)

Moreover I just added a way to serialize faiss indexes.
This is important because for big datasets the index construction can take some time.

Examples:

```python
ds = nlp.load_dataset('crime_and_punish', split='train')
ds_with_embeddings = ds.map(lambda example: {'embeddings': embed(example['line']}))
ds_with_embeddings.add_faiss_index(column='embeddings')
# query
scores, retrieved_examples = ds_with_embeddings.get_nearest_examples('embeddings', embed('my new query'), k=10)
# save index
ds_with_embeddings.get_index('embeddings').save('my_index.faiss')
```

```python
ds = nlp.load_dataset('crime_and_punish', split='train')
# load index
faiss_index = nlp.search.FaissIndex.load('my_index.faiss')
ds.add_faiss_index('embeddings', faiss_index=faiss_index)
# query
scores, retrieved_examples = ds.get_nearest_examples('embeddings', embed('my new query'), k=10)
```

Let me know what you think

Nice!

Here are a few comments:

I think it would be good to separate (1) the name of the column we use for indexing and (2) the name of the index itself, at least in our head. As I understand it, once the index is created, the column we used to create it is irrelevant so the column name will only be relevant in the `add_faiss_index` and we should be able to supply a different index name, e.g. `my_faiss_index`. When we reload an index, we don't really care about the column that was used to create it, right? so it's maybe better to have an `index_name` (which default to the column name for a simple user experience but it can also be something else and this should be clear in our head when we define the API).

I'm wondering if we should not have a triple of methods for each retrieval engine: `add_xxx_index`, `save_xxx_index` and `load_xxx_index` when `xxx` can be `faiss` or `elasticsearch`. I'm not a fan of exposing `nlp.search.FaissIndex` unless you think there is a strong reason to have the user learn this abstraction.

Last but not least, I think we should already think about hosting index on our S3. I would maybe go for something like this: host the index serialized with the cached dataset on user-provided namespaces:
```python
wiki_indexed = load_dataset('thom/wiki_indexed_with_dpr_faiss')
```

I agree, I just changed to using `index_name` and having add/save/load methods

To summarize:


```python
ds = nlp.load_dataset('crime_and_punish', split='train')
ds_with_embeddings = ds.map(lambda example: {'embeddings': embed(example['line']}))
ds_with_embeddings.add_faiss_index(column='embeddings')
# query
scores, retrieved_examples = ds_with_embeddings.get_nearest_examples('embeddings', embed('my new query'), k=10)
# save index
ds_with_embeddings.save_faiss_index('embeddings', 'my_index.faiss')
```

```python
ds = nlp.load_dataset('crime_and_punish', split='train')
# load index
ds.load_faiss_index('embeddings', 'my_index.faiss')
# query
scores, retrieved_examples = ds.get_nearest_examples('embeddings', embed('my new query'), k=10)
```

Good to me. I understand that for now there is no check that the index matches the dataset on loading.
Maybe just add a basic test on the number of examples?

Ok I think this one is ready now

Looks like the CI is having troubles to pass because of `tests/test_dataset_common.py::AWSDatasetTest::test_builder_configs_{<insert_rando_dataset_name_here>}`, `requests.exceptions.ConnectionError` :/