{"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\nThis notebook trains from the XLM-Roberta large model which was finetuned with masked language modelling on the jigsaw test dataset [Link](https://www.kaggle.com/riblidezso/finetune-xlm-roberta-on-jigsaw-test-data-with-mlm).\n\nThis notebook also implements a few improvements compared to a previous starter notebook that I shared\n\n* 1, It trains on translated data\n* 2, It uses different learning rate for the head layer and the transformer\n* 3, It restores the model weights after training to the checkpoint which had the highest validation score\n\nSuggestions/improvements are appreciated!\n\n---\n\n### References:\n\n- The shared XLM-Roberta large model, finetuned on the Jigsaw multilingual test data with masked language modelling Notebook [link]() / Dataset [link](https://www.kaggle.com/riblidezso/jigsaw-mlm-finetuned-xlm-r-large)\n- My previous starter notebook [link](https://www.kaggle.com/riblidezso/tpu-custom-tensoflow2-training-loop)\n- This notebook uses the translated versions of the training dataset too, big thanks to Michael Kazachok! [link](https://www.kaggle.com/miklgr500/jigsaw-train-multilingual-coments-google-api)\n- This notebook uses different learning rate for the transformer and the head, I got the ideas from the writeup of the winning team of the Google QUEST Q&A Labeling competition  [link](https://www.kaggle.com/c/google-quest-challenge/discussion/129840), I have seen it described to be useful elsewhere too.\n- This notebook heavily 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 = 192 \nDROPOUT = 0.5 # use aggressive dropout\nBATCH_SIZE = 16 # per TPU core\nTOTAL_STEPS_STAGE1 = 2000\nVALIDATE_EVERY_STAGE1 = 200\nTOTAL_STEPS_STAGE2 = 200\nVALIDATE_EVERY_STAGE2 = 10\n\n### Different learning rate for transformer and head ###\nLR_TRANSFORMER = 5e-6\nLR_HEAD = 1e-3\n\nPRETRAINED_TOKENIZER=  'jplu/tf-xlm-roberta-large'\nPRETRAINED_MODEL = '/kaggle/input/jigsaw-mlm-finetuned-xlm-r-large'\nD = '/kaggle/input/jigsaw-multilingual-toxic-comment-classification/'\nD_TRANS = '/kaggle/input/jigsaw-train-multilingual-coments-google-api/'\n\n\nimport os\nimport numpy as np\nimport pandas as pd\nfrom sklearn.metrics import roc_auc_score\nfrom sklearn.model_selection import train_test_split\nimport tensorflow as tf\nprint(tf.__version__)\nfrom tensorflow.keras.layers import Dense, Input, Dropout\nfrom tensorflow.keras.models import Model\nfrom tensorflow.keras.optimizers import Adam\nimport transformers\nfrom transformers import TFRobertaModel, 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 - Traning data is englih + all translations. The reason to use english too is that people use english in their foreign language comments all the time.\n - Not using the full dataset, downsampling negatives to 50-50%","execution_count":null},{"metadata":{"_uuid":"80b620f4-85cd-4436-b868-acdcea7bf2b6","_cell_guid":"1cba2ff3-145e-4187-a6a4-60d28adf5141","trusted":true},"cell_type":"code","source":"def load_jigsaw_trans(langs=['tr','it','es','ru','fr','pt'], \n                      columns=['comment_text', 'toxic']):\n    train_6langs=[]\n    for i in range(len(langs)):\n        fn = D_TRANS+'jigsaw-toxic-comment-train-google-%s-cleaned.csv'%langs[i]\n        train_6langs.append(downsample(pd.read_csv(fn)[columns]))\n\n    return train_6langs\n\ndef downsample(df):\n    \"\"\"Subsample the train dataframe to 50%-50%\"\"\"\n    ds_df= pd.concat([\n        df.query('toxic==1'),\n        df.query('toxic==0').sample(sum(df.toxic))\n    ])\n    \n    return ds_df\n    \n\ntrain_df = pd.concat(load_jigsaw_trans()) \nval_df = pd.read_csv(D+'validation.csv')\ntest_df = pd.read_csv(D+'test.csv')\nsub_df = pd.read_csv(D+'sample_submission.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\n- Note it takes some time ( approx 5 minutes)\n- Note, we need to reshape the targets","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_TOKENIZER)\nX_train = regular_encode(train_df.comment_text.values, tokenizer, maxlen=MAX_LEN)\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)\n\ny_train = train_df.toxic.values.reshape(-1,1)\ny_val = val_df.toxic.values.reshape(-1,1)","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 datasets\n\n- Note, validation dataset does not contain labels, we keep track of it ourselves","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, y_train, True)\nval_dist_dataset   = create_dist_dataset(X_val)\ntest_dist_dataset  = create_dist_dataset(X_test)","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\n\nLet's use a different learning rate for the head and the transformer like the winning team of the Google QUEST Q&A Labeling competition  [link](https://www.kaggle.com/c/google-quest-challenge/discussion/129840). \n\nThe reasoning is the following, the transformer is trained for super long time and has a very good multilingual representaton, which we only want to change a little, while the head needs to be trained from scratch.\n\nWe define 2 separate optimizers for the transofmer and the head layer. This is a simple way to use different learning rate for the transformer and the head. The caffe style \"lr_multiplier\" option would be more elegant but that is not available in keras.\n\nWe add the name 'custom' to the head layer, so that we can find it later and use a different learning rate with this layer\n\n- Note: Downloading the model takes some time!","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_model_and_optimizer():\n    with strategy.scope():\n        transformer_layer = TFRobertaModel.from_pretrained(PRETRAINED_MODEL)                \n        model = build_model(transformer_layer)\n        optimizer_transformer = Adam(learning_rate=LR_TRANSFORMER)\n        optimizer_head = Adam(learning_rate=LR_HEAD)\n    return model, optimizer_transformer, optimizer_head\n\n\ndef build_model(transformer):\n    inp = Input(shape=(MAX_LEN,), dtype=tf.int32, name=\"input_word_ids\")\n    # Huggingface transformers have multiple outputs, embeddings are the first one\n    # let's slice out the first position, the paper says its not worse than pooling\n    x = transformer(inp)[0][:, 0, :]  \n    x = Dropout(DROPOUT)(x)\n    ### note, adding the name to later identify these weights for different LR\n    out = Dense(1, activation='sigmoid', name='custom_head')(x)\n    model = Model(inputs=[inp], outputs=[out])\n    \n    return model\n\n\nmodel, optimizer_transformer, optimizer_head = create_model_and_optimizer()\nmodel.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 custom training loop\n\nWe will need:\n- 1, losses, and  optionally a training AUC metric here: these need to be defined in the scope of the distributed strategy. \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- 4, A prediction loop with dstibute \n\n\nAt the end of training we restore the parameters which had the best validation score.\n\n\nFor the different learning rate we need to apply gradients in two steps, check the train_step function for details.\n\n\n\n- Note, we are using exact AUC, for the valdationdata, and approximate AUC for the training data","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_losses_and_metrics():\n    with strategy.scope():\n        loss_object = tf.keras.losses.BinaryCrossentropy(\n            reduction=tf.keras.losses.Reduction.NONE, from_logits=False)\n\n        def compute_loss(labels, predictions):\n            per_example_loss = 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_accuracy_metric = tf.keras.metrics.AUC(name='training_AUC')\n\n    return compute_loss, train_accuracy_metric\n\n\ndef train(train_dist_dataset, val_dist_dataset=None, y_val=None,\n          total_steps=2000, validate_every=200):\n    best_weights, history = None, []\n    step = 0\n    ### Training lopp ###\n    for tensor in train_dist_dataset:\n        distributed_train_step(tensor) \n        step+=1\n\n        if (step % validate_every == 0):   \n            ### Print train metrics ###  \n            train_metric = train_accuracy_metric.result().numpy()\n            print(\"Step %d, train AUC: %.5f\" % (step, train_metric))   \n            \n            ### Test loop with exact AUC ###\n            if val_dist_dataset:\n                val_metric = roc_auc_score(y_val, predict(val_dist_dataset))\n                print(\"Step %d,   val AUC: %.5f\" %  (step,val_metric))   \n                \n                # save weights if it is the best yet\n                history.append(val_metric)\n                if history[-1] == max(history):\n                    best_weights = model.get_weights()\n\n            ### Reset (train) metrics ###\n            train_accuracy_metric.reset_states()\n            \n        if step  == total_steps:\n            break\n    \n    ### Restore best weighths ###\n    model.set_weights(best_weights)\n\n\n\n@tf.function\ndef distributed_train_step(data):\n    strategy.experimental_run_v2(train_step, args=(data,))\n\ndef train_step(inputs):\n    features, labels = inputs\n    \n    ### get transformer and head separate vars\n    # get rid of pooler head with None gradients\n    transformer_trainable_variables = [ v for v in model.trainable_variables \n                                       if (('pooler' not in v.name)  and \n                                           ('custom' not in v.name))]\n    head_trainable_variables = [ v for v in model.trainable_variables \n                                if 'custom'  in v.name]\n\n    # calculate the 2 gradients ( note persistent, and del)\n    with tf.GradientTape(persistent=True) as tape:\n        predictions = model(features, training=True)\n        loss = compute_loss(labels, predictions)\n    gradients_transformer = tape.gradient(loss, transformer_trainable_variables)\n    gradients_head = tape.gradient(loss, head_trainable_variables)\n    del tape\n        \n    ### make the 2 gradients steps\n    optimizer_transformer.apply_gradients(zip(gradients_transformer, \n                                              transformer_trainable_variables))\n    optimizer_head.apply_gradients(zip(gradients_head, \n                                       head_trainable_variables))\n\n    train_accuracy_metric.update_state(labels, predictions)\n\n\n\ndef predict(dataset):  \n    predictions = []\n    for tensor in dataset:\n        predictions.append(distributed_prediction_step(tensor))\n    ### stack replicas and batches\n    predictions = np.vstack(list(map(np.vstack,predictions)))\n    return predictions\n\n@tf.function\ndef distributed_prediction_step(data):\n    predictions = strategy.experimental_run_v2(prediction_step, args=(data,))\n    return strategy.experimental_local_results(predictions)\n\ndef prediction_step(inputs):\n    features = inputs  # note datasets used in prediction do not have labels\n    predictions = model(features, training=False)\n    return predictions\n\n\ncompute_loss, train_accuracy_metric = define_losses_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 on english comments\n\n\n- Note it takes some time\n- Don't mind the warning: \"Converting sparse IndexedSlices to a dense Tensor\"","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(train_dist_dataset, val_dist_dataset, y_val,\n      TOTAL_STEPS_STAGE1, VALIDATE_EVERY_STAGE1)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"803c2d73-e750-4189-aa53-db248fc25ea7","_cell_guid":"a280adb2-4463-4ad4-8f4a-22bc43292d82","trusted":true},"cell_type":"markdown","source":"## Finetune it on the validation data","execution_count":null},{"metadata":{"_uuid":"d5b8523a-69d6-41f0-971b-8a2bfcb9e908","_cell_guid":"1bfbb20f-409d-487f-bdf0-fca60927a868","trusted":true},"cell_type":"code","source":"%%time\n\n# decrease LR for second stage in the head\noptimizer_head.learning_rate.assign(1e-4)\n\n# split validation data into train test\nX_train, X_val, y_train, y_val = train_test_split(X_val, y_val, test_size = 0.1)\n\n# make a datasets\ntrain_dist_dataset = create_dist_dataset(X_train, y_train, training=True)\nval_dist_dataset = create_dist_dataset(X_val, y_val)\n\n# train again\ntrain(train_dist_dataset, val_dist_dataset, y_val,\n      total_steps = TOTAL_STEPS_STAGE2, \n      validate_every = VALIDATE_EVERY_STAGE2)  # not validating but printing now","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"62874aa9-54a8-495b-a2a9-ac726c6dbecf","_cell_guid":"e8be7da4-a25f-482c-88fa-c9b87a8a5f1e","trusted":true},"cell_type":"markdown","source":"## Make predictions and submission","execution_count":null},{"metadata":{"_uuid":"5cd548e4-5756-4d92-b72c-f71cb72b395d","_cell_guid":"96a7be93-0dce-4d00-92c4-d4c4cb2238f8","trusted":true},"cell_type":"code","source":"%%time\nsub_df['toxic'] = predict(test_dist_dataset)[:,0]\nsub_df.to_csv('submission.csv', index=False)","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}