Error with inferencing with fine-tuned model loaded from keras load_model
Author: jamie0725
Pull request: no
Labels: name: bug

### System Info

```shell
transformers version: 4.16.2
tensorflow version: 2.9.1
```


### Who can help?

@gante 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```python
tokenizer = BertTokenizer.from_pretrained('bert-base-cased')
model = TFBertForSequenceClassification.from_pretrained('bert-base-cased', num_labels=2)

imdb = load_dataset('imdb')
def preprocess_function(examples):
    return tokenizer(examples["text"], truncation=True)
tokenized_imdb = imdb.map(preprocess_function, batched=True)

train_ds = tokenized_imdb['train']
val_test_ds = tokenized_imdb['test'].train_test_split(test_size=0.5)
val_ds, test_ds = val_test_ds['train'], val_test_ds['test']
print(len(train_ds), len(val_ds), len(test_ds))

data_collator = DataCollatorWithPadding(tokenizer=tokenizer, return_tensors="tf")

batch_size = 8
tf_train_ds = train_ds.to_tf_dataset(
    columns=['input_ids', 'attention_mask', 'token_type_ids'],
    label_cols=['label'],
    shuffle=True,
    drop_remainder=False,
    batch_size=batch_size,
    collate_fn=data_collator
)
tf_val_ds = val_ds.to_tf_dataset(
    columns=['input_ids', 'attention_mask', 'token_type_ids'],
    label_cols=['label'],
    shuffle=True,
    drop_remainder=False,
    batch_size=batch_size,
    collate_fn=data_collator
)
tf_test_ds = test_ds.to_tf_dataset(
    columns=['input_ids', 'attention_mask', 'token_type_ids'],
    label_cols=['label'],
    shuffle=False,
    drop_remainder=False,
    batch_size=batch_size,
    collate_fn=data_collator
)

num_epochs = 10
batches_per_epoch = len(tf_train_ds)
num_train_steps = batches_per_epoch * num_epochs
num_warmup_steps = int(0.1 * num_train_steps)
init_lr = 2e-5

optimizer = optimization.create_optimizer(
      init_lr=init_lr,
      num_train_steps=num_train_steps,
      num_warmup_steps=num_warmup_steps,
      optimizer_type='adamw')

model.compile(
    optimizer=optimizer,
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=[tf.keras.metrics.SparseCategoricalAccuracy()]
)

model.save("./best_model")

re_model = tf.keras.models.load_model('./best_model', custom_objects={'AdamWeightDecay': optimizer})

# it works fine with serving() method
for i, j in tf_test_ds:
    a = re_model.serving({
        'input_ids': tf.cast(i['input_ids'], tf.int32),
        'token_type_ids': tf.cast(i['token_type_ids'], tf.int32),
        'attention_mask': tf.cast(i['attention_mask'], tf.int32),
                         })
    print(a['logits'].numpy().argmax(axis=1))
    break

# it does not work with __call__ method
results = re_model.evaluate(x=tf_test_ds)
for name, value in zip(re_model.metrics_names, results):
    print("%s: %.3f" % (name, value))
```

error message:
```
ValueError: Exception encountered when calling layer "tf_bert_for_sequence_classification" (type TFBertForSequenceClassification).
    
    Could not find matching concrete function to call loaded from the SavedModel. Got:
      Positional arguments (11 total):
        * {'attention_mask': <tf.Tensor 'input_ids:0' shape=(None, None) dtype=int64>,
     'input_ids': <tf.Tensor 'input_ids_1:0' shape=(None, None) dtype=int64>,
     'token_type_ids': <tf.Tensor 'input_ids_2:0' shape=(None, None) dtype=int64>}
        * None
        * None
        * None
        * None
        * None
        * None
        * None
        * None
        * None
        * False
      Keyword arguments: {}
    
     Expected these arguments to match one of the following 2 option(s):
    
    Option 1:
      Positional arguments (11 total):
        * {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids/input_ids')}
        * None
        * None
        * None
        * None
        * None
        * None
        * None
        * None
        * None
        * False
      Keyword arguments: {}
    
    Option 2:
      Positional arguments (11 total):
        * {'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids/input_ids')}
        * None
        * None
        * None
        * None
        * None
        * None
        * None
        * None
        * None
        * True
      Keyword arguments: {}
    
    Call arguments received by layer "tf_bert_for_sequence_classification" (type TFBertForSequenceClassification):
      • args=({'input_ids': 'tf.Tensor(shape=(None, None), dtype=int64)', 'token_type_ids': 'tf.Tensor(shape=(None, None), dtype=int64)', 'attention_mask': 'tf.Tensor(shape=(None, None), dtype=int64)'},)
      • kwargs={'training': 'False'}

```
Basically, from what I could tell, the reloaded TFBertForSequenceClassification model seems to accept only `input_ids` as the only input, while before saving the model with keras, it also allows `attention_mask` and `token_type_ids`. Is it something related to input signatures of the __call__ function?

### Expected behavior

```shell
The model saved with keras.save() works normally after loaded from tf.keras.models.load_model(...).

i.e. the __call__ method should recoginize the correct input format as serving() method.
```


Comments:
Hi @jamie0725 👋  Thank you for raising this issue. We are aware of problems like this, related to serialization/deserialization of Keras models, and it is in our plans to fix it.

I'm keeping the issue open for tracking purposes :)

> Hi @jamie0725 👋 Thank you for raising this issue. We are aware of problems like this, related to serialization/deserialization of Keras models, and it is in our plans to fix it.
> 
> I'm keeping the issue open for tracking purposes :)

Thanks for notice, please let me know where there is a fix.

Meanwhile, is there any workaround I can do except for using the ``serving()`` method? Because ``serving()`` does not support keras methods like `evaluate()` etc.

There is, working with TF's concrete functions. Have a look at [this comment](https://github.com/huggingface/transformers/issues/11296#issuecomment-825798041) for an example :)

(We will probably build documentation for it soon, as part of our plans)

> There is, working with TF's concrete functions. Have a look at [this comment](https://github.com/huggingface/transformers/issues/11296#issuecomment-825798041) for an example :)
> 
> (We will probably build documentation for it soon, as part of our plans)

I tried, but it gives the same error. My code is as below:

```
@tf.function
def call_model(input_ids, attention_mask):
    return model(input_ids=input_ids, attention_mask=attention_mask)

concrete_function = call_model.get_concrete_function(tf.TensorSpec([None, None], tf.int32, name="input_ids"), tf.TensorSpec([None, None], tf.int32, name="attention_mask"))

model.save('./best_model', signatures=concrete_function)

re_model = tf.keras.models.load_model('./best_model', custom_objects={'AdamWeightDecay': optimizer})

results = re_model.evaluate(x=tf_test_ds)
for name, value in zip(re_model.metrics_names, results):
    print("%s: %.3f" % (name, value))
```

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.