{"cells":[{"metadata":{},"cell_type":"markdown","source":"https://www.kaggle.com/keremt/xlm-roberta-tpu-training-fastai-style","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import os\n\nimport numpy as np\nimport pandas as pd\nimport tensorflow as tf\nfrom tensorflow.keras.optimizers import Adam\nfrom tensorflow.keras.models import Model\nfrom tensorflow.keras.callbacks import ModelCheckpoint\nfrom kaggle_datasets import KaggleDatasets\nimport transformers\nfrom transformers import TFAutoModel, AutoTokenizer, AutoConfig\nfrom tqdm.notebook import tqdm\nfrom tokenizers import Tokenizer, models, pre_tokenizers, decoders, processors\n\nfrom fastai.text import *\nfrom tensorflow.keras.layers import Dense, Input, Dropout, Flatten\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Detect hardware, return appropriate distribution strategy\ntry:\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())\nexcept ValueError:\n    tpu = None\n\nif tpu:\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    strategy = tf.distribute.experimental.TPUStrategy(tpu)\nelse:\n    # Default distribution strategy in Tensorflow. Works on CPU and single GPU.\n    strategy = tf.distribute.get_strategy()\n\nprint(\"REPLICAS: \", strategy.num_replicas_in_sync)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"AUTO = tf.data.experimental.AUTOTUNE\n\n# Data access\n# GCS_DS_PATH = KaggleDatasets().get_gcs_path('')\n\n# Configuration\nEPOCHS = 2\nBATCH_SIZE = 16 * strategy.num_replicas_in_sync","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Create Datasets","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"from catalyst.data.sampler import DistributedSamplerWrapper, BalanceClassSampler","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"XLM_PROCESSED_PATH = Path(\"/kaggle/input/xlmrobertabase/xlm_roberta_processed/\"); XLM_PROCESSED_PATH.ls()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# def one_hot_encode(x): return np.array([[1,0], [0,1]])[x]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class JigsawArrayDataset(torch.utils.data.Dataset):\n    def __init__(self, input_ids:np.array, attention_mask:np.array, toxic:np.array=None):\n        self.input_ids = input_ids\n        self.attention_mask = attention_mask\n        self.toxic = toxic\n    \n    def __getitem__(self, idx):\n        xb = (tensor(self.input_ids[idx]), tensor(self.attention_mask[idx]))\n        yb = tensor(0.) if self.toxic is None else tensor(self.toxic[idx])\n        return xb,yb    \n        \n    def __len__(self):\n        return len(self.input_ids)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# train_ds\ntrain_input_ids = np.load(XLM_PROCESSED_PATH/'translated_train_inputs/input_ids.npy')\ntrain_attetion_mask = np.load(XLM_PROCESSED_PATH/'translated_train_inputs/attention_mask.npy').astype(np.int32)\ntrain_toxic = np.load(XLM_PROCESSED_PATH/'translated_train_inputs/toxic.npy').astype(np.float32)\ntrain_lang = np.load(XLM_PROCESSED_PATH/'translated_train_inputs/lang.npy', allow_pickle=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# labels for stratified batch sampler\ntrain_stratify_labels = array(s1+s2 for (s1,s2) in zip(train_lang, train_toxic.astype(str)))\nlabels2int = {v:k for k,v in enumerate(np.unique(train_stratify_labels))}\nlabels = [labels2int[o] for o in train_stratify_labels]\nbalanced_sampler = BalanceClassSampler(labels)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"idxs1 = list(iter(balanced_sampler))\nidxs2 = list(iter(balanced_sampler))\nidxs3 = list(iter(balanced_sampler))\nidxs4 = list(iter(balanced_sampler))\nidxs5 = list(iter(balanced_sampler))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# reshape for loss and metric\ntrain_toxic = train_toxic.reshape(-1,1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_input_ids_sampled = np.vstack([train_input_ids[idxs] for idxs in [idxs1,idxs2,idxs3,idxs4,idxs5]])\ntrain_attetion_mask_sampled = np.vstack([train_attetion_mask[idxs] for idxs in [idxs1,idxs2,idxs3,idxs4,idxs5]])\ntrain_toxic_sampled = np.vstack([train_toxic[idxs] for idxs in [idxs1,idxs2,idxs3,idxs4,idxs5]])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_input_ids_sampled.shape, train_attetion_mask_sampled.shape, train_toxic_sampled.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"del train_input_ids, train_attetion_mask, train_toxic, train_lang\ngc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = (\n    tf.data.Dataset\n    .from_tensor_slices(((train_input_ids_sampled, train_attetion_mask_sampled), train_toxic_sampled))\n    .batch(BATCH_SIZE)\n    .prefetch(AUTO)\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# def generate(self):\n#     for idx in idxs1:\n#         yield ((train_input_ids[idx], train_attetion_mask[idx]), train_toxic[idx])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# train_dataset = (\n#     tf.data.Dataset\n#     .from_generator(generate, \n#                     ((tf.int32, tf.int32), tf.float32),\n#                     ((tf.TensorShape([256]), tf.TensorShape([256])), tf.TensorShape([1]))\n#                    )\n#     .repeat()\n#     .shuffle(2048)\n#     .batch(BATCH_SIZE)\n#     .prefetch(AUTO)\n# )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"valid_input_ids = np.load(XLM_PROCESSED_PATH/'valid_inputs/input_ids.npy')\nvalid_attention_mask = np.load(XLM_PROCESSED_PATH/'valid_inputs/attention_mask.npy').astype(np.int32)\nvalid_toxic = np.load(XLM_PROCESSED_PATH/'valid_inputs/toxic.npy').astype(np.float32).reshape(-1,1)\n# valid_toxic = one_hot_encode(valid_toxic).astype(np.float32)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"valid_dataset = (\n    tf.data.Dataset\n    .from_tensor_slices(((valid_input_ids, valid_attention_mask), valid_toxic))\n    .batch(BATCH_SIZE)\n    .cache()\n    .prefetch(AUTO)\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Load Model","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"MODEL = 'jplu/tf-xlm-roberta-large'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_xlm_roberta(modelname=MODEL):        \n    conf = AutoConfig.from_pretrained(modelname)\n    conf.output_hidden_states = True\n    model = TFAutoModel.from_pretrained(modelname, config=conf)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def build_model(xlm_roberta, max_len=256, p=0.5):\n    \"\"\"\n    https://www.kaggle.com/xhlulu/jigsaw-tpu-distilbert-with-huggingface-and-keras\n    \"\"\"\n    input_ids = Input(shape=(max_len,), dtype=tf.int32, name=\"input_ids\")\n    attention_mask = Input(shape=(max_len,), dtype=tf.int32, name=\"attention_mask\")\n    \n    _, _, hidden_states = xlm_roberta([input_ids, attention_mask])\n    x = tf.concat(hidden_states[-2:], -1)\n    x = tf.concat((tf.reduce_mean(x, 1), tf.reduce_max(x, 1)), -1)    \n    x = Dropout(rate=0.5)(x)\n    out = Dense(1, activation=\"sigmoid\")(x)\n    \n    model = Model(inputs=(input_ids, attention_mask), outputs=out)\n    \n    loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=False, label_smoothing=0.1, name='binary_crossentropy')\n    \n    model.compile(Adam(lr=1e-5), loss=loss_fn, metrics=[tf.metrics.AUC()])\n    \n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"with strategy.scope():\n    xlm_roberta = get_xlm_roberta()\n    model = build_model(xlm_roberta)\nmodel.summary()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Training","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"BATCH_SIZE, EPOCHS = 128*strategy.num_replicas_in_sync, 2\nBATCH_SIZE, EPOCHS","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"checkpoint_filepath = 'bestmodel'\nmodel_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(\n    filepath=checkpoint_filepath,\n    save_weights_only=True,\n    monitor='val_auc',\n    mode='max',\n    save_best_only=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"n_steps = len(balanced_sampler) // BATCH_SIZE\n\ntrain_history = model.fit(\n    train_dataset,\n    steps_per_epoch=n_steps,\n    validation_data=valid_dataset,\n    epochs=EPOCHS,\n    callbacks=[model_checkpoint_callback]\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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":1}