{"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\nI have seen many great notebooks using Pytorch and the Keras built in training loop, but I wanted to share one which uses a custom training loop in TensorFlow 2.\n\nI hope this starter will allow more people to start experimenting with their unique ideas for tweaking.\n\nAs an example, a custom loop allows us to use the exact AUC for validation instead of the  (very convenient) [approximate value used in Keras](https://www.tensorflow.org/api_docs/python/tf/keras/metrics/AUC). The two values may differ if predictions are close to each other, and not uniformly distributed (both are happening here).\n\n** Note: This notebook tries to be simple, and only uses a small amount of data, and it does not use translated datasets or other tricks. You need to add those yourself to squeeze out a good score.**\n\nSuggestions/improvements are appreciated!\n\n---\n\n### References:\n\n\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  #Reduced for quicker execution\nLR = 1e-5\nBATCH_SIZE = 16 # per TPU core\nTOTAL_STEPS_STAGE1 = 300\nVALIDATE_EVERY_STAGE1 = 100\nTOTAL_STEPS_STAGE2 = 200\nVALIDATE_EVERY_STAGE2 = 100\n\nPRETRAINED_MODEL = 'jplu/tf-xlm-roberta-large'\nD = '/kaggle/input/jigsaw-multilingual-toxic-comment-classification/'\n\nimport os\nimport numpy as np\nimport pandas as pd\nfrom sklearn.metrics import roc_auc_score\nimport tensorflow as tf\nprint(tf.__version__)\nfrom tensorflow.keras.layers import Dense, Input\nfrom tensorflow.keras.models import Model\nimport transformers\nfrom transformers import TFAutoModel, 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","execution_count":null},{"metadata":{"_uuid":"80b620f4-85cd-4436-b868-acdcea7bf2b6","_cell_guid":"1cba2ff3-145e-4187-a6a4-60d28adf5141","trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(D+'jigsaw-toxic-comment-train.csv')\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')\n\n# subsample the train dataframe to 50%-50%\ntrain_df = pd.concat([\n    train_df.query('toxic==1'),\n    train_df.query('toxic==0').sample(sum(train_df.toxic),random_state=42)\n])\n# shufle it just to make sure\ntrain_df = train_df.sample(frac=1, random_state = 42)","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!\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_MODEL)\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- 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 = TFAutoModel.from_pretrained(PRETRAINED_MODEL)                \n        model = build_model(transformer_layer)\n        optimizer = tf.keras.optimizers.Adam(learning_rate=LR, epsilon=1e-08)\n    return model, optimizer\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    out = Dense(1, activation='sigmoid')(x)\n    model = Model(inputs=[inp], outputs=[out])\n    \n    return model\n\n\nmodel, optimizer = 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 th 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\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\n\ndef train(train_dist_dataset, val_dist_dataset=None, y_val=None,\n          total_steps=5000, validate_every=500):\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(\"     validation AUC: %.5f\" %  val_metric)   \n\n            ### Reset (train) metrics ###\n            train_accuracy_metric.reset_states()\n            \n        if step  == total_steps:\n            break\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    with tf.GradientTape() as tape:\n        predictions = model(features, training=True)\n        loss = compute_loss(labels, predictions)\n\n    gradients = tape.gradient(loss, model.trainable_variables)\n    optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n\n    train_accuracy_metric.update_state(labels, predictions)\n\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# make a new dataset for training with the validation data \n# with targets, shuffling and repeating\nval_dist_dataset_4_training = create_dist_dataset(X_val, y_val, training=True)\n\n# train again\ntrain(val_dist_dataset_4_training,\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}