如何在 MLM 任务上训练 Tensorflow 的预训练 BERT？ （仅在 Tensorflow 中使用预训练的 model）

Question

. Answers to this question are eligible for a +50 reputation bounty. Deshwal is looking for a :

There is no official documentation I found which says how can you train MLM and then fine tune as Classification model in Tensorflow. One way could be training MLM as Pytorch model and loading the same as Tensorflow model. Smallest help is appreciated on the topic. Be it training in Pytorch and loading in Tensorflow or Doing everything in Tensorflow.

Before Anyone suggests pytorch and other things, I am looking specifically for Tensorflow + pretrained + MLM task only. I know, there are lots of blogs for PyTorch and lots of blogs for fine tuning ( Classification) on Tensorflow.

Coming to the problem, I got a language model which is English + LaTex where a text data can represent any text from Physics, Chemistry, MAths and Biology and any typical example can look something like this: Link to OCR image

"Find the value of function x in the equation: \n \\( f(x)=\\left\\{\\begin{array}{ll}x^{2} & \\text { if } x<0 \\\\ 2 x & \\text { if } x \\geq 0\\end{array}\\right. \\)"

So my language model needs to understand \geq \\begin array \eng \left \right other than the English language and that is why I need to train an MLM first on pre-trained BERT or SciBERT to have both. So I went up digging the internet and found some tutorials:

1. MLM training on Tensorflow BUT from Scratch; I need pre-trained
2. MLM on pre-trained but in Pytorch ; I need Tensorflow
3. Fine Tuning with Keras; It is for classification but I want MLM

I already have a fine tuning classification model. Some of the code is as follows:

tokenizer = transformers.BertTokenizer.from_pretrained('bert-large-uncased')

def regular_encode(texts, tokenizer, maxlen=maxlen):
  enc_di = tokenizer.batch_encode_plus(texts,  return_token_type_ids=False,padding='max_length',max_length=maxlen,truncation = True,)
  return np.array(enc_di['input_ids'])

Xtrain_encoded = regular_encode(X_train.astype('str'), tokenizer, maxlen=maxlen)
ytrain_encoded = tf.keras.utils.to_categorical(y_train, num_classes=classes,dtype = 'int32')

def build_model(transformer, loss='categorical_crossentropy', max_len=maxlen, dense = 512, drop1 = 0.3, drop2 = 0.3):
    input_word_ids = tf.keras.layers.Input(shape=(max_len,), dtype=tf.int32, name="input_word_ids")
    sequence_output = transformer(input_word_ids)[0]
    cls_token = sequence_output[:, 0, :]

    #Fine Tuning Model Start
    x = tf.keras.layers.Dropout(drop1)(cls_token)  
    x = tf.keras.layers.Dense(512, activation='relu')(x)
    x = tf.keras.layers.Dropout(drop2)(x)
    out = tf.keras.layers.Dense(classes, activation='softmax')(x)
    model = tf.keras.Model(inputs=input_word_ids, outputs=out)
    return model

Only useful thing I could get was in the HuggingFace that

With the tight interoperability between TensorFlow and PyTorch models, you can even save the model and then reload it as a PyTorch model (or vice-versa)

from transformers import AutoModelForSequenceClassification

model.save_pretrained("my_imdb_model")
pytorch_model = AutoModelForSequenceClassification.from_pretrained("my_imdb_model", from_tf=True)

So maybe I could train a pytorch MLM and then load it as a tensorflow fine tuned classification model? Is there any other way?

Answer 1

I also encountered this problem, and did some research. A possible solution is to use tensorflow_hub models supporting mlm.

For example, https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-10_H-128_A-2/2 this model has a relatively complete tutorial on how to create a mlm task. But there is a problem if you are designing multitasks. Because:

import tensorflow_hub as hub

encoder = hub.load(model_path)
mlm = hub.KerasLayer(encoder.mlm, trainable=True)

I'm not very sure this can help in multitasks because I dont know whether encoder is updated. Hope this could be useful and maybe someone could also help me out. Thanks.

Some Updates: I have checked the official document for tfhub, and I guess I have mistaken what the encoder means in this model. The weights are updated in the encoder and this actually can serve multitask pretraining, will update if any result goes wrong in my experiment.

Answer 2

I think after some research, I found something. I don't know if it'll work but it uses transformer and Tensorflow with XLM . I think it'll work for BERT too.

PRETRAINED_MODEL = 'jplu/tf-xlm-roberta-large'
from tensorflow.keras.optimizers import Adam
import transformers
from transformers import TFAutoModelWithLMHead, AutoTokenizer

def create_mlm_model_and_optimizer():
    with strategy.scope():
        model = TFAutoModelWithLMHead.from_pretrained(PRETRAINED_MODEL)
        optimizer = tf.keras.optimizers.Adam(learning_rate=LR)
    return model, optimizer


mlm_model, optimizer = create_mlm_model_and_optimizer()




def define_mlm_loss_and_metrics():
    with strategy.scope():
        mlm_loss_object = masked_sparse_categorical_crossentropy

        def compute_mlm_loss(labels, predictions):
            per_example_loss = mlm_loss_object(labels, predictions)
            loss = tf.nn.compute_average_loss(
                per_example_loss, global_batch_size = global_batch_size)
            return loss

        train_mlm_loss_metric = tf.keras.metrics.Mean()
        
    return compute_mlm_loss, train_mlm_loss_metric


def masked_sparse_categorical_crossentropy(y_true, y_pred):
    y_true_masked = tf.boolean_mask(y_true, tf.not_equal(y_true, -1))
    y_pred_masked = tf.boolean_mask(y_pred, tf.not_equal(y_true, -1))
    loss = tf.keras.losses.sparse_categorical_crossentropy(y_true_masked,
                                                          y_pred_masked,
                                                          from_logits=True)
    return loss

            
            
def train_mlm(train_dist_dataset, total_steps=2000, evaluate_every=200):
    step = 0
    ### Training lopp ###
    for tensor in train_dist_dataset:
        distributed_mlm_train_step(tensor) 
        step+=1

        if (step % evaluate_every == 0):   
            ### Print train metrics ###  
            train_metric = train_mlm_loss_metric.result().numpy()
            print("Step %d, train loss: %.2f" % (step, train_metric))     

            ### Reset  metrics ###
            train_mlm_loss_metric.reset_states()
            
        if step  == total_steps:
            break


@tf.function
def distributed_mlm_train_step(data):
    strategy.experimental_run_v2(mlm_train_step, args=(data,))


@tf.function
def mlm_train_step(inputs):
    features, labels = inputs

    with tf.GradientTape() as tape:
        predictions = mlm_model(features, training=True)[0]
        loss = compute_mlm_loss(labels, predictions)

    gradients = tape.gradient(loss, mlm_model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, mlm_model.trainable_variables))

    train_mlm_loss_metric.update_state(loss)
    

compute_mlm_loss, train_mlm_loss_metric = define_mlm_loss_and_metrics()

Now train it as train_mlm(train_dist_dataset, TOTAL_STEPS, EVALUATE_EVERY)

Above ode is from this notebook and you need to do all the necessary things given exactly

The author says in the end that:

This fine tuned model can be loaded just as the original to build a classification model from it