Saving HF wrapped in Keras

Hi,

Trying to save a Keras model which has a HF model and a liner layer (Dense layer) on top of it. To save a model, Keras requires that every layer would have a serialize_layer_fn implemented. However it seems that HF models don’t include this function.

Spending some time understanding and googling this issue I came across the @keras_serializable which I assume is supposed to allow a HF layer to be serialized. However when I wrap my class with this, I get this issue:

AttributeError: Must set config_class to use @keras_serializable

Although the attached code which describes this issue does have a config_class member.

To test this behaviour I generated a code which creates a model based on BERT or some Bionlp-BERT variant. Then the code tries to train the model. After the short training period, we try to save the model, which includes the BERT model, BERT tokenizer (using save_pretrained) and the liner layer.

For me, the best way to save this model would be using the to_json function, which converts the Keras model into a serialized model. However, I couldn’t manage to get this working. Any idea how can this be done? Or any objection about saving HF model using the “to_json” method?

A possible workaround would be to use the save_pretrained method for the BERT model and tokenizer but then how would I save also the linear layer?

This issue relates to this issue https://github.com/huggingface/transformers/issues/2733.

compile_model.txt

[jplu]

Or a much nicer version IMO:

def create_model(bert_variant, max_len):
    config = BertConfig.from_pretrained(bert_variant)
    input_word_ids = tf.keras.layers.Input(shape=(max_len,), dtype=tf.int32,
                                           name="input_word_ids")
    embedding = TFBertMainLayer(config)(input_word_ids)
    cls_layer = embedding.last_hidden_state[:,0,:]
    logits = tf.keras.layers.Dense(4, name="prediction", use_bias=True)(cls_layer)
    logits = tf.keras.layers.Reshape((1, 4))(logits)
    pred_probs = tf.keras.layers.Activation(tf.keras.activations.softmax)(logits)
    model = tf.keras.Model(inputs=[input_word_ids], outputs=[pred_probs])
    loss = tf.keras.losses.CategoricalCrossentropy(from_logits=False)
    optimizer = tf.keras.optimizers.Adam(lr=5e-5)
    model.compile(optimizer=optimizer, loss=loss)

    return model

model = create_model("bert-base-cased", 128)
model.save("model_path")  

Once saved you can load and use it like this:

model = tf.keras.models.load_model("model_path")
l = [1]*128 
inp = tf.constant([l])
model(inp)

[github-actions[bot]]

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 are likely to be ignored.