[Rag] extend_enc_output fails when number of retrieved documents not equal to RagConfig.n_docs
Author: lalitpagaria
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: 3.3.1
- Platform: Linux-4.19.112+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.6.9
- PyTorch version (GPU?): 1.6.0+cu101 (False)
- Tensorflow version (GPU?): 2.3.0 (False)
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.
 
 albert, bert, GPT2, XLM: @LysandreJik 
 tokenizers: @mfuntowicz
 Trainer: @sgugger
 Speed and Memory Benchmarks: @patrickvonplaten
 Model Cards: @julien-c
 Translation: @sshleifer
 Summarization: @sshleifer
 TextGeneration: @TevenLeScao 
 examples/distillation: @VictorSanh
 nlp datasets: [different repo](https://github.com/huggingface/nlp)
 rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)
 Text Generation: @TevenLeScao
 blenderbot: @mariamabarham
 Bart: @sshleifer
 Marian: @sshleifer
 T5: @patrickvonplaten
 Longformer/Reformer: @patrickvonplaten
 TransfoXL/XLNet: @TevenLeScao 
 examples/seq2seq: @sshleifer
 examples/bert-loses-patience: @JetRunner
 tensorflow: @jplu
 examples/token-classification: @stefan-it
 documentation: @sgugger
 -->
@patrickvonplaten

## Information

Model I am using (Bert, XLNet ...): Rag

The problem arises when using:
* [X] my own modified scripts: (give details below)

The tasks I am working on is:
* [X] an official GLUE/SQUaD task: `dummy_dataset`

## To reproduce

Steps to reproduce the behavior:

1. Use caling retriever seperately example and but modify number of retrieved document parameter `n_docs` as follows -
```
docs_dict = retriever(input_ids.numpy(), question_hidden_states.detach().numpy(), return_tensors="pt", n_docs=2)
```


<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->
Code snippet -
```
tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True)
model = RagTokenForGeneration.from_pretrained("facebook/rag-token-nq", use_dummy_dataset=True)

input_dict = tokenizer.prepare_seq2seq_batch("What is capital of France?", return_tensors="pt")
input_ids = input_dict["input_ids"]

# Caling retriever seperately
question_hidden_states = model.question_encoder(input_ids)[0]

# Default value of RagConfig.n_docs is 5 hence change n_docs to different than default value
docs_dict = retriever(input_ids.numpy(), question_hidden_states.detach().numpy(), return_tensors="pt", n_docs=2)
doc_scores = torch.bmm(question_hidden_states.unsqueeze(1), docs_dict["retrieved_doc_embeds"].float().transpose(1, 2)).squeeze(1)

outputs = model.generate(context_input_ids=docs_dict["context_input_ids"], context_attention_mask=docs_dict["context_attention_mask"], doc_scores=doc_scores)
generated_string = tokenizer.batch_decode(outputs, skip_special_tokens=True)
print(generated_string)
```
Stacktrace -
```
/usr/local/lib/python3.6/dist-packages/transformers/modeling_rag.py in generate(self, input_ids, attention_mask, context_input_ids, context_attention_mask, doc_scores, max_length, min_length, early_stopping, use_cache, num_beams, bos_token_id, pad_token_id, eos_token_id, length_penalty, no_repeat_ngram_size, bad_words_ids, num_return_sequences, decoder_start_token_id, **kwargs)
   1353 
   1354         # correctly extend last_hidden_state and attention mask
-> 1355         context_attention_mask = extend_enc_output(context_attention_mask, num_beams=num_beams)
   1356         encoder_outputs["last_hidden_state"] = extend_enc_output(last_hidden_state, num_beams=num_beams)
   1357 

/usr/local/lib/python3.6/dist-packages/transformers/modeling_rag.py in extend_enc_output(tensor, num_beams)
   1346         def extend_enc_output(tensor, num_beams=None):
   1347             # split into `batch_size`, `num_beams`, `num_docs`
-> 1348             tensor = tensor[None, None, :].reshape((batch_size, 1, self.config.n_docs) + tensor.shape[1:])
   1349             # repeat same last hidden states over `num_beams` dimension
   1350             tensor = tensor.expand((batch_size, num_beams, self.config.n_docs) + tensor.shape[3:])

RuntimeError: shape '[0, 1, 5, 300]' is invalid for input of size 600
```
## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->
I think `extend_enc_output` should use tensor dimension instead of  `self.config.n_docs` as follows -
```
tensor = tensor[None, None, :].reshape((batch_size, 1, tensor.shape[0:]) + tensor.shape[1:])
```
Let me know if proposed fix is fine, I can raise PR.

Comments:
@lalitpagaria - good point! I think the cleanest solution here would actually be to add `n_docs` as a argument to all RagModel forward's function that defaults to `self.config.n_docs` (exactly as it is the case for `use_cache` or `output_attentions`) => it would be great if you could open a PR for this :-) 

@patrickvonplaten I think this approach also will not solve a situation when retriever return less docs than `n_docs`.  For example let's say a dataset only have `4` documents but `n_docs` value is `5`. 

How about solving this as follows -
1) When fetching docs from retriever use `n_docs` value (either as a argument or default `self.config.n_docs`)
2) All other places use value from number of documents returned by retriever ie dimension `0` of `context_input_ids` can be used

Let me know what do you think.

I think this is an edge case - if a dataset only has 4 documents, then the user should not define `n_docs` as 5. I agree that it would be easier to just use the incoming tensor's dimension as `n_docs`, but it would be less readable and can also lead to confusion. `n_docs` should always be the number of docs actually given to Rag. So I'd prefer to not use the retriever's dimension here. We could add an assert that the dimension 0 is always equal `n_docs` which would be helpful for the user, but I think we should still rely on the `n_docs` parameter.

Sure I will create PR (using n_docs with assert) with related tests.