Loading a faiss index KeyError
State: closed | Author: nbroad1881 (NONE)
Labels: name: documentation

I've recently been testing out RAG and DPR embeddings, and I've run into an issue that is not apparent in the documentation.

The basic steps are:

1. Create a dataset (dataset1)
2. Create an embeddings column using DPR
3. Add a faiss index to the dataset
4. Save faiss index to a file
5. Create a new dataset (dataset2) with the same text and label information as dataset1
6. Try to load the faiss index from file to dataset2
7. Get `KeyError: "Column embeddings not in the dataset"`

I've made a colab notebook that should show exactly what I did. Please switch to GPU runtime; I didn't check on CPU.

https://colab.research.google.com/drive/1X0S9ZuZ8k0ybcoei4w7so6dS_WrABmIx?usp=sharing

Ubuntu Version
VERSION="18.04.5 LTS (Bionic Beaver)"

datasets==1.4.1
faiss==1.5.3
faiss-gpu==1.7.0
torch==1.8.0+cu101
transformers==4.3.3

NVIDIA-SMI 460.56
Driver Version: 460.32.03
CUDA Version: 11.2    
Tesla K80           

I was basically following the steps here: https://huggingface.co/docs/datasets/faiss_and_ea.html#adding-a-faiss-index

I included the exact code from the documentation at the end of the notebook to show that they don't work either.


Comments: In your code `dataset2` doesn't contain the "embeddings" column, since it is created from the pandas DataFrame with columns "text" and "label".

Therefore when you call `dataset2[embeddings_name]`, you get a `KeyError`.

If you want the "embeddings" column back, you can create `dataset2` with
```python
dataset2 = load_from_disk(dataset_filename)
```
where `dataset_filename` is the place where you saved you dataset with the embeddings in the first place.

Ok in that case HF should fix their misleading example at https://huggingface.co/docs/datasets/faiss_and_ea.html#adding-a-faiss-index 

I copy-pasted it here.

> When you are done with your queries you can save your index on disk:
> 
> ```python
> ds_with_embeddings.save_faiss_index('embeddings', 'my_index.faiss')
> ```
> Then reload it later:
> 
> ```python
> ds = load_dataset('crime_and_punish', split='train[:100]')
> ds.load_faiss_index('embeddings', 'my_index.faiss')
> ```

Hi !

The code of the example is valid.
An index is a search engine, it's not considered a column of a dataset.
When you do `ds.load_faiss_index("embeddings", 'my_index.faiss')`, it attaches an index named "embeddings" to the dataset but it doesn't re-add the "embeddings" column. You can list the indexes of a dataset by using `ds.list_indexes()`.

If I understand correctly by reading this example you thought that it was re-adding the "embeddings" column.
This looks misleading indeed, and we should add a note to make it more explicit that it doesn't store the column that was used to build the index.

Feel free to open a PR to suggest an improvement on the documentation if you want to contribute :)

> If I understand correctly by reading this example you thought that it was re-adding the "embeddings" column.
Yes. I was trying to use the dataset in RAG and it complained that the dataset didn't have the right columns. No problems when loading the dataset with `load_from_disk` and then doing `load_faiss_index`

What I learned was
1. column and index are different
2. loading the index does not create a column
3. the column is not needed to be able to use the index
4. RAG needs both the embeddings column and the index

If I can come up with a way to articulate this in the right spot in the docs, I'll open a PR