List index out of range when saving finetuned Tensorflow model

Viewed 1183

I'm trying to fine-tune a pre-trained BERT model from Huggingface using Tensorflow. Everything runs smoothly and the model builds and trains without error. But when I try to save the model it stops with the error "IndexError: list index out of range". I'm using Google Colab with TPU.

Any help would be much appreciated!

Code:

import tensorflow as tf
from tensorflow.keras import activations, optimizers, losses
from transformers import TFBertModel

def create_model(max_sequence, model_name, num_labels):
    bert_model = TFBertModel.from_pretrained(model_name)
    input_ids = tf.keras.layers.Input(shape=(max_sequence,), dtype=tf.int32, name='input_ids')
    attention_mask = tf.keras.layers.Input((max_sequence,), dtype=tf.int32, name='attention_mask')
    output = bert_model([input_ids, attention_mask])[0]
    output = output[:, 0, :]
    output = tf.keras.layers.Dense(num_labels, activation='sigmoid')(output)
    model = tf.keras.models.Model(inputs=[input_ids, attention_mask], outputs=output)
    return model

with strategy.scope():
  model = create_model(20, 'bert-base-uncased', 1)
  opt = optimizers.Adam(learning_rate=3e-5)
  loss = 'binary_crossentropy'
  model.compile(optimizer=opt, loss=loss, metrics=['accuracy'])

model.fit(tfdataset_train, batch_size=32, epochs=2)
SAVE_PATH = 'path/to/save/location'
model.save(SAVE_PATH)

Error:

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-22-255116b49022> in <module>()
      1 SAVE_PATH = 'path/to/save/location'
----> 2 model.save(SAVE_PATH,save_format='tf')

50 frames
/usr/local/lib/python3.7/dist-packages/transformers/modeling_tf_utils.py in input_processing(func, config, input_ids, **kwargs)
    372                     output[tensor_name] = input
    373                 else:
--> 374                     output[parameter_names[i]] = input
    375             elif isinstance(input, allowed_types) or input is None:
    376                 output[parameter_names[i]] = input

IndexError: list index out of range

Model plotted with shapes: Tensorflow Model

3 Answers

The solution is to change:

output = bert_model([input_ids, attention_mask])[0]

to

output = bert_model.bert([input_ids, attention_mask])[0]

Reference: https://github.com/huggingface/transformers/issues/3627

I upvoted the solution you posted, but then I discovered a problem with the model while training. It doesn't converge well.

I got the same error while trying to save a fine-tuned BERT model (using huggingface and tensorflow) that uses all three input tensors - input_ids, token_type_ids, and attention_mask.

The solution that works for me is a slightly modification of user_007 answer, as follow:

output = bert_model.bert(
    input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask
)[0]

Problem solved!

Removing one of the two input layers (i.e. attention_mask) solved the problem.

Working code ->

import tensorflow as tf
from tensorflow.keras import activations, optimizers, losses
from transformers import TFBertModel

def create_model(max_sequence, model_name, num_labels):
    bert_model = TFBertModel.from_pretrained(model_name)
    input_ids = tf.keras.layers.Input(shape=(max_sequence,), dtype=tf.int32, name='input_ids')
    #attention_mask = tf.keras.layers.Input((max_sequence,), dtype=tf.int32, name='attention_mask')
    #output = bert_model([input_ids, attention_mask])[0]
    output = bert_model([input_ids])[0]
    output = output[:, 0, :]
    output = tf.keras.layers.Dense(num_labels, activation='sigmoid')(output)
    #model = tf.keras.models.Model(inputs=[input_ids, attention_mask], outputs=output)
    model = tf.keras.models.Model(inputs=[input_ids], outputs=output)
    return model

with strategy.scope():
  model = create_model(20, 'bert-base-uncased', 1)
  opt = optimizers.Adam(learning_rate=3e-5)
  loss = 'binary_crossentropy'
  model.compile(optimizer=opt, loss=loss, metrics=['accuracy'])

model.fit(tfdataset_train, batch_size=32, epochs=2)
SAVE_PATH = 'path/to/save/location'
model.save(SAVE_PATH)
Related