Allow Custom Dataset in RAG Retriever
Author: lhoestq
Pull request: yes

As asked in #7725 , #7462 and #7631 , I added a way to let users build and load their own knowledge dataset for RAG.
I also added an example script that shows how to do that from csv files.

Before merging I'd like to make sure it creates the exact same embeddings that the one that were computed by the RAG team. I might need to do adjustments to the tokenization and maybe change the DPR encoder from the one trained on Natural Questions to the one trained ont the multiset/hybrid dataset.

Any feedbacks on the example and the HFIndex changes are welcome !

More details about the changes:
Previously the HFIndex only allowed to load existing datasets ("canonical" datasets) from the datasets library. So I split it into two classes `CanonicalHFIndex` to load canonical datasets and `CustomHFIndex` for custom user-defined ones.
Moreover the `config.index_name` used to accept any canonical dataset name, or "legacy" for the index the RAG team first provided. Now `config.index_name` can also accept "custom" for custom user-defined indexed datasets.

Comments:
I took your comments into account, let me know if you have other things to improve.
Also I had to change the DPR encoder from the one trained on Natural Questions to the one trained ont the multiset/hybrid dataset to match the embeddings used by the Rag team.

can't wait to try this out. 
@lhoestq , can this functionality be adapted to address #6399 ?

Glad this got implemented! Many thanks @lhoestq . I checked out a copy, added a custom 25MB CSV file, and gave it a run:
python examples/rag/use_own_knowledge_dataset.py

Got this.
```
Traceback (most recent call last):
  File "examples/rag/use_own_knowledge_dataset.py", line 195, in <module>
    main(tmp_dir, rag_example_args, processing_args, index_hnsw_args)
  File "examples/rag/use_own_knowledge_dataset.py", line 84, in main
    ctx_tokenizer = DPRContextEncoderTokenizerFast.from_pretrained(rag_example_args.dpr_ctx_encoder_model_name)
  File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/transformers/tokenization_utils_base.py", line 1544, in from_pretrained
    list(cls.vocab_files_names.values()),
OSError: Model name 'facebook/dpr-ctx_encoder-multiset-base' was not found in tokenizers model name list (facebook/dpr-ctx_encoder-single-nq-base). We assumed 'facebook/dpr-ctx_encoder-multiset-base' was a path, a model identifier, or url to a directory containing vocabulary files named ['vocab.txt'] but couldn't find such vocabulary files at this path or url.
```

I then switched to **dpr-ctx_encoder-single-nq-base**. In this case, indexing was successful. However, after Step 3 - Load RAG, the script started downloading the 74G wiki dataset (which should not be necessary) and then errored out on line 195.

```
$ python examples/rag/use_own_knowledge_dataset.py
INFO:__main__:Step 1 - Create the dataset
Using custom data configuration default
Reusing dataset csv (/home/ioannis/.cache/huggingface/datasets/csv/default-49c04e2dbd1cfa6f/0.0.0/49187751790fa4d820300fd4d0707896e5b941f1a9c644652645b866716a4ac4)
100%|█████████████████████████████████████████| 155/155 [00:01<00:00, 82.27ba/s]
100%|█████████████████████████████████████| 18531/18531 [28:23<00:00, 10.88ba/s]
INFO:__main__:Step 2 - Index the dataset
100%|█████████████████████████████████████████| 297/297 [30:22<00:00,  6.14s/it]
INFO:__main__:Step 3 - Load RAG
Using custom data configuration psgs_w100.nq.no_index
Downloading and preparing dataset wiki_dpr/psgs_w100.nq.no_index (download: Unknown size, generated: Unknown size, post-processed: Unknown size, total: Unknown size) to /home/ioannis/.cache/huggingface/datasets/wiki_dpr/psgs_w100.nq.no_index/0.0.0/14b973bf2a456087ff69c0fd34526684eed22e48e0dfce4338f9a22b965ce7c2...
Downloading: 100%|█████████████████████████| 11.2k/11.2k [00:00<00:00, 2.65MB/s]
Downloading: 100%|███████████████████████| 78.4G/78.4G [1:54:59<00:00, 11.4MB/s]
Dataset wiki_dpr downloaded and prepared to /home/ioannis/.cache/huggingface/datasets/wiki_dpr/psgs_w100.nq.no_index/0.0.0/14b973bf2a456087ff69c0fd34526684eed22e48e0dfce4338f9a22b965ce7c2. Subsequent calls will reuse this data.
Using custom data configuration psgs_w100.nq.custom
Downloading and preparing dataset wiki_dpr/psgs_w100.nq.custom (download: Unknown size, generated: Unknown size, post-processed: Unknown size, total: Unknown size) to /home/ioannis/.cache/huggingface/datasets/wiki_dpr/psgs_w100.nq.custom/0.0.0/14b973bf2a456087ff69c0fd34526684eed22e48e0dfce4338f9a22b965ce7c2...
huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
To disable this warning, you can either:
	- Avoid using `tokenizers` before the fork if possible
	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
...
...
Downloading: 100%|██████████████████████████| 1.33G/1.33G [29:21<00:00, 753kB/s]
Downloading: 100%|██████████████████████████| 1.33G/1.33G [29:30<00:00, 749kB/s]
...
...
Traceback (most recent call last):            
  File "examples/rag/use_own_knowledge_dataset.py", line 195, in <module>
    main(tmp_dir, rag_example_args, processing_args, index_hnsw_args)
  File "examples/rag/use_own_knowledge_dataset.py", line 116, in main
    rag_example_args.rag_model_name, index_name="custom", indexed_dataset=dataset
  File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/transformers/retrieval_rag.py", line 321, in from_pretrained
    config, question_encoder_tokenizer=question_encoder_tokenizer, generator_tokenizer=generator_tokenizer
  File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/transformers/retrieval_rag.py", line 310, in __init__
    self.init_retrieval()
  File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/transformers/retrieval_rag.py", line 338, in init_retrieval
    self.index.init_index()
  File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/transformers/retrieval_rag.py", line 248, in init_index
    dummy=self.use_dummy_dataset,
  File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/datasets/load.py", line 611, in load_dataset
    ignore_verifications=ignore_verifications,
  File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/datasets/builder.py", line 476, in download_and_prepare
    dl_manager=dl_manager, verify_infos=verify_infos, **download_and_prepare_kwargs
  File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/datasets/builder.py", line 553, in _download_and_prepare
    self._prepare_split(split_generator, **prepare_split_kwargs)
  File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/datasets/builder.py", line 841, in _prepare_split
    generator, unit=" examples", total=split_info.num_examples, leave=False, disable=not_verbose
  File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/tqdm/std.py", line 1133, in __iter__
    for obj in iterable:
  File "/home/ioannis/.cache/huggingface/modules/datasets_modules/datasets/wiki_dpr/14b973bf2a456087ff69c0fd34526684eed22e48e0dfce4338f9a22b965ce7c2/wiki_dpr.py", line 124, in _generate_examples
    id, text, title = line.strip().split("\t")
ValueError: not enough values to unpack (expected 3, got 2)

```


> Glad this got implemented! Many thanks @lhoestq . I checked out a copy, added a custom 25MB CSV file, and gave it a run:
> python examples/rag/use_own_knowledge_dataset.py
> 
> Got this.
> 
> ```
> Traceback (most recent call last):
>   File "examples/rag/use_own_knowledge_dataset.py", line 195, in <module>
>     main(tmp_dir, rag_example_args, processing_args, index_hnsw_args)
>   File "examples/rag/use_own_knowledge_dataset.py", line 84, in main
>     ctx_tokenizer = DPRContextEncoderTokenizerFast.from_pretrained(rag_example_args.dpr_ctx_encoder_model_name)
>   File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/transformers/tokenization_utils_base.py", line 1544, in from_pretrained
>     list(cls.vocab_files_names.values()),
> OSError: Model name 'facebook/dpr-ctx_encoder-multiset-base' was not found in tokenizers model name list (facebook/dpr-ctx_encoder-single-nq-base). We assumed 'facebook/dpr-ctx_encoder-multiset-base' was a path, a model identifier, or url to a directory containing vocabulary files named ['vocab.txt'] but couldn't find such vocabulary files at this path or url.
> ```
> 
> I then switched to **dpr-ctx_encoder-single-nq-base**. In this case, indexing was successful. However, after Step 3 - Load RAG, the script started downloading the 74G wiki dataset (which should not be necessary) and then errored out on line 195.
> 
> ```
> $ python examples/rag/use_own_knowledge_dataset.py
> INFO:__main__:Step 1 - Create the dataset
> Using custom data configuration default
> Reusing dataset csv (/home/ioannis/.cache/huggingface/datasets/csv/default-49c04e2dbd1cfa6f/0.0.0/49187751790fa4d820300fd4d0707896e5b941f1a9c644652645b866716a4ac4)
> 100%|█████████████████████████████████████████| 155/155 [00:01<00:00, 82.27ba/s]
> 100%|█████████████████████████████████████| 18531/18531 [28:23<00:00, 10.88ba/s]
> INFO:__main__:Step 2 - Index the dataset
> 100%|█████████████████████████████████████████| 297/297 [30:22<00:00,  6.14s/it]
> INFO:__main__:Step 3 - Load RAG
> Using custom data configuration psgs_w100.nq.no_index
> Downloading and preparing dataset wiki_dpr/psgs_w100.nq.no_index (download: Unknown size, generated: Unknown size, post-processed: Unknown size, total: Unknown size) to /home/ioannis/.cache/huggingface/datasets/wiki_dpr/psgs_w100.nq.no_index/0.0.0/14b973bf2a456087ff69c0fd34526684eed22e48e0dfce4338f9a22b965ce7c2...
> Downloading: 100%|█████████████████████████| 11.2k/11.2k [00:00<00:00, 2.65MB/s]
> Downloading: 100%|███████████████████████| 78.4G/78.4G [1:54:59<00:00, 11.4MB/s]
> Dataset wiki_dpr downloaded and prepared to /home/ioannis/.cache/huggingface/datasets/wiki_dpr/psgs_w100.nq.no_index/0.0.0/14b973bf2a456087ff69c0fd34526684eed22e48e0dfce4338f9a22b965ce7c2. Subsequent calls will reuse this data.
> Using custom data configuration psgs_w100.nq.custom
> Downloading and preparing dataset wiki_dpr/psgs_w100.nq.custom (download: Unknown size, generated: Unknown size, post-processed: Unknown size, total: Unknown size) to /home/ioannis/.cache/huggingface/datasets/wiki_dpr/psgs_w100.nq.custom/0.0.0/14b973bf2a456087ff69c0fd34526684eed22e48e0dfce4338f9a22b965ce7c2...
> huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
> To disable this warning, you can either:
> 	- Avoid using `tokenizers` before the fork if possible
> 	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
> huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
> ...
> ...
> Downloading: 100%|██████████████████████████| 1.33G/1.33G [29:21<00:00, 753kB/s]
> Downloading: 100%|██████████████████████████| 1.33G/1.33G [29:30<00:00, 749kB/s]
> ...
> ...
> Traceback (most recent call last):            
>   File "examples/rag/use_own_knowledge_dataset.py", line 195, in <module>
>     main(tmp_dir, rag_example_args, processing_args, index_hnsw_args)
>   File "examples/rag/use_own_knowledge_dataset.py", line 116, in main
>     rag_example_args.rag_model_name, index_name="custom", indexed_dataset=dataset
>   File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/transformers/retrieval_rag.py", line 321, in from_pretrained
>     config, question_encoder_tokenizer=question_encoder_tokenizer, generator_tokenizer=generator_tokenizer
>   File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/transformers/retrieval_rag.py", line 310, in __init__
>     self.init_retrieval()
>   File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/transformers/retrieval_rag.py", line 338, in init_retrieval
>     self.index.init_index()
>   File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/transformers/retrieval_rag.py", line 248, in init_index
>     dummy=self.use_dummy_dataset,
>   File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/datasets/load.py", line 611, in load_dataset
>     ignore_verifications=ignore_verifications,
>   File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/datasets/builder.py", line 476, in download_and_prepare
>     dl_manager=dl_manager, verify_infos=verify_infos, **download_and_prepare_kwargs
>   File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/datasets/builder.py", line 553, in _download_and_prepare
>     self._prepare_split(split_generator, **prepare_split_kwargs)
>   File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/datasets/builder.py", line 841, in _prepare_split
>     generator, unit=" examples", total=split_info.num_examples, leave=False, disable=not_verbose
>   File "/home/ioannis/anaconda3/envs/transformers/lib/python3.6/site-packages/tqdm/std.py", line 1133, in __iter__
>     for obj in iterable:
>   File "/home/ioannis/.cache/huggingface/modules/datasets_modules/datasets/wiki_dpr/14b973bf2a456087ff69c0fd34526684eed22e48e0dfce4338f9a22b965ce7c2/wiki_dpr.py", line 124, in _generate_examples
>     id, text, title = line.strip().split("\t")
> ValueError: not enough values to unpack (expected 3, got 2)
> ```

I am facing the same issue.

Thanks for reporting I'm looking into it

> Glad this got implemented! Many thanks @lhoestq . I checked out a copy, added a custom 25MB CSV file, and gave it a run:
> python examples/rag/use_own_knowledge_dataset.py
> 
> Got this.
> 
> ```
> OSError: Model name 'facebook/dpr-ctx_encoder-multiset-base' was not found in tokenizers model name list (facebook/dpr-ctx_encoder-single-nq-base). We assumed 'facebook/dpr-ctx_encoder-multiset-base' was a path, a model identifier, or url to a directory containing vocabulary files named ['vocab.txt'] but couldn't find such vocabulary files at this path or url.
> ```
> 
> I then switched to **dpr-ctx_encoder-single-nq-base**. In this case, indexing was successful. However, after Step 3 - Load RAG, the script started downloading the 74G wiki dataset (which should not be necessary) and then errored out on line 195.
> 
> ```
> $ python examples/rag/use_own_knowledge_dataset.py
> ...
> ValueError: not enough values to unpack (expected 3, got 2)
> ```

You're having this issue because you are running the script with a version of transformers that doesn't include the changes I had to make in this PR to support custom datasets. This PR not only adds an example script, but there are also changes to make it possible in the core code of the RAG retriever.

Everything works fine if you have all the changes of this PR

> > Glad this got implemented! Many thanks @lhoestq . I checked out a copy, added a custom 25MB CSV file, and gave it a run:
> > python examples/rag/use_own_knowledge_dataset.py
> > Got this.
> > ```
> > OSError: Model name 'facebook/dpr-ctx_encoder-multiset-base' was not found in tokenizers model name list (facebook/dpr-ctx_encoder-single-nq-base). We assumed 'facebook/dpr-ctx_encoder-multiset-base' was a path, a model identifier, or url to a directory containing vocabulary files named ['vocab.txt'] but couldn't find such vocabulary files at this path or url.
> > ```
> > 
> > 
> > I then switched to **dpr-ctx_encoder-single-nq-base**. In this case, indexing was successful. However, after Step 3 - Load RAG, the script started downloading the 74G wiki dataset (which should not be necessary) and then errored out on line 195.
> > ```
> > $ python examples/rag/use_own_knowledge_dataset.py
> > ...
> > ValueError: not enough values to unpack (expected 3, got 2)
> > ```
> 
> You're having this issue because you are running the script with a version of transformers that doesn't include the changes I had to make in this PR to support custom datasets. This PR not only adds an example script, but there are also changes to make it possible in the core code of the RAG retriever.
> 
> Everything works fine if you have all the changes of this PR

Having a separate script to fine-tune with custom datasets would be super useful!

> Having a separate script to fine-tune with custom datasets would be super useful!

I am adding flags to the fine-tuning scripts to make it work with a custom retriever ;)

Ah, I think my mistake was that I was using the previous conda environment of transformers (with the new branch). Trying this again with a new env now :)

> > I just checkd out [14420dc](https://github.com/huggingface/transformers/commit/14420dcbfa0f06012152cc66a7af84d7165a2a17) (committed one hour ago) and I am still getting the same error. Also, it was this branch that was giving the previous errors, not the main transformers branch. Perhaps some of your files have not been updated in the repository? Confused :)
> > `OSError: Model name 'facebook/dpr-ctx_encoder-multiset-base' was not found in tokenizers model name list (facebook/dpr-ctx_encoder-single-nq-base). We assumed 'facebook/dpr-ctx_encoder-multiset-base' was a path, a model identifier, or url to a directory containing vocabulary files named ['vocab.txt'] but couldn't find such vocabulary files at this path or url.`
> 
> Could you try to run
> 
> ```python
> from transformers import DPRContextEncoderTokenizerFast
> tokenizer = DPRContextEncoderTokenizerFast.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
> ```
> 
> and let me know if you're having the OSError on [14420dc](https://github.com/huggingface/transformers/commit/14420dcbfa0f06012152cc66a7af84d7165a2a17) ?
> 
> Also just to help me fix this issue, could you also tell me the input of this code please:
> 
> ```python
> from transformers.tokenization_dpr_fast import CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
> print(CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES)
> # it should print {'facebook/dpr-ctx_encoder-single-nq-base': 512, 'facebook/dpr-ctx_encoder-multiset-base': 512}
> ```

I was using the old transformers conda env. Doing things over again now and will report back asap!

Ok thanks !
I deleted my comment since I just noticed yours about the wrong env.
It should work for you now in a new env

I added some tests for the distributed retriever for fine-tuning.
Unless you still have an issue @ioannist this PR should be ready.
Cc @patrickvonplaten if you want to take a look at the new changes

> Ok thanks !
> I deleted my comment since I just noticed yours about the wrong env.
> It should work for you now in a new env

I created a new conda env from 14420dcb, installed pytorch and tf 2.2 (2.0 did not work). I ran again into the OS error. Here is the output on the tests you asked me to run +1 more.

(HEAD detached at 14420dcb)

```
from transformers import DPRContextEncoderTokenizerFast
tokenizer = DPRContextEncoderTokenizerFast.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
```
No error

```
from transformers.tokenization_dpr_fast import CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
print(CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES)
```
ModuleNotFoundError: No module named 'transformers.tokenization_dpr_fast'
I checked if the tokenization_dpr_fast file is under src/transformers. It's there.

```
from transformers.tokenization_dpr import CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
print(CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES)
```
{'facebook/dpr-ctx_encoder-single-nq-base': 512}


Give me an hour to test this more. I may be messing up somewhere and don't wanna waste your time..

If you can't import `transformers.tokenization_dpr_fast` that must be an environment issue. This file was added recently on the master branch. It looks like your env imports a version of transformers that is not up-to-date. That's why you're having an OSError.

Maybe you can run `pip show transformers` to check the location of your transformers installation ?

> If you can't import `transformers.tokenization_dpr_fast` that must be an environment issue. This file was added recently on the master branch. It looks like your env imports a version of transformers that is not up-to-date. That's why you're having an OSError.
> 
> Maybe you can run `pip show transformers` to check the location of your transformers installation ?

Ok, I just did the env again for the 3rd time and it works! No idea what i messed up before. 

Step -3 Load Rag in progress :)


Slow tests pass. This one is ready to merge.


@lhoestq - feel free to merge whenever!

When I do "from transformers import DPRContextEncoder", I get an error :
File "convert_slow_tokenizer.py", line 24, in <module>
    from tokenizers.models import BPE, Unigram, WordPiece
ImportError: cannot import name 'Unigram'

Unigram is missing in tokenizers 0.8.1.rc2. Needed to update to 0.9.0

> I took your comments into account, let me know if you have other things to improve.
> Also I had to change the DPR encoder from the one trained on Natural Questions to the one trained ont the multiset/hybrid dataset to match the embeddings used by the Rag team.

@lhoestq 

Hi, can you elaborate on the change you made in the DPR bit more? My understanding is, you have pretrained the DPR with a hybrid dataset to improve the performance when encoding a custom knowledge base. 

If you have pretrained the DPR, can you please publish the code? 

Can you please refer to this issue also. 

https://github.com/huggingface/transformers/issues/8037

Thanks a lot.

There are two versions of DPR in the paper. One trained on NQ and one trained on various datasets. The authors released the code and the weights in this [this repo](https://github.com/facebookresearch/DPR).

The change I did was just to use the weight of the second one, since it was the one used for RAG.

Thanks. So as mentioned in the RAG paper, we can use the doc encoder to get
embeddings for any custom dataset.  Later we can only fine tune the BART
model and the question encoder.

On Tue, Oct 27, 2020, 01:42 Quentin Lhoest <notifications@github.com> wrote:

> There are two versions of DPR in the paper. One trained on NQ and one
> trained on various datasets. The authors released the code and the weights
> in this this repo <https://github.com/facebookresearch/DPR>.
>
> The change I did was just to use the weight of the second one, since it
> was the one used for RAG.
>
> —
> You are receiving this because you commented.
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/pull/7763#issuecomment-716521222>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AEA4FGUJ5EC6J2ROO56RXB3SMVVDPANCNFSM4SPHCGQQ>
> .
>


Yes exactly :) 

@lhoestq , I tried use_own_knowledge_dataset.py to try and retrieve passages from a custom dataset in response to queries. It works, though the relevance of the results isn't great. 
I wanted to try and use finetune.sh. Is it possible to finetune just the retriever ? Is there a sample format for the training data ? Looks like it will need the question and passages as input. Need positive and negative passages ? 
Thanks!

@mchari  He previously said (https://github.com/huggingface/transformers/issues/8037), it is possible to retrain DPR with Facebook code and then convert the checkpoint to the huge face compatible.

What if you fine-tune the RAG system with your own data letting the question encoder to get fine-tuned? It seems like the pre-training of DPR can be a hard task since the results can depend on the selection procedure of the negative samples as mentioned in the paper.