TFRag draft #1 (page BROKEN) - Should close and use #9002 instead
Author: ratthachat
Pull request: yes

# What does this PR do?
Hi guys, this is the draft WIP of TFRag. It is runnable in eager mode with mostly proper outputs.
I really need help/consult at this stage especially from @Patrick and @jplu .
In this draft, only `modeling_tf_rag.py` is new here.

## What is done and test. (working on HF 4.0.0, but not on Master due to changes in TF input)
- TFRagModel
- TFRagTokenForGeneration
- generate function with no_beam_search
- work on eager mode 
- Colab notebook to test/play around whether the code works properly : https://colab.research.google.com/drive/1CfCulkKGrneiQ0gV0Bgdo71gZ_kgMRIB?usp=sharing

## Main things not done yet
- TFRagSequenceForGeneration
- beam_search generation (may wait for TF generation refactor ?)
- Working in graph mode (due to the need of `.numpy()` for retriever calling, and this doesn't work on graph mode)
- Change input format for HF 4.1.0 (need help from @jplu)

## Need your suggestion on NEED_ADVICE, NEED_HELP
As stated, the code is mostly OK except on the points I marked TOFIX which will be clean later during finishing the draft.
However, there are 2 categories that I really need help especially from @Patrick:

1) There are some codes that work, but I am not sure if I meet Huggingface coding standard (marked by NEED_ADVICE)

2) There are 2 points that I need real help (marked by  NEED_HELP ) 
2.1) the aforementioned `.numpy()` on graph mode. 
2.2) about `.from_pretrained` : Rag has two loading methods which are `.from_pretrained_question_encoder_generator` and `.from_pretrained` . While` .from_pretrained_question_encoder_generator` works , in `.from_pretrained` there are two weights that the names do not match , which I could not find any way to fix :

```
'rag.generator.model.shared.weight', 'rag.generator.final_logits_bias' --> Pytorch name
'model.shared.weight', 'final_logits_bias' --> TF name
```
So at the moment I made UGLY fix by overwrite .from_pretrained and manually loading these two weights.


## Who can review?
TFRag : @patrickvonplaten , 
new TF input for master / 4.1.0 : @jplu 
about graph mode & retriever module : will need help from @lhoestq later once all other issues are fixed :)

Comments:
Hi guys, most commits were from the previous PR (TFDPR). I do not know how to remove them, sorry!
Only `modeling_tf_rag.py` is new here.

Awesome work @ratthachat!!!

For the input, I think the best way is to check how it is done in TF BERT for example and if you have difficulties to understand you can ask your questions here :)

About the `.numpy()` as suggested you can use the `tf.make_ndarra()` like this `tf.make_ndarray(tf.make_tensor_proto(my_tensor))`.

About the weights, I suggest you to check how it is done in TF BART there are similar weights naming.

> Awesome work @ratthachat!!!
> 
> For the input, I think the best way is to check how it is done in TF BERT for example and if you have difficulties to understand you can ask your questions here :)
> 
> About the `.numpy()` as suggested you can use the `tf.make_ndarra()` like this `tf.make_ndarray(tf.make_tensor_proto(my_tensor))`.
> 
> About the weights, I suggest you to check how it is done in TF BART there are similar weights naming.

Hi Julien @jplu , thanks for the reply!
On these points, I think I may miss something simple, but I could not solve the puzzle by myself at this moment.
1) on `tf.make_ndarray(tf.make_tensor_proto(my_tensor))` , I could make it work **only** on eager mode too, so still have the problem in graph mode (I could not make it work inside @tf.function)

2) about the input, yes I tried to replicate TF Bert & TF DPR (which is previously my implement) e.g. replace `inputs` with `input_ids` with/without default value (`None`) and use `input_processing` , but no matter what I tried , the simple call got error

```
outputs = model(inputs)
ValueError: The first argument to `Layer.call` must always be passed.
```

I really missed something simple here, I will try to play around again.

3) about the weights, yes, I tried to replicate other TF models, so all weights loading works in `.from_pretrained_question_encoder_generator`  and mostly properly loaded in `.from_pretrained`. 
Only two aforementioned weights could not have the correct name . I tried various fixing to these minor cases but could not, except my very ugly manual load.

1. Arf, I thought that this line would automatically deactivate the graph execution, but it is not the case. So to make it short, convert a graph tensor to numpy array is not possible because the graph does not execute in Python - so there is no numpy at graph execution. The only one work around would be to play with [tf.py_function](https://www.tensorflow.org/api_docs/python/tf/py_function) but you will literally kill the perf with this, even though it is your only way to go to.

2. This error means that you are not passing the first argument to a call method (basically don't pass any `input_ids`), you always have to pass it. Where in the code the error is raised?

3. What did you try more precisely?

Hi again Julien,

1. I see!  Let us work only on eager mode for now and come back later.

2.  Here's 4.1.0 colab where I try to modify 4.1.0 input on `TFRagModel` (Cell 6) 
https://colab.research.google.com/drive/1RvtOxUIravWEkwMnj48pedv2mFnlWYkY?usp=sharing
Please see Cell 9, to see various ways I try to pass the first argument. Really sorry I think I overlooked some simple things here.

3. I tried to adjust `base_prefix_name`, and also setting module `name` as discussed with Sam [here](https://discuss.huggingface.co/t/solved-issue-on-translating-dpr-to-tfdpr-on-loading-pytorch-weights-to-tf-model/1764/2): 


1. The problem is that you cannot do anything with the model if it cannot be run in graph mode (no serving, no training, no optimization, very slow) 😢 
2. Please, see how it is done in all the other TF implementations, the way you handle the inputs in `modeling_tf_rag.py` is wrong.
3. I don't really have time this week to go deeper in this, but I will take some time on Monday to do it!

A test file is missing as well, having it might help you to detect what has to be updated :)

Hi Julien!

> 1. The problem is that you cannot do anything with the model if it cannot be run in graph mode (no serving, no training, no optimization, very slow) 😢

I got it. I think we can do some work around on TFRag model training in graph mode. Instead of fitting the model with `input_ids` which will need `.numpy` and `retriever` , we can do all retriever stuffs offline (or with tf.Dataset) first to get the `context_input_ids` and then feed them directly to the trainnig loop. I will test this idea.

> 2. Please, see how it is done in all the other TF implementations, the way you handle the inputs in `modeling_tf_rag.py` is wrong.

Finally, I was able to find a single line that's wrong :) I will update the file in my local.

> 3. I don't really have time this week to go deeper in this, but I will take some time on Monday to do it!
> A test file is missing as well, having it might help you to detect what has to be updated :)

Thanks Julien. I will write the full test file soon. My reason is that I need some suggestions (as commented in the modeling file) for my current draft at this stage , so I have not written the full-fledge test file yet. The colab I posted did have some 10+ basic tests, which I still think I need to (cleanly) pass all these tests first. 


> I got it. I think we can do some work around on TFRag model training in graph mode. Instead of fitting the model with input_ids which will need .numpy and retriever , we can do all retriever stuffs offline (or with tf.Dataset) first to get the context_input_ids and then feed them directly to the trainnig loop. I will test this idea.

I was thinking the same :) to do the search offline, might be a solution. I think the very long term solution would be to use something more adapted to TensorFlow than FAISS such as [SCANN](https://github.com/google-research/google-research/tree/master/scann). (Pinging @lhoestq to know what he is thinking about this :) )

> Finally, I was able to find a single line that's wrong :) I will update the file in my local.

Nice!! Just to let you know that Friday we have also updated the way the handle the booleans, so be careful to integrate this as well.

> Thanks Julien. I will write the full test file soon. My reason is that I need some suggestions (as commented in the modeling file) for my current draft at this stage , so I have not written the full-fledge test file yet. The colab I posted did have some 10+ basic tests, which I still think I need to (cleanly) pass all these tests first.

It is ok, we are not in a hurry, take your time ^^

@ratthachat - I think you're on the correct track here :-)

TFRag will actually be the first TF composite model, so I'm quite certain you'll run into problems here that we haven't seen before. 
In general I think:

1) We should try to make the `from_pretrained` method work in a nice way (this will actually also show us how `TFEncoderDecoder` could be implemented). I'll take a look at this :-) 

2) Make `TFRagTokenForGeneration` work with integration tests that the model behaves the same way as PT using the faiss index. It's fine for me if it works only in eager mode for now. Maybe we can think about a different solution at a later stage if it's impossible to have RAG + Faiss in graph mode. I think it's a bit out-of-the-scope to integrate SCANN here with RAG.

3) Add the other functionalities. 

I'll try to help you with 1) here - will add some commits to your PR.

> We should try to make the from_pretrained method work in a nice way (this will actually also show us how TFEncoderDecoder could be implemented). I'll take a look at this :-)

I would like to remove all the `from_pretrained` calls from the model implementation, it will raises issues for some usage, such as training.

> Make TFRagTokenForGeneration work with integration tests that the model behaves the same way as PT using the faiss index. It's fine for me if it works only in eager mode for now. Maybe we can think about a different solution at a later stage if it's impossible to have RAG + Faiss in graph mode. I think it's a bit out-of-the-scope to integrate SCANN here with RAG.

The problem here is that if the model runs only in eager mode, the model won't be able to be served properly and then becomes useless. I don't see the point to have a model that runs only in your console locally :( The best solution IMHO would be to run the FAISS search offline.

Hey @ratthachat, 

I think the `from_pretrained()` functionality now works as expected. I removed some hacks and we shouldn't have to define any `from_pretrained()` method actually.

I've added two tests. 1 already passes (great job! - I didn't really change anything here...), the other one (a very difficult one based on `generate()` does not pass yet). It'll be quite difficult to make the other one pass, but if you manage you'll certainly have an in-depth knowledge of how `generate()` works.

I would recommend the following next steps for the PR:

1) Implement @jplu's simplified handling of the inputs as proposed in this PR: https://github.com/huggingface/transformers/pull/8602 . This should remove a lot of boiler plate code. If some params are not supported yet, I'm sure @jplu can help. 

2) Make the generate test pass

After this I'm happy to take another look :-) 

Lemme know if you have any problems with the weight loading. It all worked nicely for me

> Implement @jplu's simplified handling of the inputs as proposed in this PR: #8602 . This should remove a lot of boiler plate code. I some params are not supported yet, I'm sure @jplu can help.

I will be happy to help. Which ones are the "not supported yet"?

> > Implement @jplu's simplified handling of the inputs as proposed in this PR: #8602 . This should remove a lot of boiler plate code. If some params are not supported yet, I'm sure @jplu can help.
> 
> I will be happy to help. Which ones are the "not supported yet"?



> > Implement @jplu's simplified handling of the inputs as proposed in this PR: #8602 . This should remove a lot of boiler plate code. If some params are not supported yet, I'm sure @jplu can help.
> 
> I will be happy to help. Which ones are the "not supported yet"?

I think it should all work perfectly fine! Sorry, this came across a bit bad - meant to say "just in case" something doesn't work don't hesitate to ping you ;-) Wasn't sure if int inputs like `n_docs` are supported, but I think this was added as well - so it should all work fine :-) 

Thanks so much for your great help, Patrick! I will carefully look in each point you made. Full addressing will take a while, but I will be back. For now I have some initial responses: 

- (Need help the most) Unfortunately ,there's still a bug in weight loading if removing the hack (please see details in below thread)

- About graph mode & training, I think we can consistently combine we three's thoughts here by (a) Finish the code in eager mode first -- (b) Make minimal changes to support offline-mode for retriever (or maybe not change anything at all) -- (c) Make a community notebook to guide this offline-retrieved training in graph mode.  -- (d) SCANN will be an interesting long-term solution we can discuss after all these stuffs.

- I think there is similar graph-retrieval problem also in TFDPR training (which we haven't tested) , so I will also try make some example notebook to train TFDPR in graph mode using this offline principle.

- May I ask what is the meaning of these original Pytorch's 3 lines? (in `def generate()` )
```
            # retrieved_doc_embeds = retrieved_doc_embeds.to(question_hidden_states)
            # context_input_ids = context_input_ids.to(input_ids)
            # context_attention_mask = context_attention_mask.to(input_ids)
```

- About test on `generate`, I will try give it a shot. BTW, I previously test `Bart` vs. `TFBart` and found out that they produce **"different"**  `generate` results as well. 
Do you have the same experience, and will this affect RAG `generate` test ??


> I think the `from_pretrained()` functionality now works as expected. I removed some hacks and we shouldn't have to define any `from_pretrained()` method actually.
> 
> Lemme know if you have any problems with the weight loading. It all worked nicely for me

Hi Patrick, @patrickvonplaten 
Unfortunately, I found the same bug prior to my hack. 

- `from_pretrained_question_encoder_generator` <-- Work great

- `from_pretrained` <-- **BUG** (only on **_2 weights_**) : `['model.shared.weight', 'final_logits_bias']`

```
Some weights or buffers of the TF 2.0 model TFRagTokenForGeneration were not initialized from the PyTorch model and are newly initialized: ['model.shared.weight', 'final_logits_bias']
```

- (New found) local loading `from_pretrained("./rag")` <-- **BUG** on **all weights**

Please (please :) take a look at this new colab which provides "minimal" code to show the bugs (just 8 cells).
https://colab.research.google.com/drive/1s-j9PB9yzrFsL6q5rZUQyf8_Lt6jDAkL?usp=sharing
Bugs only in the last 2 cells.

> > I think the `from_pretrained()` functionality now works as expected. I removed some hacks and we shouldn't have to define any `from_pretrained()` method actually.
> > Lemme know if you have any problems with the weight loading. It all worked nicely for me
> 
> Hi Patrick, @patrickvonplaten
> Unfortunately, I found the same bug prior to my hack.
> 
> * `from_pretrained_question_encoder_generator` <-- Work great
> * `from_pretrained` <-- **BUG** (only on **_2 weights_**) : `['model.shared.weight', 'final_logits_bias']`

This is not a bug. It's fine actually. Those weights are handled differently in TF and PT, so this message is expected.

> 
> ```
> Some weights or buffers of the TF 2.0 model TFRagTokenForGeneration were not initialized from the PyTorch model and are newly initialized: ['model.shared.weight', 'final_logits_bias']
> ```
> 
> * (New found) local loading `from_pretrained("./rag")` <-- **BUG** on **all weights**

Let me look into this!

> 
> Please (please :) take a look at this new colab which provides "minimal" code to show the bugs (just 8 cells).
> https://colab.research.google.com/drive/1s-j9PB9yzrFsL6q5rZUQyf8_Lt6jDAkL?usp=sharing
> Bugs only in the last 2 cells.

@ratthachat,

I think you can ignore those warnings for now. A good next step to make sure that the `from_pretrained()` methods work correctly is to add tests that verify that after saving/loading the model yields the same output as before: 
- https://github.com/huggingface/transformers/blob/9d7d0005b046a95d9d59354714bb6c3547a612fe/tests/test_modeling_rag.py#L900

I checked and the following code works fully as expected:

```
from transformers import RagRetriever, TFRagTokenForGeneration

retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True)
model = TFRagTokenForGeneration.from_pretrained("facebook/rag-token-nq", from_pt=True, retriever=retriever)

model.save_pretrained("./rag")
model = TFRagTokenForGeneration.from_pretrained("./rag", retriever=retriever)
```

All those commands work as they should, so I think we're good for now with the `from_pretrained()`. I think the next step should be to concentrate on removing the TF input boilerplate code and then making the generation work.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.