{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Install torch_xla"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n!python pytorch-xla-env-setup.py --apt-packages libomp5 libopenblas-dev","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Import libraries and utility scripts"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import os\n\nimport dataset\nimport engine\nimport torch\nimport transformers\nimport warnings\n\nimport pandas as pd\nimport numpy as np\nimport torch.nn as nn\n\nfrom model import JigsawModel\n\nfrom sklearn import metrics\nfrom transformers import AdamW\nfrom transformers import get_linear_schedule_with_warmup\n\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.distributed.xla_multiprocessing as xmp\n\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Define config"},{"metadata":{"trusted":true},"cell_type":"code","source":"class config:\n    MAX_LEN = 192\n    TRAIN_BATCH_SIZE = 64\n    VALID_BATCH_SIZE = 4\n    EPOCHS = 1\n    LEARNING_RATE = 0.5e-5\n    BERT_PATH = \"../input/bert-base-multilingual-uncased/\"\n    MODEL_PATH = \"model.bin\"\n    TOKENIZER = transformers.BertTokenizer.from_pretrained(\n        BERT_PATH,\n        do_lower_case=True\n    )\n    JIGSAW_DATA_PATH = \"../input/jigsaw-multilingual-toxic-comment-classification/\"\n    TRAINING_FILE_1 = os.path.join(\n        JIGSAW_DATA_PATH, \n        \"jigsaw-toxic-comment-train.csv\"\n    )\n    TRAINING_FILE_2 = os.path.join(\n        JIGSAW_DATA_PATH, \n        \"jigsaw-unintended-bias-train.csv\"\n    )\n    VALIDATION_FILE = os.path.join(\n        JIGSAW_DATA_PATH, \n        \"validation.csv\"\n    )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Load model and datasets"},{"metadata":{"trusted":true},"cell_type":"code","source":"MX = JigsawModel(config.BERT_PATH)\n\ndf_train1 = pd.read_csv(\n    config.TRAINING_FILE_1, \n    usecols=[\"comment_text\", \"toxic\"]\n).fillna(\"none\")\n\ndf_train2 = pd.read_csv(\n    config.TRAINING_FILE_2, \n    usecols=[\"comment_text\", \"toxic\"]\n).fillna(\"none\")\n\ndf_valid = pd.read_csv(config.VALIDATION_FILE)\n\ndf_train = pd.concat([df_train1, df_train2], axis=0).reset_index(drop=True)\ndf_train = df_train.sample(frac=1).reset_index(drop=True).head(200000)\n\ndf_train = df_train.reset_index(drop=True)\ndf_valid = df_valid.reset_index(drop=True)\n\ntrain_targets = df_train.toxic.values\nvalid_targets = df_valid.toxic.values","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Main training function"},{"metadata":{"trusted":true},"cell_type":"code","source":"def run():\n    train_dataset = dataset.JigsawTraining(\n        comment_text=df_train.comment_text.values,\n        targets=train_targets,\n        config=config\n    )\n\n    train_sampler = torch.utils.data.distributed.DistributedSampler(\n          train_dataset,\n          num_replicas=xm.xrt_world_size(),\n          rank=xm.get_ordinal(),\n          shuffle=True)\n\n    train_data_loader = torch.utils.data.DataLoader(\n        train_dataset,\n        batch_size=config.TRAIN_BATCH_SIZE,\n        sampler=train_sampler,\n        drop_last=True,\n        num_workers=2\n    )\n\n    valid_dataset = dataset.JigsawTraining(\n        comment_text=df_valid.comment_text.values,\n        targets=valid_targets,\n        config=config\n    )\n\n    valid_sampler = torch.utils.data.distributed.DistributedSampler(\n          valid_dataset,\n          num_replicas=xm.xrt_world_size(),\n          rank=xm.get_ordinal(),\n          shuffle=False)\n\n    valid_data_loader = torch.utils.data.DataLoader(\n        valid_dataset,\n        batch_size=config.VALID_BATCH_SIZE,\n        sampler=valid_sampler,\n        drop_last=False,\n        num_workers=1\n    )\n\n    device = xm.xla_device()\n    model = MX.to(device)\n    \n    param_optimizer = list(model.named_parameters())\n    no_decay = [\"bias\", \"LayerNorm.bias\", \"LayerNorm.weight\"]\n    optimizer_parameters = [\n        {\n            'params': [\n                p for n, p in param_optimizer if not any(\n                    nd in n for nd in no_decay\n                )\n            ], \n            'weight_decay': 0.001\n        },\n        {\n            'params': [\n                p for n, p in param_optimizer if any(\n                    nd in n for nd in no_decay\n                )\n            ],\n            'weight_decay': 0.0\n        },\n    ]\n\n    num_train_steps = int(\n        len(df_train) / config.TRAIN_BATCH_SIZE / xm.xrt_world_size() * config.EPOCHS\n    )\n    optimizer = AdamW(\n        optimizer_parameters, \n        lr=config.LEARNING_RATE * xm.xrt_world_size()\n    )\n    scheduler = get_linear_schedule_with_warmup(\n        optimizer,\n        num_warmup_steps=0,\n        num_training_steps=num_train_steps\n    )\n\n    best_auc = 0\n    for epoch in range(config.EPOCHS):\n        para_loader = pl.ParallelLoader(train_data_loader, [device])\n        engine.train_fn(\n            para_loader.per_device_loader(device), \n            model, \n            optimizer, \n            device, \n            scheduler\n        )\n        \n        para_loader = pl.ParallelLoader(valid_data_loader, [device])\n        outputs, targets = engine.eval_fn(\n            para_loader.per_device_loader(device), \n            model, \n            device\n        )\n\n        targets = np.array(targets) >= 0.5\n        auc = metrics.roc_auc_score(targets, outputs)\n        print(f'[xla:{xm.get_ordinal()}]: AUC={auc}')\n        if auc > best_auc:\n            xm.save(model.state_dict(), config.MODEL_PATH)\n            best_auc = auc","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Multi-processing wrapper"},{"metadata":{"trusted":true},"cell_type":"code","source":"def _mp_fn(rank, flags):\n    torch.set_default_tensor_type('torch.FloatTensor')\n    a = run()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Process spawner for training on TPUs"},{"metadata":{"trusted":true},"cell_type":"code","source":"FLAGS={}\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')","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}