modify mlm example to support all backends

[Mrutyunjay01]

With reference to the draft PR, here's an attempt to make to example truly backend-agnostic. In the process, I raised the following issues: keras-team/keras#18410, and keras-team/keras#19665. While the latter got resolved, the former persists still. Kindly review the changes, and let me know how shall we proceed with the issues mentioned.

cc: @fchollet

[Mrutyunjay01]

Here, shall I subclass the MaskedLanguageModel instead, and get rid of the tf custom training step?

[fchollet]

I believe you should be able to remove the train_step() and instead implement the compute_loss() method. If you write the compute_loss() method using only keras.ops functions, it will work with all backends.

[Mrutyunjay01]

Given that no additional logic is required here, and using SparseCategoricalCrossEntropy as Loss function; does it require an override for compute_loss()?

[fchollet]

Thanks for the PR!

[fchollet]

If it works with all backends, you can just remove this line.

[fchollet]

I believe you should be able to remove the train_step() and instead implement the compute_loss() method. If you write the compute_loss() method using only keras.ops functions, it will work with all backends.

[Mrutyunjay01]

added subclassing implementation. Certain obstacles remain as mentioned:

1. In BertEncoderLayer class, as it is inherited from layers.Layer, is it possible to access to the sub-layers here? as:

mlm_model = models.load_model("bert_mlm.keras", custom_objects={"MaskedLanguageModel": MaskedLanguageModel})

#mlm_model.layers : 
#[<Embedding name=word_embedding, built=True>,
#<PositionEmbedding name=position_embedding, built=True>,
#<BertEncoderLayer name=bert_encoder_layer, built=True>,
#<Dense name=mlm_cls, built=True>]

mlm_model.get_layer("bert_encoder_layer").get_layer("encoder_0_ffn_layernormalization")

This throws error saying

AttributeError: 'BertEncoderLayer' object has no attribute 'get_layer'

2. The reason I want to do the above is to reproduce the following:

# Load pretrained bert model
mlm_model = models.load_model(
    "bert_mlm_imdb.keras", custom_objects={"MaskedLanguageBertModel": MaskedLanguageBertModel}
)
pretrained_bert_model = Model(
    mlm_model.input, mlm_model.get_layer("encoder_0_ffn_layernormalization").output
)

This was possible in functional declaration of the Model, but not in sub-classing (in last commit).

Kindly suggest w.r.t last commit

Thanks!