GitHub issue #298: Add searchable datasets
State: closed | Created: 2020-06-23T07:33:03Z

# 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 :)