{"cells":[{"metadata":{"_uuid":"7e00a1af-ce5e-41c8-bc67-b7a18bd5db75","_cell_guid":"1ac1c5d8-a575-4562-b79d-71328adf01a5","trusted":true},"cell_type":"markdown","source":"## About this notebook\n\n\nDomain specific masked language model (MLM) fine tuning was reported to improve the performance of BERT models [(arxiv/1905.05583)](https://arxiv.org/abs/1905.05583)\n\nMLM fine tuning is also regularly mentioned as one of the things winning teams have tried in previous kaggle NLP competitions (sometimes it is reported to help, sometimes not).\n\n\nHuggingface provides a pytorch script to fine tune models, however there is no TensorFlow solution which just works out of the box.\n\n\nHere I share a notebook with my implementation of MLM fine tuning on TPUs with a custom training loop in TensorFlow 2.\n\nIn my experience, starting from an MLM finetuned model makes training faster and more stable, and possibly more accurate.\n\nSuggestions/improvements are appreciated!\n\n---\n\n### References:\n\n- This notebook relies on the great [notebook]((https://www.kaggle.com/xhlulu//jigsaw-tpu-xlm-roberta) by, Xhulu: [@xhulu](https://www.kaggle.com/xhulu/) \n- The tensorflow distrubuted training tutorial: [Link](https://www.tensorflow.org/tutorials/distribute/custom_training)","execution_count":null},{"metadata":{"_uuid":"07d7bc95-0377-4db3-aa05-2bd6cf989e78","_cell_guid":"c4ad26dd-8d2a-43e9-a3ea-0a4342a16e52","trusted":true},"cell_type":"code","source":"MAX_LEN = 128\nBATCH_SIZE = 16 # per TPU core\nTOTAL_STEPS = 2000  # thats approx 4 epochs\nEVALUATE_EVERY = 200\nLR =  1e-5\n\nPRETRAINED_MODEL = 'jplu/tf-xlm-roberta-large'\nD = '/kaggle/input/jigsaw-multilingual-toxic-comment-classification/'\n\n\nimport os\nimport numpy as np\nimport pandas as pd\nimport tensorflow as tf\nprint(tf.__version__)\nfrom tensorflow.keras.optimizers import Adam\nimport transformers\nfrom transformers import TFAutoModelWithLMHead, AutoTokenizer\nimport logging\n# no extensive logging \nlogging.getLogger().setLevel(logging.NOTSET)\n\nAUTO = tf.data.experimental.AUTOTUNE","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"f0c9b5d5-e529-441d-8f37-7f64a1daccb9","_cell_guid":"dda6b430-ea9c-4772-8bd1-0938cfd6e63b","trusted":true},"cell_type":"markdown","source":"## Connect to TPU","execution_count":null},{"metadata":{"_uuid":"5a5fd45e-8f97-44d8-95f6-29cc8563dc3d","_cell_guid":"86a303dd-5d1d-4d3f-bf28-4a02b6816e9b","trusted":true},"cell_type":"code","source":"def connect_to_TPU():\n    \"\"\"Detect hardware, return appropriate distribution strategy\"\"\"\n    try:\n        # TPU detection. No parameters necessary if TPU_NAME environment variable is\n        # set: this is always the case on Kaggle.\n        tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n        print('Running on TPU ', tpu.master())\n    except ValueError:\n        tpu = None\n\n    if tpu:\n        tf.config.experimental_connect_to_cluster(tpu)\n        tf.tpu.experimental.initialize_tpu_system(tpu)\n        strategy = tf.distribute.experimental.TPUStrategy(tpu)\n    else:\n        # Default distribution strategy in Tensorflow. Works on CPU and single GPU.\n        strategy = tf.distribute.get_strategy()\n\n    global_batch_size = BATCH_SIZE * strategy.num_replicas_in_sync\n\n    return tpu, strategy, global_batch_size\n\n\ntpu, strategy, global_batch_size = connect_to_TPU()\nprint(\"REPLICAS: \", strategy.num_replicas_in_sync)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"2111bb0e-44ed-4db7-aba8-d2ff3507c454","_cell_guid":"9e748f90-73ca-4a3e-a3a5-aef269ce89c4","trusted":true},"cell_type":"markdown","source":" ## Load text data into memory\n \n - Now we will only use the multilingual test and validation data.","execution_count":null},{"metadata":{"_uuid":"80b620f4-85cd-4436-b868-acdcea7bf2b6","_cell_guid":"1cba2ff3-145e-4187-a6a4-60d28adf5141","trusted":true},"cell_type":"code","source":"val_df = pd.read_csv(D+'validation.csv')\ntest_df = pd.read_csv(D+'test.csv')","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"5e6d306f-f2ca-4452-9ace-661e6df8bae8","_cell_guid":"d070fa12-8165-41c2-9828-9690c8efa826","trusted":true},"cell_type":"markdown","source":"## Tokenize  it with the models own tokenizer\n","execution_count":null},{"metadata":{"_uuid":"7d5605a3-df77-4574-9899-89529d0a6e61","_cell_guid":"900f2d01-b87c-42e7-8cc2-8dcba95ffa54","trusted":true},"cell_type":"code","source":"%%time\n\ndef regular_encode(texts, tokenizer, maxlen=512):\n    enc_di = tokenizer.batch_encode_plus(\n        texts, \n        return_attention_masks=False, \n        return_token_type_ids=False,\n        pad_to_max_length=True,\n        max_length=maxlen\n    )\n    \n    return np.array(enc_di['input_ids'])\n    \n\ntokenizer = AutoTokenizer.from_pretrained(PRETRAINED_MODEL)\nX_val = regular_encode(val_df.comment_text.values, tokenizer, maxlen=MAX_LEN)\nX_test = regular_encode(test_df.content.values, tokenizer, maxlen=MAX_LEN)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Mask inputs and create appropriate labels for masked language modelling\n\nFollowing the BERT recipee: 15% used for prediction. 80% of that is masked. 10% is random token, 10% is just left as is.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def prepare_mlm_input_and_labels(X):\n    # 15% BERT masking\n    inp_mask = np.random.rand(*X.shape)<0.15 \n    # do not mask special tokens\n    inp_mask[X<=2] = False\n    # set targets to -1 by default, it means ignore\n    labels =  -1 * np.ones(X.shape, dtype=int)\n    # set labels for masked tokens\n    labels[inp_mask] = X[inp_mask]\n    \n    # prepare input\n    X_mlm = np.copy(X)\n    # set input to [MASK] which is the last token for the 90% of tokens\n    # this means leaving 10% unchanged\n    inp_mask_2mask = inp_mask  & (np.random.rand(*X.shape)<0.90)\n    X_mlm[inp_mask_2mask] = 250001  # mask token is the last in the dict\n\n    # set 10% to a random token\n    inp_mask_2random = inp_mask_2mask  & (np.random.rand(*X.shape) < 1/9)\n    X_mlm[inp_mask_2random] = np.random.randint(3, 250001, inp_mask_2random.sum())\n    \n    return X_mlm, labels\n\n\n# use validation and test data for mlm\nX_train_mlm = np.vstack([X_test, X_val])\n# masks and labels\nX_train_mlm, y_train_mlm = prepare_mlm_input_and_labels(X_train_mlm)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"13500d69-d19f-4ca0-8958-d4d8af5610e8","_cell_guid":"0725dd72-142e-4139-8e12-103302d53bc9","trusted":true},"cell_type":"markdown","source":"## Create distributed tensorflow dataset\n","execution_count":null},{"metadata":{"_uuid":"8b748024-314d-4256-b67b-0bc71ecac19d","_cell_guid":"757b2c93-9c07-441d-ba2e-47b8607b3795","trusted":true},"cell_type":"code","source":"def create_dist_dataset(X, y=None, training=False):\n    dataset = tf.data.Dataset.from_tensor_slices(X)\n\n    ### Add y if present ###\n    if y is not None:\n        dataset_y = tf.data.Dataset.from_tensor_slices(y)\n        dataset = tf.data.Dataset.zip((dataset, dataset_y))\n        \n    ### Repeat if training ###\n    if training:\n        dataset = dataset.shuffle(len(X)).repeat()\n\n    dataset = dataset.batch(global_batch_size).prefetch(AUTO)\n\n    ### make it distributed  ###\n    dist_dataset = strategy.experimental_distribute_dataset(dataset)\n\n    return dist_dataset\n    \n    \ntrain_dist_dataset = create_dist_dataset(X_train_mlm, y_train_mlm, True)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"7ae0b584-8ec8-4e10-928b-6c05ea2ea515","_cell_guid":"bcfe8f7d-cb29-414b-9858-b36dc57a9af1","trusted":true},"cell_type":"markdown","source":"## Build model from pretrained transformer\n\nWe just use the pretrained model loaded as a model with a language modelling head.\n\n\n\n- Note: Downloading the model takes some time!\n- Note reported number of total params, and the numbers printed in the column Param # do not match, as the weights of the mbedding matrix are shared between the encoder and the decoder.","execution_count":null},{"metadata":{"_uuid":"9f6161c2-a256-4285-9a0e-f23700f4b195","_cell_guid":"3ba0b8c0-3e4c-4afe-8766-844c6b6a8c0f","trusted":true},"cell_type":"code","source":"%%time\n\ndef create_mlm_model_and_optimizer():\n    with strategy.scope():\n        model = TFAutoModelWithLMHead.from_pretrained(PRETRAINED_MODEL)\n        optimizer = tf.keras.optimizers.Adam(learning_rate=LR)\n    return model, optimizer\n\n\nmlm_model, optimizer = create_mlm_model_and_optimizer()\nmlm_model.summary()","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"9ce1d72f-d16c-4cde-9edc-d63926623155","_cell_guid":"0e02e579-4010-4a52-b4ba-f7157b5c9f42","trusted":true},"cell_type":"markdown","source":"### Define stuff for the masked language modelling, and the custom training loop\n\nWe will need:\n- 1, MLM loss, and a loss metric. These need to be defined in the scope of the distributed strategy. The MLM loss is only evauluated at the 15% of the tokens, so we need to make a custom masked sparse categorical crossentropy. Labels with -1 value are ignored during loss calculation.\n- 2, A full training loop\n- 3, A distributed train step called in the training loop, which uses a single replica train step\n","execution_count":null},{"metadata":{"_uuid":"50c766ac-68bd-4884-ba54-657f5d385786","_cell_guid":"20b20e6a-16c4-475f-92ef-2a7809d21621","trusted":true},"cell_type":"code","source":"def define_mlm_loss_and_metrics():\n    with strategy.scope():\n        mlm_loss_object = masked_sparse_categorical_crossentropy\n\n        def compute_mlm_loss(labels, predictions):\n            per_example_loss = mlm_loss_object(labels, predictions)\n            loss = tf.nn.compute_average_loss(\n                per_example_loss, global_batch_size = global_batch_size)\n            return loss\n\n        train_mlm_loss_metric = tf.keras.metrics.Mean()\n        \n    return compute_mlm_loss, train_mlm_loss_metric\n\n\ndef masked_sparse_categorical_crossentropy(y_true, y_pred):\n    y_true_masked = tf.boolean_mask(y_true, tf.not_equal(y_true, -1))\n    y_pred_masked = tf.boolean_mask(y_pred, tf.not_equal(y_true, -1))\n    loss = tf.keras.losses.sparse_categorical_crossentropy(y_true_masked,\n                                                          y_pred_masked,\n                                                          from_logits=True)\n    return loss\n\n            \n            \ndef train_mlm(train_dist_dataset, total_steps=2000, evaluate_every=200):\n    step = 0\n    ### Training lopp ###\n    for tensor in train_dist_dataset:\n        distributed_mlm_train_step(tensor) \n        step+=1\n\n        if (step % evaluate_every == 0):   \n            ### Print train metrics ###  \n            train_metric = train_mlm_loss_metric.result().numpy()\n            print(\"Step %d, train loss: %.2f\" % (step, train_metric))     \n\n            ### Reset  metrics ###\n            train_mlm_loss_metric.reset_states()\n            \n        if step  == total_steps:\n            break\n\n\n@tf.function\ndef distributed_mlm_train_step(data):\n    strategy.experimental_run_v2(mlm_train_step, args=(data,))\n\n\n@tf.function\ndef mlm_train_step(inputs):\n    features, labels = inputs\n\n    with tf.GradientTape() as tape:\n        predictions = mlm_model(features, training=True)[0]\n        loss = compute_mlm_loss(labels, predictions)\n\n    gradients = tape.gradient(loss, mlm_model.trainable_variables)\n    optimizer.apply_gradients(zip(gradients, mlm_model.trainable_variables))\n\n    train_mlm_loss_metric.update_state(loss)\n    \n\ncompute_mlm_loss, train_mlm_loss_metric = define_mlm_loss_and_metrics()","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"e2997f2d-ae14-4bd3-b109-82e6700abb33","_cell_guid":"f7958337-2407-4734-98a0-51c89f4f9dc6","trusted":true},"cell_type":"markdown","source":"## Finally train it\n\n\n- Note it takes some time\n","execution_count":null},{"metadata":{"_uuid":"b30a95a7-ae7e-40e8-846b-9d31bfe997ad","_cell_guid":"d3d0e97f-d7d4-4d3f-8c6e-925669565248","trusted":true},"cell_type":"code","source":"%%time\ntrain_mlm(train_dist_dataset, TOTAL_STEPS, EVALUATE_EVERY)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"62874aa9-54a8-495b-a2a9-ac726c6dbecf","_cell_guid":"e8be7da4-a25f-482c-88fa-c9b87a8a5f1e","trusted":true},"cell_type":"markdown","source":"## Save finetuned model\n\nThis fine tuned model can be loaded just as the original to build a classification model from it.","execution_count":null},{"metadata":{"_uuid":"5cd548e4-5756-4d92-b72c-f71cb72b395d","_cell_guid":"96a7be93-0dce-4d00-92c4-d4c4cb2238f8","trusted":true},"cell_type":"code","source":"mlm_model.save_pretrained('./')","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}