Cannot save GPT2 model with signature
Author: ZhangTianrong
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: 4.5.0
- Platform: Windows-10-10.0.19041-SP0
- Python version: 3.7.4
- PyTorch version (GPU?): 1.8.1+cu111 (True)
- Tensorflow version (GPU?): 2.4.1 (True)
- Using GPU in script?: Yes
- 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.

Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @Rocketknight1

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

Model hub:

- for issues with a model report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

@VictorSanh @n1t0 @Pierrci 

## Information

Model I am using (Bert, XLNet ...): GPT2

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

I am trying to follow this [post](https://blog.tensorflow.org/2020/05/how-hugging-face-achieved-2x-performance-boost-question-answering.html) where @Pierrci illustrated how to convert a distilled BERT model into Tensorflow saved model and serve it with Tensorflow.js in the end. I would like to do similar stuff with a distilgpt2 model.

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [ ] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

```python
import tensorflow as tf
from transformers import TFGPT2LMHeadModel
from transformers import GPT2Tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")
model = TFGPT2LMHeadModel.from_pretrained("distilgpt2",  pad_token_id=tokenizer.eos_token_id)
callable = tf.function(model.call)
concrete_function = callable.get_concrete_function([tf.TensorSpec([None, 384], tf.int32, name="input_ids"), tf.TensorSpec([None, 384], tf.int32, name="attention_mask")])
tf.saved_model.save(model, 'distilgpt2_sig', signatures=concrete_function)
```

and the error messages are as follow

```
ValueError: Got a non-Tensor value (<tf.Tensor 'StatefulPartitionedCall:1' shape=(2, None, 12, 384, 64) dtype=float32>, <tf.Tensor 'StatefulPartitionedCall:2' shape=(2, None, 12, 384, 64) dtype=float32>, <tf.Tensor 'StatefulPartitionedCall:3' shape=(2, None, 12, 384, 64) dtype=float32>, <tf.Tensor 'StatefulPartitionedCall:4' shape=(2, None, 12, 384, 64) dtype=float32>, <tf.Tensor 'StatefulPartitionedCall:5' shape=(2, None, 12, 384, 64) dtype=float32>, <tf.Tensor 'StatefulPartitionedCall:6' shape=(2, None, 12, 384, 64) dtype=float32>, <tf.Tensor 'StatefulPartitionedCall:7' shape=(2, None, 12, 384, 64) dtype=float32>, <tf.Tensor 'StatefulPartitionedCall:8' shape=(2, None, 12, 384, 64) dtype=float32>, <tf.Tensor 'StatefulPartitionedCall:9' shape=(2, None, 12, 384, 64) dtype=float32>, <tf.Tensor 'StatefulPartitionedCall:10' shape=(2, None, 12, 384, 64) dtype=float32>, <tf.Tensor 'StatefulPartitionedCall:11' shape=(2, None, 12, 384, 64) dtype=float32>, <tf.Tensor 'StatefulPartitionedCall:12' shape=(2, None, 12, 384, 64) dtype=float32>) for key 'past_key_values' in the output of the function __inference_call_90110 used to generate the SavedModel signature 'serving_default'. Outputs for functions used as signatures must be a single Tensor, a sequence of Tensors, or a dictionary from string to Tensor.
```

I can save the model if I don't specify `signitures` but in that case the input shape is default to [-1,5].

<!-- 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.-->

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->

I expect the model to be saved without a problem.


Comments:
Tbh, I've never really worked with `get_concrete_function()`, etc.... @Rocketknight1 - do you have an idea by any chance?

> Tbh, I've never really worked with `get_concrete_function()`, etc.... @Rocketknight1 - do you have an idea by any chance?

I just realized I should tag the authors of the post I read about in the issue. I have edited the issue.

Hi, I'm the TF maintainer! There are two problems here. The first is that the first two arguments to `TFGPT2LMHeadModel` are not `input_ids` and `attention_mask`, they are `input_ids` and `past`, see [here](https://huggingface.co/transformers/model_doc/gpt2.html#tfgpt2lmheadmodel). Also, `TFGPT2LMHeadModel` returns a tuple/dict of Tensors. Concrete functions do not support that - you need to pick which one you want. Try something like this, which should work (if you want an output other than "logits", you can just change that bit):

```
import tensorflow as tf
from transformers import TFGPT2LMHeadModel
from transformers import GPT2Tokenizer

@tf.function
def call_model(input_ids, attention_mask):
    return model(input_ids=input_ids, attention_mask=attention_mask)['logits']

tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")
model = TFGPT2LMHeadModel.from_pretrained("distilgpt2",  pad_token_id=tokenizer.eos_token_id)
concrete_function = call_model.get_concrete_function(tf.TensorSpec([None, 384], tf.int32, name="input_ids"), tf.TensorSpec([None, 384], tf.int32, name="attention_mask"))
tf.saved_model.save(model, 'distilgpt2_sig', signatures=concrete_function)
```

> Hi, I'm the TF maintainer! There are two problems here. The first is that the first two arguments to `TFGPT2LMHeadModel` are not `input_ids` and `attention_mask`, they are `input_ids` and `past`, see [here](https://huggingface.co/transformers/model_doc/gpt2.html#tfgpt2lmheadmodel). Also, `TFGPT2LMHeadModel` returns a tuple/dict of Tensors. Concrete functions do not support that - you need to pick which one you want. Try something like this, which should work (if you want an output other than "logits", you can just change that bit):
> 
> ```
> import tensorflow as tf
> from transformers import TFGPT2LMHeadModel
> from transformers import GPT2Tokenizer
> 
> @tf.function
> def call_model(input_ids, attention_mask):
>     return model(input_ids=input_ids, attention_mask=attention_mask)['logits']
> 
> tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")
> model = TFGPT2LMHeadModel.from_pretrained("distilgpt2",  pad_token_id=tokenizer.eos_token_id)
> concrete_function = call_model.get_concrete_function(tf.TensorSpec([None, 384], tf.int32, name="input_ids"), tf.TensorSpec([None, 384], tf.int32, name="attention_mask"))
> tf.saved_model.save(model, 'distilgpt2_sig', signatures=concrete_function)
> ```

That works! Thank you for the help. I am not familiar with TF especially things like `get_concrete_function`, I didn't know you can define a function outside the model and then save it.