Add TFRag
Author: ratthachat
Pull request: yes

# What does this PR do?
This is a reopen PR of TFRag Draft version  (https://github.com/huggingface/transformers/pull/8892 )
which somehow seems broken and not accessible at the moment . 

## Things done
- `TFRagModel`
- `TFRagSequenceForGeneration`
- `TFRagTokenForGeneration`
- beam_search generation 
- "Work-around" example in graph mode training (The full graph-mode training need of `.numpy()` for retriever calling, and this doesn't work on graph mode) --> using `context_input_ids` in place of `input_ids`
- Complete test on `TFRag` 

## Things not yet done ... 
- Integrate with T5 as generator <-- In the next PR

## Who can review?
@jplu @patrickvonplaten 

Comments:
@jplu Thanks so much for your kind reviews! I will improve the code as you suggested.
@patrickvonplaten I have confirmed that my TF and Pytorch have equivalent `generate` output on `num_beam=1` (greedy) on all (15) test cases . 

Nevertheless, I just confirmed that the labels in the official test file is based on beam search where `num_beam=4`  and this TFRag does not have `beam_search` yet since I am not sure whether I should wait for `tf_generation` refactor . 

In the current Pytorch test, if explicitly set `num_beam=1` we will get exact same result as my TF implementation.

So to pass the current official test completely I will have to adapt `beam_search` into TFRag which I will try :)
For now, I slightly modify test cases to match Pytorch greedy output (we can revert back later when I finish `beam_search` )

This is very much WIP -> need some more time for the `save/from_pretrained()`

**UPDATED** Dec, 23 2020 : Finish `TFRagSequenceForGeneration`  

---
Sorry that I forgot to mention that the latest updated is still in very WIP (phase-2, after phase-1 which is core-part of `TFRagModel` and `TFRagToken` ), so not ready for review .

As discussed with Patrick, there is still an "output_truncation" issue (perhaps due to Caching & TFBart refactor) that Patrick will help take a look. 

**We will have to finish in Phase-2** (reason of TFRagToken tests fail at the moment )  :
[ x] - bugs in the new test `test_rag_token_inference_nq_checkpoint() `
[ x] - `beam_search` of `TFRagTokenForGeneration`  and 
[ x ] -  `TFRagSequenceForGeneration` <-- finished!
 [ ] - Apply jplu comments to clean up the code

which will take some more time :) :)  

Hi @jplu , @patrickvonplaten , regarding the graph issue, I just successfully made a[ colab notebook](https://colab.research.google.com/drive/1s-j9PB9yzrFsL6q5rZUQyf8_Lt6jDAkL?usp=sharing) capable of training `TFRagSequenceForGeneration` in graph mode  using `context_input_ids` as inputs (instead of `input_ids` ) ... 
Hopefully this is a reasonable work around on TFRag training.

Now `TFRagToken` has the same output_truncation [issue as Pytorch 's 9098](https://github.com/huggingface/transformers/pull/9098) . If Patrick help me solve this, I will be able to finish the beam search ... which should finish all the main parts of TFRag ... 

(After this, the remaining is to clean codes as suggested by you guys, complete all fast tests, and fix the ugly hacks [ ie. pass `test_rag_token_inference_nq_checkpoint` without hacking ] )

Hey @ratthachat,

great work! I'll take care of the new (and hopefully last :D) `save_/from_pretrained()` issue today. And I'll also make sure that `TFRagTokenGeneration` works properly for `greedy_search`! I'll leave `beam_search` then for you :-) 

Okey, I fixed the `from_pretrained` bug when `from_pt=True` and also fixed `greedy search` for `TFRagToken`. Thanks for the very in-detail error descriptions @ratthachat. 

I think the only thing left to do for now is to correctly implement `beam_search` for `TFRagToken`. As I understood it, you'd like to give it a try. Let me know if you want me to tackle this or if you need help :-)

I extended your "load from pt" test slightly to make sure that weights are now always (I really hope so :D) correctly loaded. Also, I added one test for `greedy_search` for `RagToken`

Hi @patrickvonplaten , I agree and all T5 related tests are deleted.

Tests related to TFBart are still not passed which likely due to TFBart bug.
More precisely, TFBart vs. Torch's Bart forward pass return `generator_enc_last_hidden_state.shape`
differently if we provide `decoder_input_ids`

```
texts = "My friends are cool but they eat too many carbs. I really want them to be healthy, so I buy them vegetable."
texts2 = "My friends are cool."
inputs = tokenizer([texts], max_length=1024, return_tensors='tf')
inputs2 = tokenizer([texts2], max_length=1024, return_tensors='tf')

input_ids=inputs['input_ids']
input_ids2=inputs2['input_ids']
out = model(input_ids,decoder_input_ids=None)
print(out.encoder_last_hidden_state.shape) # RETURN (1, 27, 1024)

out = model(input_ids,decoder_input_ids=input_ids2)
print(out.encoder_last_hidden_state.shape) # RETURN (1, 7, 1024)

```

If we run the same snippet in Pytorch, they will both return `(1, 27, 1024)` . (tested on both official released and master)
So likely this is TFBart bug, **and it makes 4 TFRag fast-tests fail**. 
(shape of `out.encoder_last_hidden_state.shape` is not as expected)

> Hi @patrickvonplaten , I agree and all T5 related tests are deleted.
> 
> Tests related to TFBart are still not passed which likely due to TFBart bug.
> More precisely, TFBart vs. Torch's Bart forward pass return `generator_enc_last_hidden_state.shape`
> differently if we provide `decoder_input_ids`
> 
> ```
> texts = "My friends are cool but they eat too many carbs. I really want them to be healthy, so I buy them vegetable."
> texts2 = "My friends are cool."
> inputs = tokenizer([texts], max_length=1024, return_tensors='tf')
> inputs2 = tokenizer([texts2], max_length=1024, return_tensors='tf')
> 
> input_ids=inputs['input_ids']
> input_ids2=inputs2['input_ids']
> out = model(input_ids,decoder_input_ids=None)
> print(out.encoder_last_hidden_state.shape) # RETURN (1, 27, 1024)
> 
> out = model(input_ids,decoder_input_ids=input_ids2)
> print(out.encoder_last_hidden_state.shape) # RETURN (1, 7, 1024)
> ```
> 
> If we run the same snippet in Pytorch, they will both return `(1, 27, 1024)` . (tested on both official released and master)
> So likely this is TFBart bug, **and it makes 4 TFRag fast-tests fail**.
> (shape of `out.encoder_last_hidden_state.shape` is not as expected)

Hey @ratthachat, 

> Hi @patrickvonplaten , I agree and all T5 related tests are deleted.
> 
> Tests related to TFBart are still not passed which likely due to TFBart bug.
> More precisely, TFBart vs. Torch's Bart forward pass return `generator_enc_last_hidden_state.shape`
> differently if we provide `decoder_input_ids`
> 
> ```
> texts = "My friends are cool but they eat too many carbs. I really want them to be healthy, so I buy them vegetable."
> texts2 = "My friends are cool."
> inputs = tokenizer([texts], max_length=1024, return_tensors='tf')
> inputs2 = tokenizer([texts2], max_length=1024, return_tensors='tf')
> 
> input_ids=inputs['input_ids']
> input_ids2=inputs2['input_ids']
> out = model(input_ids,decoder_input_ids=None)
> print(out.encoder_last_hidden_state.shape) # RETURN (1, 27, 1024)
> 
> out = model(input_ids,decoder_input_ids=input_ids2)
> print(out.encoder_last_hidden_state.shape) # RETURN (1, 7, 1024)
> ```
> 
> If we run the same snippet in Pytorch, they will both return `(1, 27, 1024)` . (tested on both official released and master)
> So likely this is TFBart bug, **and it makes 4 TFRag fast-tests fail**.
> (shape of `out.encoder_last_hidden_state.shape` is not as expected)

Hey @ratthachat,

sorry to answer so late! I don't think that there is a TFBart bug to be honest. When inputting `decoder_input_ids` the behavior you described above is expected and is the same for PyTorch. 

If you copy & paste the following code on master:

```python
#!/usr/bin/env python3
from transformers import AutoModelForSeq2SeqLM, TFAutoModelForSeq2SeqLM
from transformers import AutoTokenizer

model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-base")
tf_model = TFAutoModelForSeq2SeqLM.from_pretrained("facebook/bart-base")

tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")

input_str_short = "this is a string"
input_str_long = "this is a veeeeeeery veeeeeeeeery long string!!!"

output_shape = model(input_ids=tokenizer(input_str_short, return_tensors="pt").input_ids, decoder_input_ids=tokenizer(input_str_long, return_tensors="pt").input_ids)[0].shape

tf_output_shape = tf_model(input_ids=tokenizer(input_str_short, return_tensors="tf").input_ids, decoder_input_ids=tokenizer(input_str_long, return_tensors="tf").input_ids)[0].shape

assert output_shape == tf_output_shape, "Output shapes have to be the same"
```

you'll see that no assertion error is thrown. 

I think you might have to adapt a couple of TFRag tests to make it work. There also might be a small chance that you have to rebase your current PR to master because there is a weird version of Bart in this PR (but I doubt that a bit to be honest). 

Please let me know if you need help for the tests! Think we are almost finished !!! :-) 

Hi @patrickvonplaten , sorry that I did not explain clear enough.
The test I exactly adapted from Pytorch test the shape of `last_hidden_states` not the shape of `logits` . 

Ie. please see
https://github.com/ratthachat/transformers/blob/tfrag-draft-new/tests/test_modeling_tf_rag.py#L375

```
self.assertEqual(
                outputs.generator_enc_last_hidden_state.shape,
                (n_docs * decoder_input_ids.shape[0], self.max_combined_length, config.generator.hidden_size),
            )
```

From your example change output from `[0] (logits)` to `[2] (last_hidden_states)` , **we indeed got assertion error**. 

```
from transformers import AutoModelForSeq2SeqLM, TFAutoModelForSeq2SeqLM
from transformers import AutoTokenizer

model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-base")
tf_model = TFAutoModelForSeq2SeqLM.from_pretrained("facebook/bart-base")

tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")

input_str_short = "this is a string"
input_str_long = "this is a veeeeeeery veeeeeeeeery long string!!!"

output_shape = model(input_ids=tokenizer(input_str_short, return_tensors="pt").input_ids, decoder_input_ids=tokenizer(input_str_long, return_tensors="pt").input_ids)[2].shape

tf_output_shape = tf_model(input_ids=tokenizer(input_str_short, return_tensors="tf").input_ids, decoder_input_ids=tokenizer(input_str_long, return_tensors="tf").input_ids)[2].shape

assert output_shape == tf_output_shape, "Output shapes have to be the same"
```

```
AssertionError                            Traceback (most recent call last)
<ipython-input-17-04a80fc987f7> in <module>()
     14 tf_output_shape = tf_model(input_ids=tokenizer(input_str_short, return_tensors="tf").input_ids, decoder_input_ids=tokenizer(input_str_long, return_tensors="tf").input_ids)[2].shape
     15 
---> 16 assert output_shape == tf_output_shape, "Output shapes have to be the same"

AssertionError: Output shapes have to be the same
```

BTW, about the rebase, I really want to do it, but I could not solve the merging conflicts .

> Hi @patrickvonplaten , sorry that I did not explain clear enough.
> The test I exactly adapted from Pytorch test the shape of `last_hidden_states` not the shape of `logits` .
> 
> Ie. please see
> https://github.com/ratthachat/transformers/blob/tfrag-draft-new/tests/test_modeling_tf_rag.py#L375
> 
> ```
> self.assertEqual(
>                 outputs.generator_enc_last_hidden_state.shape,
>                 (n_docs * decoder_input_ids.shape[0], self.max_combined_length, config.generator.hidden_size),
>             )
> ```
> 
> From your example change output from `[0] (logits)` to `[2] (last_hidden_states)` , **we indeed got assertion error**.
> 
> ```
> from transformers import AutoModelForSeq2SeqLM, TFAutoModelForSeq2SeqLM
> from transformers import AutoTokenizer
> 
> model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-base")
> tf_model = TFAutoModelForSeq2SeqLM.from_pretrained("facebook/bart-base")
> 
> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
> 
> input_str_short = "this is a string"
> input_str_long = "this is a veeeeeeery veeeeeeeeery long string!!!"
> 
> output_shape = model(input_ids=tokenizer(input_str_short, return_tensors="pt").input_ids, decoder_input_ids=tokenizer(input_str_long, return_tensors="pt").input_ids)[2].shape
> 
> tf_output_shape = tf_model(input_ids=tokenizer(input_str_short, return_tensors="tf").input_ids, decoder_input_ids=tokenizer(input_str_long, return_tensors="tf").input_ids)[2].shape
> 
> assert output_shape == tf_output_shape, "Output shapes have to be the same"
> ```
> 
> ```
> AssertionError                            Traceback (most recent call last)
> <ipython-input-17-04a80fc987f7> in <module>()
>      14 tf_output_shape = tf_model(input_ids=tokenizer(input_str_short, return_tensors="tf").input_ids, decoder_input_ids=tokenizer(input_str_long, return_tensors="tf").input_ids)[2].shape
>      15 
> ---> 16 assert output_shape == tf_output_shape, "Output shapes have to be the same"
> 
> AssertionError: Output shapes have to be the same
> ```
> 
> BTW, about the rebase, I really want to do it, but I could not solve the merging conflicts .

Hey @ratthachat, please note that [2] are the `hidden_states` and not the `last_hidden_state`. The last hidden_state is as expected the same.

```python
from transformers import AutoModel, TFAutoModel
from transformers import AutoTokenizer

model = AutoModel.from_pretrained("facebook/bart-base")
tf_model = TFAutoModel.from_pretrained("facebook/bart-base")

tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")

input_str_short = "this is a string"
input_str_long = "this is a veeeeeeery veeeeeeeeery long string!!!"

output_shape = model(input_ids=tokenizer(input_str_short, return_tensors="pt").input_ids, decoder_input_ids=tokenizer(input_str_long, return_tensors="pt").input_ids).last_hidden_state.shape

tf_output_shape = tf_model(input_ids=tokenizer(input_str_short, return_tensors="tf").input_ids, decoder_input_ids=tokenizer(input_str_long, return_tensors="tf").input_ids).last_hidden_state.shape

assert output_shape == tf_output_shape, "Output shapes have to be the same"
```

I'm not 100% sure if the merge was completely correct, but let's first focus on making all TFRag tests pass. The other test we can fix later :-) 

I apologize @patrickvonplaten . I think I am now a bit confused.
In TFRag, `model.generator` is created using `TFAutoModelForSeq2SeqLM` (the same as Pytorch).
https://github.com/ratthachat/transformers/blob/tfrag-draft-new/src/transformers/models/rag/modeling_tf_rag.py#L519
https://github.com/ratthachat/transformers/blob/tfrag-draft-new/src/transformers/models/rag/modeling_rag.py#L356

In the above example, you used  `TFAutoModel`, so it's not the same.
So if I changed to `TFAutoModelForSeq2SeqLM` and check `encoder_last_hidden_state.shape` , 
**[all fast tests, test this `encoder_last_hidden_state.shape`  attribute]**
I still got assertion error. 

I am not sure what's going on and I may miss something simple here. I apologize again.

```
from transformers import AutoModelForSeq2SeqLM, TFAutoModelForSeq2SeqLM
from transformers import AutoTokenizer

model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-base")
tf_model = TFAutoModelForSeq2SeqLM.from_pretrained("facebook/bart-base")

tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")

input_str_short = "this is a string"
input_str_long = "this is a veeeeeeery veeeeeeeeery long string!!!"

output_shape = model(input_ids=tokenizer(input_str_short, return_tensors="pt").input_ids, decoder_input_ids=tokenizer(input_str_long, return_tensors="pt").input_ids).encoder_last_hidden_state.shape
tf_output_shape = tf_model(input_ids=tokenizer(input_str_short, return_tensors="tf").input_ids, decoder_input_ids=tokenizer(input_str_long, return_tensors="tf").input_ids).encoder_last_hidden_state.shape

print(output_shape, tf_output_shape)

assert output_shape == tf_output_shape, "Output shapes have to be the same"
```

```
torch.Size([1, 6, 768]) (1, 17, 768)
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-4-9790bfa59614> in <module>()
     16 # print(output.keys(), tf_output.keys())
     17 
---> 18 assert output_shape == tf_output_shape, "Output shapes have to be the same"

AssertionError: Output shapes have to be the same
```

Patrick, another issue is that after rebase, there is an error on `load_weight_prefix` which we invented for TFRag's name, so now the basic building block does not work.

`TypeError: ('Keyword argument not understood:', 'load_weight_prefix') `


> Patrick, another issue is that after rebase, there is an error on `load_weight_prefix` which we invented for TFRag's name, so now the basic building block does not work.
> 
> `TypeError: ('Keyword argument not understood:', 'load_weight_prefix') `

Yeah sorry, I made a quick & dirty rebase so there might be errors! It would be awesome if you could fix them (if there are easy to fix)

> from transformers import AutoModelForSeq2SeqLM, TFAutoModelForSeq2SeqLM
> from transformers import AutoTokenizer
> 
> model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-base")
> tf_model = TFAutoModelForSeq2SeqLM.from_pretrained("facebook/bart-base")
> 
> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
> 
> input_str_short = "this is a string"
> input_str_long = "this is a veeeeeeery veeeeeeeeery long string!!!"
> 
> output_shape = model(input_ids=tokenizer(input_str_short, return_tensors="pt").input_ids, decoder_input_ids=tokenizer(input_str_long, return_tensors="pt").input_ids).encoder_last_hidden_state.shape
> tf_output_shape = tf_model(input_ids=tokenizer(input_str_short, return_tensors="tf").input_ids, decoder_input_ids=tokenizer(input_str_long, return_tensors="tf").input_ids).encoder_last_hidden_state.shape
> 
> print(output_shape, tf_output_shape)
> 
> assert output_shape == tf_output_shape, "Output shapes have to be the same"

I see thanks for the very descriptive error description! You're completely right -> that's a bug, great catch! I'll fix

> > from transformers import AutoModelForSeq2SeqLM, TFAutoModelForSeq2SeqLM
> > from transformers import AutoTokenizer
> > model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-base")
> > tf_model = TFAutoModelForSeq2SeqLM.from_pretrained("facebook/bart-base")
> > tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
> > input_str_short = "this is a string"
> > input_str_long = "this is a veeeeeeery veeeeeeeeery long string!!!"
> > output_shape = model(input_ids=tokenizer(input_str_short, return_tensors="pt").input_ids, decoder_input_ids=tokenizer(input_str_long, return_tensors="pt").input_ids).encoder_last_hidden_state.shape
> > tf_output_shape = tf_model(input_ids=tokenizer(input_str_short, return_tensors="tf").input_ids, decoder_input_ids=tokenizer(input_str_long, return_tensors="tf").input_ids).encoder_last_hidden_state.shape
> > print(output_shape, tf_output_shape)
> > assert output_shape == tf_output_shape, "Output shapes have to be the same"
> 
> I see thanks for the very descriptive error description! You're completely right -> that's a bug, great catch! I'll fix

Ok merged it: https://github.com/huggingface/transformers/pull/9944. Could you merge master once again into your PR and see whether the tests work now? :-) 

```bash
git fetch upstream master
git merge upstream/master
```

I don' think there will be any merge conflicts. Lemme know if you need any help :-) 

Thanks so much Patrick. Tomorrow, I will try my best to fix the "load_weights_prefix" issue and will come back ❤️ 

@patrickvonplaten I am sorry - bad news. 
Even though now I think all tests should be passed, I could not find a way to fix the `load_weights_prefix` issue above arised after conflict fixing 2 days ago.

It seems that this `load_weights_prefix` is sent as `kwarg` to all required functions correctly but failed with Keras not allowing this argument .
(I attach the full `TraceBack` below).

At first, I thought that this might be due to the recent TF2.4 upgrade, but I tried downgrade back to TF2.3 and still got the same error. Could you please help take a look?

To reproduce, simply initiate the model:
```
from transformers import RagTokenizer, RagRetriever 
from transformers.models.rag.modeling_tf_rag import TFRagModel, TFRagSequenceForGeneration, TFRagTokenForGeneration

PATH = "facebook/rag-token-nq"
tokenizer = RagTokenizer.from_pretrained(PATH)
retriever = RagRetriever.from_pretrained(PATH, index_name="exact", use_dummy_dataset=True) 

model = TFRagModel.from_pretrained_question_encoder_generator('facebook/dpr-question_encoder-single-nq-base', "facebook/bart-base", generator_from_pt=True, question_encoder_from_pt=True, retriever=retriever) 
```

Produced the following TraceBack:

```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-10-a201b416721e> in <module>()
      1 
----> 2 model = TFRagModel.from_pretrained_question_encoder_generator('facebook/dpr-question_encoder-single-nq-base', "facebook/bart-base", generator_from_pt=True, question_encoder_from_pt=True, retriever=retriever)

8 frames
/usr/local/lib/python3.6/dist-packages/transformers/models/rag/modeling_tf_rag.py in from_pretrained_question_encoder_generator(cls, question_encoder_pretrained_model_name_or_path, generator_pretrained_model_name_or_path, retriever, *model_args, **kwargs)
    366                 name="generator",
    367                 load_weight_prefix=cls.load_weight_prefix,
--> 368                 **kwargs_generator,
    369             )
    370 

/usr/local/lib/python3.6/dist-packages/transformers/models/auto/modeling_tf_auto.py in from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs)
   1105         if type(config) in TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.keys():
   1106             return TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING[type(config)].from_pretrained(
-> 1107                 pretrained_model_name_or_path, *model_args, config=config, **kwargs
   1108             )
   1109         raise ValueError(

/usr/local/lib/python3.6/dist-packages/transformers/modeling_tf_utils.py in from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs)
   1245 
   1246         # Instantiate model.
-> 1247         model = cls(config, *model_args, **model_kwargs)
   1248 
   1249         if from_pt:

/usr/local/lib/python3.6/dist-packages/transformers/models/bart/modeling_tf_bart.py in __init__(self, config, load_weight_prefix, *inputs, **kwargs)
   1246     def __init__(self, config, load_weight_prefix=None, *inputs, **kwargs):
   1247         super().__init__(config, *inputs, **kwargs)
-> 1248         self.model = TFBartModel(config, load_weight_prefix=load_weight_prefix, name="model")
   1249         self.use_cache = config.use_cache
   1250         # final_bias_logits is registered as a buffer in pytorch, so not trainable for the the sake of consistency.

/usr/local/lib/python3.6/dist-packages/transformers/models/bart/modeling_tf_bart.py in __init__(self, config, *inputs, **kwargs)
   1139 class TFBartModel(TFBartPretrainedModel):
   1140     def __init__(self, config: BartConfig, *inputs, **kwargs):
-> 1141         super().__init__(config, *inputs, **kwargs)
   1142 
   1143         self.model = TFBartMainLayer(config, name="model")

/usr/local/lib/python3.6/dist-packages/transformers/modeling_tf_utils.py in __init__(self, config, *inputs, **kwargs)
    629 
    630     def __init__(self, config, *inputs, **kwargs):
--> 631         super().__init__(*inputs, **kwargs)
    632         if not isinstance(config, PretrainedConfig):
    633             raise ValueError(

/usr/local/lib/python3.6/dist-packages/tensorflow/python/training/tracking/base.py in _method_wrapper(self, *args, **kwargs)
    455     self._self_setattr_tracking = False  # pylint: disable=protected-access
    456     try:
--> 457       result = method(self, *args, **kwargs)
    458     finally:
    459       self._self_setattr_tracking = previous_value  # pylint: disable=protected-access

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py in __init__(self, *args, **kwargs)
    260     # self.non_trainable_weights
    261     generic_utils.validate_kwargs(kwargs, {'trainable', 'dtype', 'dynamic',
--> 262                                            'name', 'autocast'})
    263     super(Model, self).__init__(**kwargs)
    264     # By default, Model is a subclass model, which is not in graph network.

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/utils/generic_utils.py in validate_kwargs(kwargs, allowed_kwargs, error_message)
    776   for kwarg in kwargs:
    777     if kwarg not in allowed_kwargs:
--> 778       raise TypeError(error_message, kwarg)
    779 
    780 

TypeError: ('Keyword argument not understood:', 'load_weight_prefix')
```

> @patrickvonplaten I am sorry - bad news.
> Even though now I think all tests should be passed, I could not find a way to fix the `load_weights_prefix` issue above arised after conflict fixing 2 days ago.
> 
> It seems that this `load_weights_prefix` is sent as `kwarg` to all required functions correctly but failed with Keras not allowing this argument .
> (I attach the full `TraceBack` below).
> 
> At first, I thought that this might be due to the recent TF2.4 upgrade, but I tried downgrade back to TF2.3 and still got the same error. Could you please help take a look?
> 
> To reproduce, simply initiate the model:
> 
> ```
> from transformers import RagTokenizer, RagRetriever 
> from transformers.models.rag.modeling_tf_rag import TFRagModel, TFRagSequenceForGeneration, TFRagTokenForGeneration
> 
> PATH = "facebook/rag-token-nq"
> tokenizer = RagTokenizer.from_pretrained(PATH)
> retriever = RagRetriever.from_pretrained(PATH, index_name="exact", use_dummy_dataset=True) 
> 
> model = TFRagModel.from_pretrained_question_encoder_generator('facebook/dpr-question_encoder-single-nq-base', "facebook/bart-base", generator_from_pt=True, question_encoder_from_pt=True, retriever=retriever) 
> ```
> 
> Produced the following TraceBack:
> 
> ```
> ---------------------------------------------------------------------------
> TypeError                                 Traceback (most recent call last)
> <ipython-input-10-a201b416721e> in <module>()
>       1 
> ----> 2 model = TFRagModel.from_pretrained_question_encoder_generator('facebook/dpr-question_encoder-single-nq-base', "facebook/bart-base", generator_from_pt=True, question_encoder_from_pt=True, retriever=retriever)
> 
> 8 frames
> /usr/local/lib/python3.6/dist-packages/transformers/models/rag/modeling_tf_rag.py in from_pretrained_question_encoder_generator(cls, question_encoder_pretrained_model_name_or_path, generator_pretrained_model_name_or_path, retriever, *model_args, **kwargs)
>     366                 name="generator",
>     367                 load_weight_prefix=cls.load_weight_prefix,
> --> 368                 **kwargs_generator,
>     369             )
>     370 
> 
> /usr/local/lib/python3.6/dist-packages/transformers/models/auto/modeling_tf_auto.py in from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs)
>    1105         if type(config) in TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.keys():
>    1106             return TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING[type(config)].from_pretrained(
> -> 1107                 pretrained_model_name_or_path, *model_args, config=config, **kwargs
>    1108             )
>    1109         raise ValueError(
> 
> /usr/local/lib/python3.6/dist-packages/transformers/modeling_tf_utils.py in from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs)
>    1245 
>    1246         # Instantiate model.
> -> 1247         model = cls(config, *model_args, **model_kwargs)
>    1248 
>    1249         if from_pt:
> 
> /usr/local/lib/python3.6/dist-packages/transformers/models/bart/modeling_tf_bart.py in __init__(self, config, load_weight_prefix, *inputs, **kwargs)
>    1246     def __init__(self, config, load_weight_prefix=None, *inputs, **kwargs):
>    1247         super().__init__(config, *inputs, **kwargs)
> -> 1248         self.model = TFBartModel(config, load_weight_prefix=load_weight_prefix, name="model")
>    1249         self.use_cache = config.use_cache
>    1250         # final_bias_logits is registered as a buffer in pytorch, so not trainable for the the sake of consistency.
> 
> /usr/local/lib/python3.6/dist-packages/transformers/models/bart/modeling_tf_bart.py in __init__(self, config, *inputs, **kwargs)
>    1139 class TFBartModel(TFBartPretrainedModel):
>    1140     def __init__(self, config: BartConfig, *inputs, **kwargs):
> -> 1141         super().__init__(config, *inputs, **kwargs)
>    1142 
>    1143         self.model = TFBartMainLayer(config, name="model")
> 
> /usr/local/lib/python3.6/dist-packages/transformers/modeling_tf_utils.py in __init__(self, config, *inputs, **kwargs)
>     629 
>     630     def __init__(self, config, *inputs, **kwargs):
> --> 631         super().__init__(*inputs, **kwargs)
>     632         if not isinstance(config, PretrainedConfig):
>     633             raise ValueError(
> 
> /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/tracking/base.py in _method_wrapper(self, *args, **kwargs)
>     455     self._self_setattr_tracking = False  # pylint: disable=protected-access
>     456     try:
> --> 457       result = method(self, *args, **kwargs)
>     458     finally:
>     459       self._self_setattr_tracking = previous_value  # pylint: disable=protected-access
> 
> /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py in __init__(self, *args, **kwargs)
>     260     # self.non_trainable_weights
>     261     generic_utils.validate_kwargs(kwargs, {'trainable', 'dtype', 'dynamic',
> --> 262                                            'name', 'autocast'})
>     263     super(Model, self).__init__(**kwargs)
>     264     # By default, Model is a subclass model, which is not in graph network.
> 
> /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/utils/generic_utils.py in validate_kwargs(kwargs, allowed_kwargs, error_message)
>     776   for kwarg in kwargs:
>     777     if kwarg not in allowed_kwargs:
> --> 778       raise TypeError(error_message, kwarg)
>     779 
>     780 
> 
> TypeError: ('Keyword argument not understood:', 'load_weight_prefix')
> ```

Should be fine now - can you check? 

@patrickvonplaten all fast and slow tests are now pass 😄  
Just that in one slow test : `test_rag_token_generate_batch` , my colab P100 ran out of memory if we used all 15 inputs.
If I reduce the size of inputs  to 5, the test passed. 
(any 5 of 15 will passed indicated that all outputs are correct)

@ratthachat - you've really done an amazing job here! The PR looks very nice to me overall. 
One thing, I'd like to change before trying to merge is to delete the `_generate_beam_search` and `_generate_no_beam_search` methods and use the default ones instead. I can definitely help you get this done here. Do you know what the differences are in `_generate_beam_search` that you added to `modeling_tf_rag.py` compared to the one in `modeling_tf_utils.py`? Happy to help you here

Apart from that, I only left a couple of nits.

Hi Patrick, thanks for all your super kind helps all these time! 😄 ❤️ 
I improved docstrings as suggested.

About` _generate_beam_search` and `_generate_no_beam_search`, actually there are exactly **20 lines differences**.
I made a notebook to show **20 lines differences**  --  Version1 (`TFRag`) and Version2 (`generation_tf_utils.py`) 
https://www.kaggle.com/ratthachat/generate-beam-and-no-searches/
 (Please clicking version on the top-right and see `diff` )

Mainly, I simply fix both functions to accept `**kwarg` arguments from `TFRag` (in particular `kwargs["encoder_outputs"]` ).
However, I did not directly change/PR `generation_tf_utils.py` for two reasons : 

(1) I am not sure if this changes will affect other TF models or not and I don't have enough resources to check them all
(2) As we once discussed, there will be big 'generation refactor' in `generation_tf_utils.py` like Pytorch soon, and it should be a great chance there to fix & test TFRag at that time together with other TF models.

What do you think ?



All the slow tests passed for TFBart & TFRag -> PR is ready for review. Started running the whole suite of SLOW tests, just to be sure - will report any unexpected behavior here.

@LysandreJik @sgugger @jplu  - it would be great if you can review the PR.

> I am very much not in favor of changing the core modeling utils to fit the needs of one model, so really dislike the change related to `from_pretrained`. I understand the problems of scope, but it seems to me that it's only there to be able to write the convenience init `from_encoder_decoder_pretrained` and the like which is not strictly necessary (one can load the encoder and decoder out of the model then instantiate it by passing the encoder and the decoder, it's just three lines of code instead of one).
> 
> I would favor this option for now while we look into different solutions for the weight loading.

I also think that it's not clean at all to introduce this hack to be able to load the pretrained weights correctly, but there is no way around it really. 

It is not possible to load the encoder and decoder separately and to pass them into init if one wants to save the model afterward. Even with the current hack this is not possible because there is no way of knowing what the correct scope of the overall model is when the submodel is loaded separately. Try this in the branch of this PR:

```python
from transformers import TFBartForConditionalGeneration, TFDPRQuestionEncoder, TFRagModel, RagRetriever

retriever = RagRetriever.from_pretrained("facebook/rag-sequence-nq", index_name="exact", use_dummy_dataset=True)

encoder = TFDPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
generator = TFBartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")

rag = TFRagModel(question_encoder=encoder, generator=generator, retriever=retriever)

rag.save_pretrained("rag_temp")

new_rag = TFRagModel.from_pretrained("rag_temp", retriever=retriever)  # ERROR => the weights are randomly initialized here
```

This means that ```from_encoder_decoder_pretrained``` is not just a convenience function, but actually the only way to correctly load two "sub" models into a composite model class. Otherwise the weights are not saved correctly and can then not be loaded again.

Also one other thing I want to point out here is that the change/hack added to `TFPretrainedModel`'s `from_pretrained` method in `modeling_tf_utils.py` is not just done for RAG but will make composite TF model classes such as `TFEncoderDecoderModel` possible.
At the same time, I also think it's a big hack that is not beautiful in any way. 

At the moment, I don't really see another solution here because we cannot just overwrite the `from_pretrained(...)` method for TFRag since the `load_preifx_weight` is passed to BART's and DPR's `from_pretrained`'s method. What we could do instead however is to add a whole new `from_pretrained_with_prefix(....)` function to `modeling_tf_utils.py` instead of changing the existing method. 

What do you think? @sgugger @LysandreJik & @jplu 

> At the moment, I don't really see another solution here because we cannot just overwrite the from_pretrained(...) method for TFRag since the load_preifx_weight is passed to BART's and DPR's from_pretrained's method. What we could do instead however is to add a whole new from_pretrained_with_prefix(....) function to modeling_tf_utils.py instead of changing the existing method.

From a short term view, I don't see another solution either. We clearly have an issue with the weights naming for Seq2Seq models and the found workaround we had until now (the `tf.compat.v1...`) reaches its limits with RAG as it requires changes where it should not.

For me this has to be rethought because I think that all the convert issues have to be handled in the `modeling_tf_pytorch_utils.py` script and not elsewhere, and we should stop to force TF to have the same names than PT but more handle how to convert "proper" TF weight names to "proper" PT weight names (and the other way around). I think that if we continue without having a better design for this we go to a more complex implementation and then understanding. We should also keep in mind that the day TF removes the V1 compatibility, none of these models will work as expected. Hence, I would accept this as a temporary solution, but we clearly need to review this TF naming part thoroughly.

As discussed a bit offline, @LysandreJik will take a final review and then we'll merge the approach of this PR.

We should integrate the suggestion from @sgugger which is to change `prefix` to a private function arg `_prefix` to keep the design flexible for future changes.

Once @LysandreJik has done the review, I'll do a last refactor & then we can merge I think @ratthachat :-)

Thanks so much everyone for the merge! Especially @jplu who gave insightful comments on several earlier versions and @patrickvonplaten who has collaborated and has greatly helped in all aspects of this works!! 
(Honestly it's not possible without Patrick's help)

About the single mismatched generated answer ("step by step" vs. "evolution"), I will investigate this point further. Strangely, in earlier versions all tests are passed meaning all outputs are equivalent.