Cannot serialize Whisper decoder layer in a keras model

[perretv]

System Info

• transformers version: 4.30.2
• Platform: Linux-4.19.0-24-cloud-amd64-x86_64-with-glibc2.28
• Python version: 3.11.3
• Huggingface_hub version: 0.15.1
• Safetensors version: 0.3.1
• PyTorch version (GPU?): not installed (NA)
• Tensorflow version (GPU?): 2.12.0 (True)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using GPU in script?: yes
• Using distributed or parallel set-up in script?: no

Who can help?

@Rocketknight1

Information

• The official example scripts
• My own modified scripts

Tasks

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

Reproduction

Changes introduced in #23760 and released in transformers v4.30.0 are breaking the ability to serialize a keras model that contains a Whisper decoder layer.
Here is a minimal reproducible example:

from transformers import TFWhisperModel
import tensorflow as tf

whisper = TFWhisperModel.from_pretrained("openai/whisper-tiny")
inp = tf.keras.Input((80, 3000))
stack = whisper.get_encoder()(inp)
decoder_input_ids = tf.ones((tf.shape(inp)[0], 1), dtype=tf.int32)* whisper.config.decoder_start_token_id
stack = whisper.get_decoder()(input_ids=decoder_input_ids, encoder_hidden_states=stack.last_hidden_state)
model = tf.keras.Model(inp, stack)
model.summary()
model.save("whisper-tiny-custom")

With transformers>=4.30.0, this minimal reproducible example will raise the error:

OperatorNotAllowedInGraphError: Exception encountered when calling layer 'decoder' (type TFWhisperDecoder).

Using a symbolic `tf.Tensor` as a Python `bool` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

Call arguments received by layer 'decoder' (type TFWhisperDecoder):
  • self=tf.Tensor(shape=(1, 1), dtype=int32)
  • input_ids=None
  • attention_mask=None
  • position_ids=None
  • encoder_hidden_states=tf.Tensor(shape=(None, 1500, 384), dtype=float32)
  • head_mask=None
  • cross_attn_head_mask=None
  • past_key_values=None
  • inputs_embeds=None
  • use_cache=None
  • output_attentions=None
  • output_hidden_states=None
  • return_dict=None
  • training=True

Expected behavior

Keras model serialization should succeed.

[Rocketknight1]

Hi @perretv yes, this looks like a regression. Investigating now, hopefully we can make a quick patch!

[Rocketknight1]

Hi @perretv, the fix has now been merged. You can pip install git+https://github.com/huggingface/transformers.git to install from main and use it immediately. It'll be included in the next 4.31 release of transformers, after which you can go back to normal pip installs.