RAG : Adding end to end training for the retriever (both question encoder and doc encoder) 
Author: shamanez
Pull request: no
Labels: name: Feature request

# 🚀 Feature request

As mentioned in this recent paper [End-to-End Training of Neural Retrievers for Open-Domain Question Answering](https://arxiv.org/abs/2101.00408), we can get better results for QA tasks if we fine-tune the retriever section in an end-to-end manner. 


## Paper's method

Fine-tune both doc-encoder and question encoder and update the pre computed index embedding in every 500 steps.   


Comments:
 @lhoestq


Interesting :) I don't think I'll be able to work on this in the short term but if anyone wants to give it a try maybe I can help with some indications

@lhoestq 

I kind of figured out a way to do this. I need a little clarification from you regarding the distributed retriever.  As mentioned in this [line](https://github.com/huggingface/transformers/blob/master/examples/research_projects/rag/finetune_rag.py#L84), we use **CustomAcc** class  to load knowledgebase and load faiss index as a separate process. 

I want to re-execute the above-mentioned process after several training steps. Let's say 1000.

with Pytorch lighting,



    def optimizer_step(self, epoch_nb, batch_nb, optimizer, optimizer_i, opt_closure):
        if self.trainer.global_step < 500:
             ****** run init_ddp_connection functiontion inside CustomAccel class*****
            1. Reinitialize knowledgebase dataset
            2. Relaod faiss index

 




Hi @shamanez You can reload an updated index during train time this way:
1. recompute all the embeddings of your knowledge source using the context encoder (costly, depending on your knowledge source size)
2. recreate the FAISS index (which can also be costly)

The REALM model does this kind of back and forth between training and indexing, you may want to check out how they did that in the paper.

I think one approach would be to extend the [RayRetriever](https://github.com/huggingface/transformers/blob/master/examples/research_projects/rag/distributed_ray_retriever.py) since you can define several workers depending on what you want to do (query the index, compute the embeddings or update the index). It's something that I feel is more natural to do with Ray than with pytorch distributed.

Ok. I was thinking the same steps looking at the REALM paper. In their code implementation, they run three separate processes and communicate in between processes when they need to compute new embedding, load them and finally feed them to the reader model. It only works with a single GPU.

Anyways I get when using RAG with PyTorch distributed, the loading operation is done before the trainer. So I was thinking can we execute that within the training loop. Anyways I get what you say. 


@amogkam  can you help with this?




@lhoestq 

I kind of tried to implement this with PyTorch distributed retriever method.  So ideally I want to re-load the knowledge-based and rea-load the indexes inside the training step (assuming I have an updated doc-encoder). Here is my implementation. Can you please let me know whether it is correct. 

```
def training_step(self, batch, batch_idx) -> Dict:

        if not  batch_idx==0 and batch_idx%10000==0:
                  self.model.retriever.re_index()

```
The reindex is a simple method inside **distributed_pytorch_retriever.py**  that only re-load the dataset and idex  in the main process.

```
def re_index(self):
        # initialize retriever only on the main worker
        if self._is_main():
            logger.info("re initializing the index")
            self.index.re_init_index()

```




Here my assumption is, we have already started a separate process with the custom accel. Now we are changing something inside it. 

What do u this of it?



It looks good :) Although I haven't tested things like that so let me know how it goes !

also one detail: maybe you wanted to write `batch_idx%10000 == 0` instead of `batch_idx%10000`

ah yeah, thanks for the update. So I am on it. I Will update you soon :) 

@lhoestq  I have another question regarding the **load_dataset** function:  

Prior to starting the DDP process, the code loads the indexed dataset by accessing the saved file in the hard disk with the **load_from_disk** .([this line](https://github.com/huggingface/transformers/blob/641f418e102218c4bf16fcd3124bfebed6217ef6/src/transformers/models/rag/retrieval_rag.py#L397)).   


During the training what if the data file (.arrow files) change?  Here, the entire data structure is the same, it is just the values that change. 


In this kind of scenario do we have to use the load_dataset function again or it will automatically access the updated file?  

You would need to create a new arrow file and load the new arrow file.
If you overwrite the arrow file that is currently loaded I'm pretty sure things won't get updated properly.

yeah, that is what I actually observed.   Btw  I have implemented the end-to-end case with RAY. Currently doing the final testing. Will do a pull request if it is possible. 

This is really cool thanks !