Title: "NameError: name 'faiss' is not defined" on `.add_faiss_index` when `device` is not None
Author: alvarobartt (MEMBER)
State: closed
Created: 2022-05-05T15:09:45Z
Labels: bug

## Describe the bug

When using `datasets` to calculate the FAISS indices of a dataset, the exception `NameError: name 'faiss' is not defined` is triggered when trying to calculate those on a device (GPU), so `.add_faiss_index(..., device=0)` fails with that exception.

All that assuming that `datasets` is properly installed and `faiss-gpu` too, as well as all the CUDA drivers required.

## Steps to reproduce the bug

```python
# Sample code to reproduce the bug
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")

from datasets import load_dataset
ds = load_dataset('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()})

ds_with_embeddings.add_faiss_index(column='embeddings', device=0) # default `device=None`
```

## Expected results

A new column named `embeddings` in the dataset that we're adding the index to.

## Actual results

An exception is triggered with the following message `NameError: name 'faiss' is not defined`.

## Environment info

- `datasets` version: 2.1.0
- Platform: Linux-5.13.0-1022-azure-x86_64-with-glibc2.31
- Python version: 3.9.12
- PyArrow version: 7.0.0
- Pandas version: 1.4.2

Comments:
So I managed to solve this by adding a missing `import faiss` in the `@staticmethod` defined in https://github.com/huggingface/datasets/blob/f51b6994db27ea69261ef919fb7775928f9ec10b/src/datasets/search.py#L305, triggered from https://github.com/huggingface/datasets/blob/f51b6994db27ea69261ef919fb7775928f9ec10b/src/datasets/search.py#L249 when trying to `ds_with_embeddings.add_faiss_index(column='embeddings', device=0)` with the code above.

As it seems that the `@staticmethod` doesn't recognize the `import faiss` defined in https://github.com/huggingface/datasets/blob/f51b6994db27ea69261ef919fb7775928f9ec10b/src/datasets/search.py#L261, so whenever the value of `device` is not None in https://github.com/huggingface/datasets/blob/71f76e0bdeaddadedc4f9c8d15cfff5a36d62f66/src/datasets/search.py#L438, that exception is triggered.

So on, adding `import faiss` inside https://github.com/huggingface/datasets/blob/71f76e0bdeaddadedc4f9c8d15cfff5a36d62f66/src/datasets/search.py#L305 right after the check of `device`'s value, solves the issue and lets you calculate the indices in GPU.

I'll add the code in a PR linked to this issue in case you want to merge it!

Adding here the complete error traceback!

```
Traceback (most recent call last):
  File "/home/alvarobartt/lol.py", line 12, in <module>
    ds_with_embeddings.add_faiss_index(column='embeddings', device=0) # default `device=None`
  File "/home/alvarobartt/.local/lib/python3.9/site-packages/datasets/arrow_dataset.py", line 3656, in add_faiss_index
    super().add_faiss_index(
  File "/home/alvarobartt/.local/lib/python3.9/site-packages/datasets/search.py", line 478, in add_faiss_index
    faiss_index.add_vectors(self, column=column, train_size=train_size, faiss_verbose=True)
  File "/home/alvarobartt/.local/lib/python3.9/site-packages/datasets/search.py", line 281, in add_vectors
    self.faiss_index = self._faiss_index_to_device(index, self.device)
  File "/home/alvarobartt/.local/lib/python3.9/site-packages/datasets/search.py", line 327, in _faiss_index_to_device
    faiss_res = faiss.StandardGpuResources()
NameError: name 'faiss' is not defined
```

Closed as https://github.com/huggingface/datasets/pull/4288 already merged! :hugs: