{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"VERSION = \"nightly\"  #@param [\"1.5\" , \"20200325\", \"nightly\"]\n!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 --version $VERSION","execution_count":null,"outputs":[]},{"metadata":{"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":"Here are all of our imports. You will note that there are already an optimization applied in order for PyTorch XLA to train.\n\n`XLA_USE_BF16` is an environment variable that tells PyTorch XLA to automatically use [bfloat16](https://cloud.google.com/blog/products/ai-machine-learning/bfloat16-the-secret-to-high-performance-on-cloud-tpus).","execution_count":null},{"metadata":{"_uuid":"993abe0b-2561-4abc-a6a2-b83d8dd4c846","_cell_guid":"16a3df23-8416-46b5-8661-c52345005b6d","trusted":true},"cell_type":"code","source":"import os\nos.environ['XLA_USE_BF16']=\"1\"\nos.environ['XLA_TENSOR_ALLOCATOR_MAXSIZE'] = '100000000'\n\nimport numpy as np\nimport pandas as pd\nfrom scipy import stats\n\n\nfrom tqdm import tqdm\nfrom collections import OrderedDict, namedtuple\n\n# torch imports\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset\nfrom torch.utils.data import DataLoader\nfrom torch.optim import lr_scheduler\n\nimport joblib\nimport sys\nimport logging\nimport gc\nimport random\nimport time\n\n# transformers imports\nfrom transformers import AdamW\nfrom transformers import get_linear_schedule_with_warmup\nfrom transformers import get_cosine_schedule_with_warmup\nfrom transformers import get_cosine_with_hard_restarts_schedule_with_warmup\nfrom transformers import get_constant_schedule\nfrom transformers import XLMRobertaConfig\nfrom transformers import XLMRobertaTokenizer\nfrom transformers import XLMRobertaModel\n\n\n# torch_xla imorts\nimport torch_xla\nimport torch_xla.debug.metrics as met\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.utils.utils as xu\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.xla_multiprocessing as xmp\nimport torch_xla.test.test_utils as test_utils\nimport torch_xla.utils.serialization as xser\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\n\nfrom sklearn.metrics import roc_auc_score \nfrom sklearn.model_selection import StratifiedKFold","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(f\"{xm.xrt_world_size()}\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def seed_everything(seed_value):\n    random.seed(seed_value)\n    np.random.seed(seed_value)\n    torch.manual_seed(seed_value)\n    os.environ['PYTHONHASHSEED'] = str(seed_value)\n    \n    if torch.cuda.is_available(): \n        torch.cuda.manual_seed(seed_value)\n        torch.cuda.manual_seed_all(seed_value)\n        torch.backends.cudnn.deterministic = True\n        torch.backends.cudnn.benchmark = True\n\nseed = 999\nseed_everything(seed)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class AverageMeter:\n    \"\"\"\n    Computes and stores the average and current value\n    \"\"\"\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n        \n\nclass ArrayDataset(Dataset):\n    def __init__(self,*arrays):\n        assert all(arrays[0].shape[0] == array.shape[0] for array in arrays)\n        self.arrays = arrays\n    \n    def __getitem__(self, index):\n        return tuple(torch.from_numpy(np.array(array[index])) for array in self.arrays)\n    \n    def __len__(self):\n        return self.arrays[0].shape[0]\n    \n    \nclass JigsawDataset(object):\n    def __init__(self, input_ids=None, token_type_ids=None, attention_mask=None, target=None):\n        self.ids = input_ids\n        self.mask = attention_mask\n        self.token_ids = token_type_ids\n        self.target = target\n        \n    def __len__(self):\n        return self.ids.shape[0]\n    \n    def __getitem__(self, item):\n        return {\n            'ids': torch.from_numpy(np.array(self.ids[item])),\n            'mask': torch.from_numpy(np.array(self.mask[item])),\n            'target': torch.from_numpy(np.array(self.target[item]))}\n    \n \n\n\n# MODELS\n    \nclass XLMRobertaLargeTC(nn.Module):\n    def __init__(self):\n        super(XLMRobertaLargeTC, self).__init__()\n        config = XLMRobertaConfig.from_pretrained('xlm-roberta-large', output_hidden_states=True)\n        self.xlm_roberta = XLMRobertaModel.from_pretrained('xlm-roberta-large', config=config)\n        \n        self.fc = nn.Linear(config.hidden_size, 1)\n        self.dropout = nn.Dropout(p=0.2)\n        \n        # initialize weight\n        nn.init.normal_(self.fc.weight, std=0.02)\n        nn.init.normal_(self.fc.bias, 0)\n        \n        \n    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None):\n        _, o2, _ = self.xlm_roberta(\n            input_ids=input_ids, \n            attention_mask=attention_mask)\n        \n        o2 = self.dropout(o2)\n        logits = self.fc(o2)\n        \n        return logits\n    \n    \nclass XLMRobertaBaseTC(nn.Module):\n    def __init__(self):\n        super(XLMRobertaBaseTC, self).__init__()\n        config = XLMRobertaConfig.from_pretrained('xlm-roberta-base', output_hidden_states=True)\n        self.xlm_roberta = XLMRobertaModel.from_pretrained('xlm-roberta-base', config=config)\n        \n        self.fc = nn.Linear(config.hidden_size, 1)\n        self.dropout = nn.Dropout(p=0.2)\n        \n        # inititalize weights\n        nn.init.normal_(self.fc.weight, std=0.02)\n        nn.init.normal_(self.fc.bias, 0)\n        \n    def forward(self, input_ids=None, token_type_ids=None, attention_mask=None):\n        _, o2, _ = self.xlm_roberta(\n            input_ids=input_ids, \n            attention_mask=attention_mask)\n        \n        o2 = self.dropout(o2)\n        logits = self.fc(o2)\n        \n        return logits\n    \n    \nclass BertMultilingualCaseTC(nn.Module):\n    def __init__(self):\n        super(BertMultilingualCaseTC, self).__init__()\n        config = BertModel.from_pretrained('bert-base-multilingual-cased', output_hidden_states=True)\n        self.bert_model = BertModel.from_pretrained('bert-base-multilingual-cased', config=config)\n        self.fc = nn.Linear(config.hidden_size, 1)\n        \n        # initialize weights\n        nn.init.normal_(self.fc.weight, std=0.02)\n        nn.init.normal_(self.fc.bias, 0)\n        \n        self.dropout = nn.Dropout(p=0.2)\n        \n    def forward(self, input_ids=None, token_type_ids=None, attention_mask=None):\n        _, o2, _ = self.bert_model(\n            input_ids=input_ids, \n            attention_mask=attention_mask)\n        \n        o2 = self.dropout(o2)\n        logits = self.fc(o2)\n        \n        return logits\n    \n    \nclass BertMultilingualUncasedTC(nn.Module):\n    def __init__(self):\n        super(BertMultilingualUncasedTC, self).__init__()\n        config = BertModel.from_pretrained('bert-base-multilingual-uncased', output_hidden_states=True)\n        self.bert_model = BertModel.from_pretrained('bert-base-multilingual-uncased', config=config)\n        self.fc = nn.Linear(config.hidden_size, 1)\n        \n        # initliaze weights\n        nn.init.normal_(self.fc.weight, std=0.02)\n        nn.init.normal_(self.fc.bias, 0)\n        \n        self.dropout = nn.Dropout(p=0.2)\n        \n    def forward(self, input_ids=None, token_type_ids=None, attention_mask=None):\n        _, o2, _ = self.bert_model(input_ids=input_ids, attention_mask=attention_mask)\n        o2 = self.dropout(o2)\n        logits = self.fc(o2)\n        \n        return logits","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch_xla.version as xv\nprint('PYTORCH:', xv.__torch_gitrev__)\nprint('XLA:', xv.__xla_gitrev__)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"!free -h","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def loss_fn(outputs, targets):\n    return nn.BCEWithLogitsLoss()(outputs, targets.view(-1, 1))\n\n\ndef reduce_fn(vals):\n    return sum(vals) / len(vals)\n\n\ndef train_loop_fn(data_loader, model, optimizer, device, scheduler=None):\n    model.train()\n    \n    train_loss = []\n    \n    for bi, data in enumerate(data_loader):\n\n        ids = data['ids']\n        mask = data['mask']\n        targets = data['target']\n\n        ids = ids.to(device, dtype=torch.long)\n        mask = mask.to(device, dtype=torch.long)\n        targets = targets.to(device, dtype=torch.float)\n\n        optimizer.zero_grad()\n        \n        outputs = model(input_ids=ids, attention_mask=mask)\n        loss = loss_fn(outputs, targets)\n        \n        train_loss.append(loss.item())\n        \n        \n        if bi % 500 == 0:\n            loss_reduced = xm.mesh_reduce('loss_reduce', loss, reduce_fn)\n            xm.master_print(f'bi={bi}, loss={loss_reduced:.5f}')\n        \n        loss.backward()\n        \n        xm.optimizer_step(optimizer)\n        \n        if scheduler is not None:\n            scheduler.step()\n            \n    return train_loss\n    \n    \ndef eval_loop_fn(data_loader, model, device):\n    model.eval()\n    \n    fin_targets = []\n    fin_outputs = []\n    \n    for bi, data in enumerate(data_loader):\n        ids = data['ids']\n        mask = data['mask']\n        targets = data['target']\n\n        ids = ids.to(device, dtype=torch.long)\n        mask = mask.to(device, dtype=torch.long)\n        targets = targets.to(device, dtype=torch.float)\n\n        outputs = model(input_ids=ids, attention_mask=mask)\n\n        targets_np = targets.cpu().detach().numpy().tolist()\n        outputs_np = outputs.cpu().detach().numpy().tolist()\n        \n        fin_targets.extend(targets_np)\n        fin_outputs.extend(outputs_np)    \n        \n        del targets_np, outputs_np\n        \n        gc.collect()\n        \n    return fin_outputs, fin_targets","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _run():\n    MAX_LEN = 192\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=TRAIN_BATCH_SIZE,\n        sampler=train_sampler,\n        drop_last=True,\n        num_workers=0)\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=4,\n        sampler=valid_sampler,\n        drop_last=False,\n        num_workers=0)\n\n    device = xm.xla_device()\n    \n    model = mx.to(device)\n    \n    # print only once\n    if fold == 0:\n        xm.master_print('done loading model')\n\n    param_optimizer = list(model.named_parameters())\n    \n    no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n    \n    optimizer_grouped_parameters = [\n        {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.001},\n        {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}]\n\n    lr = 0.5e-5 * xm.xrt_world_size()\n    \n    num_train_steps = int(len(train_dataset) / TRAIN_BATCH_SIZE / xm.xrt_world_size() * EPOCHS)\n    \n    optimizer = AdamW(optimizer_grouped_parameters, lr=lr)\n    \n    scheduler = get_cosine_schedule_with_warmup(\n        optimizer,\n        num_warmup_steps=0,\n        num_training_steps=num_train_steps)\n    \n    xm.master_print(f'num_train_steps = {num_train_steps}, world_size={xm.xrt_world_size()}')\n\n    \n    \n    for epoch in range(EPOCHS):\n        gc.collect()\n        \n        para_loader = pl.ParallelLoader(train_data_loader, [device])\n        \n        # print only once\n        if epoch == 0: \n            xm.master_print('parallel loader created... training now')\n            \n        \n        # train mode/function\n        train_loss = train_loop_fn(para_loader.per_device_loader(device), model, optimizer, device, scheduler=scheduler)\n        \n        del para_loader\n        \n        para_loader = pl.ParallelLoader(valid_data_loader, [device])\n        \n        # eval mode/function\n        o, t = eval_loop_fn(para_loader.per_device_loader(device), model, device)\n        \n        del para_loader\n        \n        auc = roc_auc_score(np.array(t) >= 0.5, o)\n        auc_reduced = xm.mesh_reduce('auc_reduce',auc,reduce_fn)\n        xm.master_print(f'Epoch: {epoch+1}/{EPOCHS} | train loss: {np.mean(train_loss):.4f} | val auc: {auc_reduced:.4f}')\n        \n        gc.collect()\n        \n    xser.save(model.state_dict(), f\"f{fold+1}_xlm_roberta_large.pth\", master_only=True)\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# load pre-encoded binary files\n\nroot_dir = '../input/jtc2020-xlm-roberta-large-tokenizer/'\n\nINPUT_IDS = np.load(root_dir + 'xlm_roberta_large_train_input_ids.npy', mmap_mode='r')\nATTENTION_MASK = np.load(root_dir + 'xlm_roberta_large_train_attention_mask.npy', mmap_mode='r')\nTARGETS = np.load(root_dir + 'train_toxic.npy', mmap_mode='r')\n\n\n# didn't shuffle it earlier\nshuffler = np.random.permutation(len(INPUT_IDS))\nINPUT_IDS = INPUT_IDS[shuffler]\nATTENTION_MASK = ATTENTION_MASK[shuffler]\nTARGETS = TARGETS[shuffler]\n\n\n#we're going to use the validation set provided in the competition as testset\ntest_input_ids = np.load(root_dir + 'xlm_roberta_large_valid_input_ids.npy', mmap_mode='r')\ntest_attention_mask = np.load(root_dir + 'xlm_roberta_large_valid_attention_mask.npy', mmap_mode='r')\ntest_targets = np.load(root_dir + 'valid_toxic.npy', mmap_mode='r')\n\ntest_input_ids = test_input_ids\ntest_attention_mask = test_attention_mask\ntest_targets = test_targets\n\n# for testing code over smaller sample\ninput_ids = INPUT_IDS[:307200]\nattention_mask = ATTENTION_MASK[:307200]\ntargets = TARGETS[:307200]\n\ndel INPUT_IDS, ATTENTION_MASK, TARGETS\ngc.collect()\n\n\nEPOCHS = 1\nTRAIN_BATCH_SIZE = 16\nVALID_BATCH_SIZE = 32\n\n\nskf = StratifiedKFold(n_splits=2, shuffle=True, random_state=seed)\n\nfor fold, (train_idx, valid_idx) in enumerate(skf.split(input_ids, targets)):\n    train_ids = input_ids[train_idx]\n    train_mask = attention_mask[train_idx]\n    train_trg = targets[train_idx]\n\n    train_dataset = JigsawDataset(input_ids=train_ids, attention_mask=train_mask, target=train_trg)\n    valid_dataset = JigsawDataset(input_ids=test_input_ids, attention_mask=test_attention_mask, target=test_targets)\n\n    mx = xmp.MpModelWrapper(XLMRobertaLargeTC())\n    \n    del train_ids, train_mask, train_trg\n    gc.collect()\n    \n    xm.master_print(f\"Fold: {fold+1}\")\n    xm.master_print(\"\")\n    # Start training processes\n    def _mp_fn(rank, flags):\n        a = _run()\n\n    FLAGS={}\n    start_time = time.time()\n    xmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')\n    \nprint('Time taken: ',time.time()-start_time)","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":4}