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](https://github.com/huggingface/transformers/files/5981406/compile_model.txt)



Hello @saboof!

First of all, which version of Transformers and TF are you using? And can you share with us a Colab from which we can reproduce the issue. Thanks!!<|||||>version of transformers: transformers==4.2.2
version of TF2: tensorflow==2.3.1

after editing, the code is attached to the main issue <|||||>Ok, Thanks for sharing! I now see better what you are trying to do. The results you get are different mostly because how the models are implemented don't support what you are trying to do and then some unexpected behavior might occur.

Usually, a model should not belong to another model, and using a tokenizer  inside a model is currently not recommended because not stable. Sorry, for your use case you will have to use `TFBertModel` and tokenize/pad your documents outside the model.

It is in our plans to give the possibility to integrate the tokenization process directly inside the model, and then be also part of a SavedModel, but we don't have an ETA yet. Sorry :(<|||||>hi @jplu 

Thanks for the quick reply :), I don't think your comment answers my question.
Even when I take the tokenizer outside the bert class model, we still get the same error: 

AttributeError: Must set `config_class` to use @keras_serializable<|||||>Sorry I should have been much clearrer. More specifically, you cannot use `TFBertMainLayer` publicly. You have to directly use `TFBertModel`.

Also, in case you don't know, TF doesn't recommend to use a model inside another model because some unstable behaviour might occur.<|||||>Thanks,

But I don't use `TFBertMainLayer` at all, the opposite, I use `TFBertModel`. 

Regarding your second point I'll try to take the BERT model outside the general model and only pass it as an argument. 
Generally -  how would you go about creating additional linear layer on top of BERT CLS embedding? 


<|||||>> Regarding your second point I'll try to take the BERT model outside the general model and only pass it as an argument.

This might bring unexpected behaviour as well.

> Generally - how would you go about creating additional linear layer on top of BERT CLS embedding?

I suggest you to take a look at how the `TFBertForSequenceClassification` is built :)<|||||>Great, I'll take a look at this, thanks @jplu <|||||>Looking at the code of `TFBertForSequenceClassification` -  
https://github.com/huggingface/transformers/blob/d1b14c9b548de34b6606946482946008622967db/src/transformers/models/bert/modeling_bert.py#L1496 
It has a linear layer and the BERT model within the same model, (although written in Pytorch). This doesn't give me insights as to how to combine these two models/layers in the same Keras model so we would be able to save them both in the same command. 

The only option that I see now is to completely adopt the `TFBertForSequenceClassification` as my full model. Unfortunately, This doesn't give the dynamic architecture generation which I would like to achieve.   <|||||>What do you mean by dynamic architecture? Can you detail a bit more about what you would like to do?<|||||>For example, having two types of outputs based on the same CLS embedding. 

The first type would be based on a single linear layer and the second type would have two layers with some activation function between the layers.

sketch of the model.

Input -> BERT CLS embedded vector -> output based on a single layer -> another output based on a second layer

More generally, we would like to get the CLS embedded vector and to add on top of that whatever other layers. 
These layers should be part of the Keras model so they all will be saved as part of the model.

I hope this is somehow clearer now. 
<|||||>Ok this is much clearer now, thanks!!

The best way to do that would be to build your model as we do for ours, it means something like:

```python
class MyModel(TFBertPreTrainedModel):
  def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)

        self.bert = TFBertMainLayer(config, name="bert")
        self.my_first_layer = .....
        self.my_second_layer = ....

  def call(inputs, training=None):
        bert_output = bert(inputs)
        # below all the processes that you have to implement for your models
        ....
```

Once you have your model you can instantiate it like our models with ```model = MyModel.from_pretrained("model_path")```<|||||>Hi @jplu 

Again, thanks for the help. 

Sorry, I think we are back in square one. 
My initial issue was how to save a Keras model which includes a HF model within it?

Attached is an example of a code that generates a Keras model which incorporates a BERT model.   
[example.txt](https://github.com/huggingface/transformers/files/6015722/example.txt)

How would you go about saving this model in such a way that I'd be able to load the model with the additional layer on top of it? can I save the tokenizer in the same place? 





<|||||>> Sorry, I think we are back in square one. My initial issue was how to save a Keras model which includes a HF model within it?

Once again, you should not integrate a model inside another model, What I proposed you to do is exactly what you would like to do because `TFBertMainLayer == TFBertModel`.

> How would you go about saving this model in such a way that I'd be able to load the model with the additional layer on top of it? can I save the tokenizer in the same place?

I would go to what I proposed :)<|||||>Thanks @jplu 

For the best of my understanding, your proposed code doesn't instantiate a Keres model. right?
If I get your suggestions correctly what you say is that I can't wrap HF model inside a Keras model since integrating a model inside another model, is discouraged by Keras. 

<|||||>Yes! I proposed you to use such a way to do that is 100% equivalent to what you would like.<|||||>Thanks! <|||||>Keras discouraging nested models is new information for me. 
And this is weird as in their Keras official examples they do exactly that.
They have a nested HF model inside a Keras model.

https://keras.io/examples/nlp/text_extraction_with_bert/#preprocess-the-data
Please look under the function create_model(). 

Do you have any pointers in which Keras discouraging nested models is mentioned? <|||||>What you point out is a totally different way to what you are doing. In the link you proposed, the model is built in a functional manner which is totally different of building a model in a subclassing manner.

Please use the TFBertMainLayer, this is the exact same thing.<|||||>Sorry, I can't find the difference between the Keras example and my previous example under the create_model function.
Attached is the function I used in my example file.

```
def create_model(bert_variant, max_len):
    ## BERT encoder
    encoder = TFBertModel.from_pretrained(bert_variant)

    ## findings classifer Model
    input_word_ids = tf.keras.layers.Input(shape=(max_len,), dtype=tf.int32,
                                           name="input_word_ids")
    embedding = encoder([input_word_ids]).last_hidden_state
    cls_layer = embedding[:,0,:]

    logits = layers.Dense(4, name="prediction", use_bias=True)(cls_layer)
    logits = layers.Reshape((1, 4))(logits)

    pred_probs = 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

```
 
<|||||>Sorry, this way of creating a model is fine, I was referencing to your first example script called `compile_model.txt`.

Now to properly save the model you can do something like this:
```python
def create_model(bert_variant, max_len):
    ## BERT encoder
    encoder = TFBertModel.from_pretrained(bert_variant).bert # ====> look at this change here :)

    ## findings classifer Model
    input_word_ids = tf.keras.layers.Input(shape=(max_len,), dtype=tf.int32,
                                           name="input_word_ids")
    embedding = encoder([input_word_ids]).last_hidden_state
    cls_layer = embedding[:,0,:]

    logits = layers.Dense(4, name="prediction", use_bias=True)(cls_layer)
    logits = layers.Reshape((1, 4))(logits)

    pred_probs = 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")

loaded_model = tf.keras.models.load_model("model_path")
```<|||||>Or a much nicer version IMO:
```python
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:
```python
model = tf.keras.models.load_model("model_path")
l = [1]*128 
inp = tf.constant([l])
model(inp)
```<|||||>Thanks, @jplu <|||||>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/master/CONTRIBUTING.md) are likely to be ignored.