RAG
Author: ola13
Pull request: yes

# Intro
This pull request implements a RAG-sequence and a RAG-token models, as defined in [the paper](https://arxiv.org/pdf/2005.11401.pdf).

RAG (for Retrieval Augmented Generation) is an architecture combining a retriever with a generator model. During a forward pass, the input sequence is used as a question to the retriever, which surfaces relevant context documents. The documents are then prepended to the input and such contextualized input is passed to the generator. In the paper, we experiment with DPR-based retrieval and BART generator.

# Implementation
RAG is a seq2seq model which encapsulates three core components:
- a retriever - a wrapper around a faiss index of the documents,
- a question encoder which encodes the input sequence before passing it to the retriever,
- a generator which learns to generate the output from the contextualized input

as well as respective tokenizers (we need to be able to decode the input sequence, encode it with the question  encoder tokenizer, and then encode contextualized input with the generator tokenizer again). 

---
We implement two variants of the model, both presented in the paper:
- `RagSequence`, which uses `DPRQuestionEncoder` as the question encoder. As for the generator, two compatible architectures have been tested: `BartForConditionalGeneration`  and `T5ForConditionalGeneration`.
- `RagToken`, which uses `DPRQuestionEncoder` and `BartForConditionalGeneration` as the generator.

---
Key files in the pull request:
- `modeling_rag.py`, `tokenization_rag.py`, `configuration_rag.py` the core model implementation
-  `retrieval_rag.py` - a distributed retriever built on top of the `torch.distributed` communication package. The retriever is an interface between the model and the faiss index of the encoded documents. During training, all workers initialize their own instance of the retriever, however, only the main worker loads the index into memory, which prevents OOMs on machines with multiple GPUs (we store the index in RAM). The index itself is based on the `nlp.Datasets`. We also implement a variant compatible with indices built using the original DPR implementation (https://github.com/facebookresearch/DPR)
- `eval_rag.py` - an evaluation script which allows to perform the evaluation end to end (measures the exact match and F1 on the downstream task) as well as the evaluation of the retrieval component alone (measures precision@k).
- `finetune.py` - a training script for finetuning RAG models.

# Testing
We have successfully managed to reproduce original paper results on Natural Questions for a couple of scenarios:
- converting original `fairseq` checkpoints to `HuggingFace` and evaluating them on `HuggingFace`
- converting original `fairseq` checkpoints to `HuggingFace` and continuing fine-tuning on `HuggingFace`
- training from scratch on `HuggingFace`

# Pretrained, ready-to-use models (after PR is merged).

- RagToken: https://huggingface.co/facebook/rag-token-nq
- RagSequence: https://huggingface.co/facebook/rag-sequence-nq
- RagTokenBase: https://huggingface.co/facebook/rag-token-base
- RagSequneceBase: https://huggingface.co/facebook/rag-sequence-base

# Future PR
- [ ] Add and test distributed Pytorch and possible Ray - Retriever @lhoestq 
- [ ] Upload more rag model combinations.
- [ ] Clean examples


Comments:
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/6813?src=pr&el=h1) Report
> Merging [#6813](https://codecov.io/gh/huggingface/transformers/pull/6813?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/244e1b5ba331cb4c1ed96d88d0895c252567f7f3?el=desc) will **decrease** coverage by `0.85%`.
> The diff coverage is `82.89%`.

[![Impacted file tree graph](https://codecov.io/gh/huggingface/transformers/pull/6813/graphs/tree.svg?width=650&height=150&src=pr&token=9qOlN6Hb1c)](https://codecov.io/gh/huggingface/transformers/pull/6813?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #6813      +/-   ##
==========================================
- Coverage   78.81%   77.95%   -0.86%     
==========================================
  Files         174      178       +4     
  Lines       33670    34125     +455     
==========================================
+ Hits        26537    26603      +66     
- Misses       7133     7522     +389     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/6813?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/tokenization\_rag.py](https://codecov.io/gh/huggingface/transformers/pull/6813/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fcmFnLnB5) | `69.76% <69.76%> (ø)` | |
| [src/transformers/modeling\_rag.py](https://codecov.io/gh/huggingface/transformers/pull/6813/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19yYWcucHk=) | `76.98% <76.98%> (ø)` | |
| [src/transformers/testing\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6813/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90ZXN0aW5nX3V0aWxzLnB5) | `67.28% <77.77%> (+0.40%)` | :arrow_up: |
| [src/transformers/retrieval\_rag.py](https://codecov.io/gh/huggingface/transformers/pull/6813/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9yZXRyaWV2YWxfcmFnLnB5) | `91.27% <91.27%> (ø)` | |
| [src/transformers/\_\_init\_\_.py](https://codecov.io/gh/huggingface/transformers/pull/6813/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9fX2luaXRfXy5weQ==) | `99.37% <100.00%> (+0.01%)` | :arrow_up: |
| [src/transformers/configuration\_auto.py](https://codecov.io/gh/huggingface/transformers/pull/6813/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9jb25maWd1cmF0aW9uX2F1dG8ucHk=) | `96.25% <100.00%> (+0.09%)` | :arrow_up: |
| [src/transformers/configuration\_dpr.py](https://codecov.io/gh/huggingface/transformers/pull/6813/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9jb25maWd1cmF0aW9uX2Rwci5weQ==) | `100.00% <100.00%> (ø)` | |
| [src/transformers/configuration\_rag.py](https://codecov.io/gh/huggingface/transformers/pull/6813/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9jb25maWd1cmF0aW9uX3JhZy5weQ==) | `100.00% <100.00%> (ø)` | |
| [src/transformers/file\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6813/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9maWxlX3V0aWxzLnB5) | `83.37% <100.00%> (+0.53%)` | :arrow_up: |
| [src/transformers/modeling\_auto.py](https://codecov.io/gh/huggingface/transformers/pull/6813/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19hdXRvLnB5) | `83.96% <100.00%> (+1.58%)` | :arrow_up: |
| ... and [24 more](https://codecov.io/gh/huggingface/transformers/pull/6813/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/6813?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/huggingface/transformers/pull/6813?src=pr&el=footer). Last update [3ebb1b3...db3e5e0](https://codecov.io/gh/huggingface/transformers/pull/6813?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


@lhoestq thanks for the comments! I did consider moving retrieval outside of the model - the benefit of this that I see would be that we would move all training-related logic (e.g. handling distributed processing in Retriever) from `transformers` to `examples`.

That said, I'm still in favor of keeping the call to `contextualize` as part of the forward pass. Here's my thinking:
- retrieval is more than just data pre-processing step, it is a core part of the model's architecture. E.g. we can't pre-compute retrieved docs for a batch of data beforehand as the question encoder will be updated at every step of training, so the set of retrieved docs would be changing dynamically. If we move retrieval outside of the model people may be tempted to do that.
- we would need to call `contextualize` before every forward pass on the model, so not only in finetuning, but also e.g. in evaluation code. On top of that anyone who would want to run the model for demo purposes would have to instantiate the retriever first and remember to call `contextualize`, instead of doing the two simple steps that other HF models require (encoding the sequence and running the model) - we could potentially consider making contextualization a part of tokenizer's `encode` method (not sure this would be intuitive for people used to HF's APIs) - however, the retrieval logic would still remain in `transformers` then
- In terms of flexibility - I think with the current approach it'd still be possible for people to build different retrievers and pass them to the model

What do you think? I'd be curious to know what others think about it, cc @patrick-s-h-lewis, @thomwolf

So we've been brainstorming with @patrickvonplaten and @lhoestq on this yesterday and here is a proposal.

The general context is that we expect to include more hybrid models with a retrieval component in the future (like REALM, MARGE, knn-LM) so it's nice if we can setup an API which is general enough to include a diversity of models.

Here is the general idea: we can add a new base class to the library (in addition to the `Tokenizer` and `Model`) which we can call a `Retriever` for instance. The `Retriever` class:
- won't include trainable components and will thus be framework independent like the tokenizer,
- does both the retrieval of the documents given an encoding and the post processing (in our case the retokenization),
- it is sub-classed to be model specific (`RAGRetriever`) like the models and tokenizers,
- we can add an `AutoRetriever` version later.

We probably want to keep it's API fairly generic for now since this field is only beginning to be explored. The `Retriever` can just have a `__call__` method (like the tokenizers) which expect an encoding and does retrieval and postprocessing.

Then (copying @patrickvonplaten's idea) we would be to have a `RetrievalGenerationModel` that comprises the trainable (and PyTorch/TF specific) elements, i.e. the encoder and generator models. In its forward pass the `RetrievalGenerationModel` would have besides the usual `input_ids` one more input argument which is of class `Retriever`. The RetrivalGenerationModel would just call Retriever.forward(encoded_input) and expect an encoding that the self.generator could then be conditioned on.

What do you think? We would help you implement this option of course since it impacts more significantly the library.

I'm just wondering if TF will like a model which accept a class as input. What do you think @jplu @patrickvonplaten?
We could also have a method in the class to set the retriever instead of dynamically providing it. 

Never tried, but I doubt it should be possible in compiled mode. I think what you propose would be a better way to go.

> Never tried, but I doubt it should be possible in compiled mode. I think what you propose would be a better way to go.

@jplu - I think there was a slight misunderstanding. 
With the proposed approach we actually would pass an instantiation of a class as an argument to the forward pass of the `RetrievalGenerationModel` -> so before following this path we should check if this can nicely be done with TF...

Oh ok! I thought the question was about to pass one class that contains all the arguments. My bad 😢 

So, after reading your explanation I can say, yes it is doable!

Hey @ola13,

Thanks for your comment, this is indeed a very important aspect that I didn't really think of before. 
With @lhoestq, we have been brainstorming a bit and thought maybe a slighly different design could make sense:

```python
#!/usr/bin/env python3

class RetrievalGenerationModel(PretrainedModel):
    
    def __init__(self, config: RetrievalGenerationConfig, encoder: PretrainedModel, retrieval_generator: PretrainedModel):
        if encoder is not None and retrieval_generator is not None: 
            self.encoder  = encoder
            self.retrieval_generator = retrieval_generator
            self.config = RetrievalGenerationConfig.from_encoder_generator_config(self.encoder.config, self.retrieval_generator.config)

        assert config is not None
        super().__init__(config)

        if encoder is None:
            self.encoder = AutoModel.from_config(config.encoder)
        if  retrieval_generator is None:
            self.retrieval_generator.from_config(config.generator)

    @classmethod
    def from_pretrained_encoder_generator(cls, encoder_model_id, generator_model_id):
        encoder = AutoModel.from_pretrained(...) # load any query encoding model
        retrieval_generator = AutoRetrievalGeneratorModel.from_pretrained(...) # this would be a new class that contains any model that can be used as the `retrieval_generator` model.
        return cls(encoder=encoder, retrieval_generator=retrieval_generator)


    def forward(input_ids, retriever: PretrainedRetriever):
        # 1. Every retriever model encodes the query -> any `AutoModel` can be used here
        input_ids_encodings = self.encoder(input_ids) # model with weights

        # 2. Use costumized retriever (tokenizer-like) class instance, like `RAGRetriever` that 
        #      - query the index
        #      - reformats the document outputs
        #      - tokenizes the document outpus
        retrieved_docs_input_ids, retrieved_docs_encodings = retriever(input_ids_encodings, input_ids) # tokenizer like postprocessor that returns the tokenized docs input and the docs encodings

        # 3. Now the retrieval_generator requires a specific forward pass which accepts at least four kinds of tensors: 1) the input_ids (query), 2) the encoded_input_ids (encoded query), 3) retrieved_docs_input (tokenized context) and 4) retrieved_docs_encodings
        output_ids = self.retrieval_generator(input_ids, encoded_query, retrieved_docs_input_ids, retrieved_docs_encodings) # any `AutoRetrievalGeneratorModel` can be used here

class RagRetrievalGenerator(PretrainedModel):

    def __init__(self):
        self.generator = AutoModelForSeq2Seq.from_pretrained(...) # e.g. Bart

    def forward(input_ids, encodings, docs_input_ids, docs_encodings):
        doc_scores = torch.bmm(encodings.unsqueeze(1), docs_encodings.transpose(1, 2)).squeeze(1)
            ....
        output_ids = self.generator.generate(...)

class RAGRetriever(PretrainedRetriever)
    """
    This retriever is framework independant (for both TF and PT) 
    similar to a tokenizer
    """

     def __init__(self):
         self.docs = nlp.load_dataset(...)
         ...

     def __call__(input_ids_encodings, input_ids):
        # no tensor operations happen here
        ...

class DPRRetrivalGenerator(PretrainedModel):

    def __init__(self):
        self.genator = AutoModelForQuestionsAnswering.from_pretrained(...) # QA model 


    def forward(input_ids, encodings, docs_input_ids, docs_encodings):
           
        concated_qa_input = torch.cat([input_ids, docs_input_ids], dim=-1)
        output_ids = self.generator(concated_qa_input)


class DPRRetriever(PretrainedRetriever)
    """
    This retriever is framework independant (for both TF and PT) 
    similar to a tokenizer
    """

     def __init__(self):
         self.docs = nlp.load_dataset(...)
         ...

     def __call__(input_ids_encodings, input_ids):
        # no tensor operations happen here
        ...
```

Hopefully this is somewhat understandable @ola13 @thomwolf ...

@lhoestq and I think that for each RetrivalAugmentedModel we need 2 specific parts:

1) A specific Retriever: how are documents retrieved and formated and tokenized -> e.g. `RAGRetriever`
2) A specific Generator: Here we can also have multiple possibilities: DPR uses a `AutoModelForQuestionAnswering` while RAG uses a `AutoModelForSeq2Seq`

So with this framework we would have to introduce 1 general class that would be used for all RetrievalAugementedModels, called `RetrievalGenerationModel` (or whatever name fits better) and 2 architecture specific classes `RAGRetriever` and `RagRetrievalGenerator`.

Would be keen to hear your thoughts :-) 

Hey @patrickvonplaten, makes sense and in fact it's not very different from how we structured the code already the key differences that I see are:
- we move re-tokenization between query_encoder and generator to the Retriever (so respective tokenizers will be encapsulated by the Retriever not a model class as we currently do it)
- we move retrieval score calculation to the model so that no tensor operations happen in the retriever

which both should be pretty straightforward to implement.

The one thing that I'm still on the fence about is passing a `retriever` to each `forward` pass on a `RetrievalGenerationModel`, instead of making it a member of `RetrievalGenerationModel` class. Why do you feel the former is preferable over the latter?

Yeah, good point! It's a bit weird to pass a class instance just to make a forward pass with it. 

My main reason is the following: 

Currently, the library makes a very clear distinction between `config`, `tokenizer` and `model` which are all independent of each other. Each of these classes have a seperate `.from_pretrained()` and `.save_pretrained()` method where as the `PretrainedModel.save_pretrained(...)` and `PretrainedModel.from_pretrained(...)` internally call `PretrainedConfig.save_pretrained(...)` and `PretrainedConfig.save_pretrained(...)`, but **never** the `PretrainedTokenizer.from_pretrained(...)` an d`PretrainedTokenizer.save_pretrained(...)` methods. For a `RetrievalGenerationModel` I would like to reuse `PretrainedModel`'s `from_pretrained(...)` and `save_pretrained(...)` methods which means that a tokenizer instance should not be part of the model because other wise we would have to tweak this function (which I don't think is a good idea). 
Also, this will make the `RetrievalGenerationModel` a "clean" and relatively light `Model` object without any string processing logic in it whatsoever which is more in line with other `PretrainedModel` classes. 

@patrickvonplaten, got it, yeah makes sense! We would still want to call `PretrainedTokenizer.from_pretrained(...)` when initializing `RagRetriever` but I guess this should be fine?

Okay, so I would propose to do the following - I will refactor this PR to follow the design we discussed. It seems though that implementing the generic `Retriever` logic as discussed earlier by @thomwolf would require extra effort and time, and is not necessarily within the scope of this PR. In the interest of time, we could land this PR and then proceed with generalizing the retrieval logic? I'm then happy to work with the RAG implementation to make it compatible.

Exactly! I was thinking that we either create a genereric `PretrainedRetriever` class with a `from_pretrained()` method that calls the tokenizer `from_pretrained()` methods or add `from_pretrained()` method directly to `RagRetriever`. Maybe @lhoestq and @thomwolf have better insight on the "tokenizer" side. 

@ola13 maybe we can wait quickly if @lhoestq and @thomwolf are fine with the design as discussed above :-) 

Sounds awesome to me!

Hey I just refactored the model following suggestions above. One point is that I had to modify `generation_utils.py` to account for a model which takes a `retriever` as an argument to the encoder. Let me know what you think!

Hi, a question - to use RAG I need a couple of non-standard dependencies (faiss, psutil, nlp) - can I define a special test environment which would install those for rag tests? any pointers on how to handle this?

> Hey I just refactored the model following suggestions above. One point is that I had to modify `generation_utils.py` to account for a model which takes a `retriever` as an argument to the encoder. Let me know what you think!

Awesome ! I'll take a look. Also cc @patrickvonplaten 

> Hi, a question - to use RAG I need a couple of non-standard dependencies (faiss, psutil, nlp) - can I define a special test environment which would install those for rag tests? any pointers on how to handle this?

Maybe @LysandreJik knows more about how to handle tests with dependencies ?


Hey @ola13, 

I think the general code design is exactly what we have imagined to go for, defining a `RagRetriever` and passing the `retriever` to the forward pass, so this is great! 

Regarding the test dependencies, you can add the libraries here: https://github.com/huggingface/transformers/blob/d6c08b07a087e83915b4b3156bbf464cebc7b9b5/setup.py#L92 and it should automatically be installed for testing on circle ci :-) `psutil` is already in the test dependency

@ola13 - it would be awesome if you could add one "full" integration test with hardcoded input and output under @slow 

By that I mean, *e.g.* hardcoding an input question "Why does it rain", loading a relevant dataset using the `HfIndex` and the full pretrained encoder and generator model and hardcoding the expected output answer in thet test. I think all operations are deterministic (beam search, etc...), so no random seeds have to be set.

 This way we have one test where we can be sure that the model works as expected and every change to the model in the future can be checked against that.

The tests you have in `test_modeling_rag.py` so far look great. We could also add a full `RagModel` test by defining a dummy dataset that will be instantiated from a hardcoded dict at test time and instantiating a very light `RagRetriever` at test time this way. But we can manually add those tests later, they are not super important.

In terms of a timeline, it would be be awesome if you manage to make the `test_modeling_rag.py` tests pass and if you could add one "full" integration test showing reasonable results. After this is finished, I think the best idea is if we add some changes on top of your PR (this should take another 1,2 days) and then merge the model into the lib :-) 

Thanks a mille for your awesome work so far!!!

Hey @patrickvonplaten, sounds good! yes definitely adding an integration test was on my agenda, right now having merged the `master` I'm also dealing with some issues arising after the refactor from https://github.com/huggingface/transformers/commit/afc4ece462ad83a090af620ff4da099a0272e171#diff-72b038fcff0de4ae5e094e3cde9471f1 as we were relying on the old structure of `past`. I'm hoping to be done with both of these things by tomorrow :) 

Hi, I just added an integration test for RAG using the dummy variant of `wiki_dpr`. However, I had to locally hack `datasets` to make it run locally, as there seems to be a discrepancy between the dummy index name hardcoded in `wiki_dpr.py` here: https://github.com/huggingface/datasets/blob/37d4840a39eeff5d472beb890c8f850dc7723bb8/datasets/wiki_dpr/wiki_dpr.py#L72 (expecting  `dummy.psgs_w100.nq.IndexHNSWFlat-IP-train.faiss`) and what's available on HF's google cloud bucket:
```
~$ gsutil ls -r gs://huggingface-nlp/datasets/wiki_dpr/*
gs://huggingface-nlp/datasets/wiki_dpr/
gs://huggingface-nlp/datasets/wiki_dpr/dummy_psgs_w100_with_nq_embeddings_IndexFlatIP-train.faiss
gs://huggingface-nlp/datasets/wiki_dpr/psgs_w100.nq.IVFPQ4096_HNSW32_PQ64-IP-train.faiss
gs://huggingface-nlp/datasets/wiki_dpr/psgs_w100_with_nq_embeddings_IVFPQ4096_HNSW32,PQ64-IP-train.faiss
```

cc @lhoestq - this would have to be fixed quickly, alternatively I could use full `wiki_dpr` in tests, but that's 78GB, not sure if it makes sense.

Let me know what you think!

> cc @lhoestq - this would have to be fixed quickly, alternatively I could use full `wiki_dpr` in tests, but that's 78GB, not sure if it makes sense.

I fixed it, dummy.psgs_w100.nq.IndexHNSWFlat-IP-train.faiss is now available on gcs



Previous RAG code is now saved in this PR: #7200

Last fail is due to time-out. All import tests are passing => merging to master.