You should be able to achieve the same with just SparseCategoricalCrossentropy by using the sample_weight argument. https://keras.io/api/losses/probabilistic_losses/#sparsecategoricalcrossentropy-class

I have tried using SparseCategoricalCrossentropy sample_weight argument but it's giving nan loss (because I have used -1 (negative/ unmasked token) in labels tensor). Should I use non-negative id for unmasked token and then use sample_weight?

Instead, please use the train_step overriding workflow. This will be much simpler and won't need to rely on experimental_run_v2. https://keras.io/guides/customizing_what_happens_in_fit/

I have implemented this using train_step overriding workflow.

Dataset is already a base class in TF. I would recommend simply refactoring this as separate utility functions, since there is no need for an OO approach here.

Ok, I have implemented it. Please review functions

You should define this in __init__ (or compile)

A problem here is that you never reset_states() on the loss tracker. So what you display is an average of batch losses since the start of training.

You should create the property metrics:

    @property
    def metrics(self):
        # We list our `Metric` objects here so that `reset_states()` can be
        # called automatically at the start of each epoch
        # or at the start of `evaluate()`.
        return [loss_tracker]

I think it's fine to keep it implemented like this -- thanks for investigating the sample_weights arg

A concern I have here: this example is basically more about "how to use HuggingFace Transformers" than about masked language modeling or Keras. The pipeline, tokenizer, and TFAutoModelWithLMHead are black boxes.

The purpose of our code examples is to be Keras-centric and as didactic as possible, i.e. they should explain the task and problem domain, similar to what we do here for object detection, where everything is implemented from scratch.

Would you consider implementing these component from scratch in Keras, e.g. with the TextVectorization layer for tokenization and with a Transformer block built with the new MultiHeadAttention layer (in tf-nightly)? No pretraining necessary.

The model could be similar to this, expect you wouldn't even have to implement MultiHeadAttention.

Hi, @fchollet Thanks for your comment.

True, I feel the same, this example looks like 'how to use HuggingFace Transformers'.

Yes, I would like to consider implementing these components from scratch in Keras, e.g. with the TextVectorization layer for tokenization and with a Transformer block built with the new MultiHeadAttention layer.

Should I use this for the classification Problems or Pretraining Problems?

I have started according to your suggestion to use MultiHead layer and TextVectorizationLayer. And it's almost done for both pretraning and fine tuning.

I have one doubt, TextVectorizationLayer is not running on TPU. Same notebook running on CPU n GPU, but not on TPU

Are you using TextVectorization as part of the tf.data pipeline? Can you share your notebook?

No, I am not using tf.data.

def get_vectorize_layer(texts, vocab_size, max_seq, special_tokens=['[MASK]']):

  vectorize_layer = preprocessing.TextVectorization(
    max_tokens=vocab_size,
    output_mode='int',
    standardize = None,
    output_sequence_length=max_seq)
  
  vectorize_layer.adapt(texts) # error in adapt

  vocab = vectorize_layer.get_vocabulary()
  vocab = vocab[2:vocab_size-len(special_tokens)] + ['[MASK]']
  vectorize_layer.set_vocabulary(vocab)
  return vectorize_layer


vectorize_layer = get_vectorize_layer(data.text.values.tolist(), 20000, 196, special_tokens=['[MASK]'])

Error is:

NotFoundError: 'OptimizeDatasetV2' is neither a type of a primitive operation nor a name of a function registered in binary running on n-86c78cbc-w-0. Make sure the operation or function is registered in the binary running in this process.

Tensorflow: 2.4.0-dev20200915
Reference Notebook: https://colab.research.google.com/drive/1QiQmA2M2WkwsjiHHxBhPM0Jw3_LBJnGV?usp=sharing

Note: Same code working with tensorflow:2.3.0

This is a bug likely to a recently introduced optimization setting for tf.data datasets (which are used internally in adapt). We're looking into it -- thank you for the report.

Since we will fix this before the next release, I recommend that you keep developing this example with TextVectorization, using GPU only for now. You could try to use an older tf-nightly version that wouldn't include the bug (though I am not sure if you'll find one that also has the MultiHeadAttention layer).

I am developing with GPU compatible. I have tried almost random 5-7 different tf-nightly version. All have same issues.

My notebook is almost ready. going to commit in a few mins.
Thanks!

In practice, this can just be mask_token='[MASK]'.

If you put it as the start of the list you won't need to compute mask_token_id by hand (it will be 0)

I am a little bit confused. My understanding is index 0 is reserved for padding token. (Reference)

And I feel there is no argument called mask_token in the class constructor of TextVectorization (Reference)

I need to insert special tokens ('[MASK] is this case') into the vocabulary and I would like to use zero index for padding.

Am I thinking right?

Got it. The mask token is normally index 0. So you could just insert it at index 0, which is what I was saying. Note that there is a mask token already, by default it is the empty string. You could simply replace [MASK] with the empty string in your code and you wouldn't need to configure it (and its index would be 0).

Note, the purpose of the default mask token is to handle padding, so your data will already contain 0s for sections that are padded (sequences shorter than MAX_LEN). Not sure if that forces you to use a separate mask token...

I agree with you that it should be configurable in the constructor of TextVectorization. We will add it (alongside a couple other arguments). Internally TextVectorization uses a StringLookup layer which allows this customization. However if we configure it in this way it would still be used for padding as well.

You should be able to remove this statement

ok I will refactor.

You should be able to remove this statement

Still nit: use

from tensorflow.keras import layers

then layers.Embedding, etc.

ok I will refactor in the same way.

Any call to fit() should include a way to monitor results on an evaluation set.

Why only train for 1 epoch?

Note: since you're training using a generator, consider using multiple workers for better performance.

I used one epoch because of time. It is taking 22min per epoch. I will use multiple workers for better performance and will do for 3 epochs.

Actually I am not able to install (running without gpu) tf-nightly-gpu on colab or kaggle. I tried to install but I unable to get GPU support.

Why no reduction? since you're using compile/fit the results will be averaged for reduction anyway (since some form of reduction is necessary).

Style nit: variable names should be lowercase. Please always use fully-spelled out variable names for better readability, instead of abbreviations or single letters.

Prefer using include_labels=True which is more explicit

Is there any way to handle this using tf.data.Dataset instead?

Ok I will try to implement in tf.data.Dataset

Got it. The mask token is normally index 0. So you could just insert it at index 0, which is what I was saying. Note that there is a mask token already, by default it is the empty string. You could simply replace [MASK] with the empty string in your code and you wouldn't need to configure it (and its index would be 0).

Note, the purpose of the default mask token is to handle padding, so your data will already contain 0s for sections that are padded (sequences shorter than MAX_LEN). Not sure if that forces you to use a separate mask token...

I agree with you that it should be configurable in the constructor of TextVectorization. We will add it (alongside a couple other arguments). Internally TextVectorization uses a StringLookup layer which allows this customization. However if we configure it in this way it would still be used for padding as well.

For each subsection, please add a text paragraph explaining what you're doing:

• How do you prepare the dataset? what do your samples end up looking like?
• What does the masked language model do (what are its inputs are outputs)?
• What does the classification model do?
• What does the end to end model do?
• How to interpret evaluation metrics?

Prefer obtaining the data from its original source, e.g. https://keras.io/examples/nlp/text_classification_from_scratch/#load-the-data-imdb-movie-review-sentiment-classification

Please expand this description, provide concrete examples

This would be far more readable as a Sequential or a Functional model. It's just a list of layers applied in sequence, but this class makes it a puzzle to decode.

You could specify this in 8-9 lines in the Functional API.

def bert_module(query, key, value):
   x = layers.MultiHeadAttention(...)(query, key, value)
   x = layers.Dropout(...)(x)
   x = layers.LayerNormalization(epsilon=1e-6)(x)
   x  = layers.Dense(...)(x)
   x  = layers.Dense(...)(x)
   x = layers.Dropout(...)(x)
   x = layers.LayerNormalization(epsilon=1e-6)(x)
   return x

Agreed

You can create this in the constructor. It should be managed by the model.

I had included it in the constructor but I was facing an issue in saving (.save('model.h5)) because it included the Mean Layer(loss tracker mean metric).

No need for the print statement here.

Why doesn't predict already return a softmax?

Prefer doing everything in NumPy, for consistency

Please use string formatting

Use lowercase variable name