Inconsistent input signatures for gpt2-medium tensorflow
Author: Spycsh
Pull request: no

### System Info

- `transformers` version: 4.33.0.dev0
- Platform: Linux-5.16.0-rc8-intel-next-01534-g53cb5f883cf7-x86_64-with-glibc2.10
- Python version: 3.8.0
- Huggingface_hub version: 0.17.1
- Safetensors version: 0.3.3
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.0.1+cpu (False)
- Tensorflow version (GPU?): 2.12.0 (False)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@gante @Rocketknight1 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Here are the minimal reproduction steps:

1. Using transformers 4.33.0, run following code

```
import tensorflow as tf
import transformers

model = transformers.TFAutoModelForCausalLM.from_pretrained("gpt2")
model.save("cur")
loaded = tf.saved_model.load("cur")
print(loaded.signatures["serving_default"])
```

And you can see the input signature as:

```
ConcreteFunction signature_wrapper(*, attention_mask, input_ids, token_type_ids)
```

2. Reinstall old version transformers 4.29.2, run nearly the same code (just change the name 'cur' to 'prev' here),

```
import tensorflow as tf
import transformers

model = transformers.TFAutoModelForCausalLM.from_pretrained("gpt2")
model.save("prev")
loaded = tf.saved_model.load("prev")
print(loaded.signatures["serving_default"])
```

you will get

```
ConcreteFunction signature_wrapper(*, attention_mask, input_ids)
```


Could you give me some hints why the latest version contains the `token_type_ids` in the input signature? How to get rid of it? I tried to pass None to parameter in the feed_dict but it did not work.

### Expected behavior

I expect the input signature should be consistent to the old version (4.29.2), which only input_ids and attention_mask is needed.

Comments:
Hi @Spycsh, this was caused by a refactor in our TF input signatures, where they're now autodetected in almost all cases. As a result, GPT2 got `token_type_ids` in its input signature because the model does actually support them as an input. However, it's easy to save with a different signature if required!

Try this code snippet:

```python
import tensorflow as tf
import transformers

model = transformers.TFAutoModelForCausalLM.from_pretrained("gpt2")
default_signature = model.input_signature  # This is just a dict
del default_signature['token_type_ids']  # Remove any keys you don't want
serving_default = model.serving.get_concrete_function(default_signature)

model.save("cur", signatures={"serving_default": serving_default})
```

And you're done! By changing the `default_signature` dict here you can use whatever signature you want, adding or removing keys or setting custom dtypes, etc.




Thanks!

No probs! I'm going to close the issue now, but feel free to comment or reopen it if you hit any other trouble with our input sigs.

Hi @Rocketknight1 , can we directly update the input signatures of the model without save and load? If we save and load that model, the model will be a tensorflow.python.saved_model.load.Loader._recreate_base_user_object.<locals>._UserObject object, which will be hard to deal with in our further logics.

I suspect not, unfortunately! When a TF model is traced and compiled and saved as SavedModel, it's probably quite hard to do 'model surgery' on it afterwards to remove the token_type_ids. It might be easier to create a new model in Python using `TFAutoModelForCausalLM` and then maybe load the weights from the checkpoint into the Python model, and then finally save the Python model with the desired signature?

Thanks @Rocketknight1 , maybe the only thing from our side that we are able to do is to accept that now the signature is `signature_wrapper(*, attention_mask, input_ids, token_type_ids)`. Saving another signature to it will change the model type and is infeasible from our side. You say that:

> GPT2 got token_type_ids in its input signature because the model does actually support them as an input.

My question is, how to use the current infer signature to simulate the previous invocation. IMO, we should give a feed_dict like {"input_ids": xxx, "attention_mask": xxx, "token_type_ids": None} to the concrete function, right? However, we cannot pass a `None` to the `token_type_ids` as a parameter of the serving function because it is not allowed to do so in TF v1.

Overall, could you give me some hints on whether you think it is possible to mimic the original call that set the token_type_ids to None. More concretely, if we use the original GPT2 model forward call, we should use it like 
```
model(input_ids=xxx,attention_mask=xxx, token_type_ids=None)
```

So what should be the equivalent of that by using the current infer signature function `signature_wrapper(*, attention_mask, input_ids, token_type_ids)`, when TF does not support a feed_dict `token_type_ids` value to be None?

Here is a real example code with latest Transformers that you can review to understand the problem, where we now have to feed token_type_ids as None to the serving function https://github.com/intel/neural-compressor/blob/master/examples/tensorflow/nlp/large_language_models/quantization/ptq/smoothquant/main.py#L82

Welcome to any suggestions!

Hi @Spycsh, I dug into this and it's actually trickier than I thought. For most models that support `token_type_ids`, when the value isn't set then the default is just to create an all-zeros array, something like `tf.zeros_like(input_ids)`. Therefore, for those models you could easily mimic `token_type_ids=None` by instead just passing a `tf.zeros()` array.

However, GPT-2 is an old model now, and so its implementation of token type IDs is quite odd! Specifically, the input embeddings are also used as the token type embeddings, but when `token_type_ids=None`, then token type embeddings are completely skipped. This means that **any** token_type_ids you pass will affect the embeddings, and so there is no tensor value you can pass that will yield the same results as `token_type_ids=None`.

Given this, there's not a lot you can do when a model has already been exported with `token_type_ids` in its input signature, except to load the weights from it into a new `TFGPT2Model` and then save that model with the correct signature. I can definitely see how annoying this is in the specific case of GPT-2, so I'll probably make a PR to remove `token_type_ids` from its default input signature, so at least this stops happening in future.

Thanks @Rocketknight1 ! I will try to verify our example on your branch:)

@Rocketknight1 , our example works correctly with the branch to your latest [PR](https://github.com/huggingface/transformers/pull/26962) that explicitly set the signature. Thanks and hope to see it merged!

@spycsh merged!