Fix TF s2s models
Author: jplu
Pull request: yes

# What does this PR do?

This PR aims to fix the Seq2Seq models in order to make them able to be served through TF Serving.  The problem is stated by @patrickvonplaten in #9313. The reason why it failed was because we use a model as a layer in the `TFXXXForConditionalGeneration` models. The tracing mechanism of TensorFlow when building a graph calls one by one all the layers for building the graph. In order to know what are the inputs needed by each layer, the tracing mechanism check if a layer has a custom input signature, if not, it takes as default a signature where only the first argument is mandatory. Here stand the problem, the Seq2Seq models needs have two mandatory arguments (`input_ids` and `decoder_input_ids` or `inputs_embeds` and `decoder_inputs_embeds`) and then the tracing fails.

The fix to this problem is to manually set the expected input signature of the base model when instantiating it in `__init__`. To be harmonized with the required serving, the same signature is used.




Comments:
@patrickvonplaten Should I remove the following hack in BART?
```python
if inputs["decoder_input_ids"] is None and inputs["input_ids"] is not None:
            inputs["decoder_input_ids"] = shift_tokens_right(
                inputs["input_ids"], self.config.pad_token_id, self.config.eos_token_id
            )
```

> In general I really don't like the tf.cond(condition, do_fn_one, do_fn_two) design. I think I understand that it is sometimes necessary, but I really like to keep the usage of this function to a minimum in general. The functional approach is very different to our general library design and make the code much much harder to read. It always creates an abstraction by having to wrap parts of the code into a function with no args, like def attn_mask_from_inp_ids() which is not easy to follow and to me always looks like a hack. In Bart we manage to do this part of the code without the usage of tf.cond and Bart has the same exact logic as LED has there -> so we can make it easier I think.

I understand your point and agree with you and I share you opinion on this, and unfortunately if you come to control flow (conditions and loops) there are some strict rules that one cannot overcome. `tf.cond` is somehow mandatory for autograph.

A solution I think that should work would be to force the `layer.call` function to be run in graph mode with `@tf.function` which takes care of making itself the translation of all these conditions and loops. This work in some cases, let's see if it works... Does it sounds a proper solution for you?

> @patrickvonplaten Should I remove the following hack in BART?
> 
> ```python
> if inputs["decoder_input_ids"] is None and inputs["input_ids"] is not None:
>             inputs["decoder_input_ids"] = shift_tokens_right(
>                 inputs["input_ids"], self.config.pad_token_id, self.config.eos_token_id
>             )
> ```

Please don't - it's needed for some use-cases in Bart and for backward comp

Actually, one thing I'd like to know more in general about our models in TF is the following: 

"Can we use normal if-else statements in the forward pass"? 

I always thought that the answer is:

"Yes we can as long as the output type and shape of each case is the same" 

So for me statements like:

```python 
if shape_list(input_ids) > n:
    attention_mask = torch.zeros(shape_list(input_ids))
else:
   attention_mask = torch.ones(shape_list(input_ids))
```
(this code snippet doesn't exist -> it's just an example)

are totally fine. Is the assumption correct @jplu ?

Or can we in general **never** use normal if-else statements in TF's forward pass and have to rely on `tf.cond(....)`? This would really surprise me as we're having tons of if statements everywhere in the TF code...



The general answer is yes, but it has some conditions. If you run this condition in eager mode, it will works by default (you can basically do almost anything in eager mode)

If you run this condition in graph mode you have two solution to make it works:
1.  Either use `tf.cond`
2.  Or to wrap your condition into a function decorated with `tf.function`. This will have to effect to apply the Autograph library over the content of your decorated function. Autograph will automatically converts `if-then` clauses, loops, `break`, `return`, `continue`, and more. You can have more information here https://www.tensorflow.org/guide/function#autograph_transformations

Now that we all agree on a solution, I will apply it for all the models 👍 

Ok, LGTM!! Feel free to merge whenever you feel it^^