{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install --upgrade pip\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 --apt-packages libomp5 libopenblas-dev","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\nfrom tqdm import tqdm\nimport os\nimport time\n\nimport warnings\nwarnings.filterwarnings('ignore')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# training data paths\nTRAIN_INPUT_ID_PATH = '/kaggle/input/jigsaw-compressed-bert-tokens/train-df-compressed-input-ids.npz'\nTRAIN_TOKEN_ID_PATH = '/kaggle/input/jigsaw-compressed-bert-tokens/train-df-compressed-token-type-ids.npz'\nTRAIN_ATTENTION_PATH = '/kaggle/input/jigsaw-compressed-bert-tokens/train-df-compressed-attention-mask.npz'\nTRAIN_TARGETS_PATH = '/kaggle/input/jigsaw-compressed-bert-tokens/train-df-compressed-targets.npz'\n\n# validation data paths\nVALID_INPUT_ID_PATH = '/kaggle/input/jigsaw-compressed-bert-tokens/valid-df-compressed-input-ids.npz'\nVALID_TOKEN_ID_PATH = '/kaggle/input/jigsaw-compressed-bert-tokens/valid-df-compressed-token-type-ids.npz'\nVALID_ATTENTION_PATH = '/kaggle/input/jigsaw-compressed-bert-tokens/valid-df-compressed-attention-mask.npz'\nVALID_TARGETS_PATH = '/kaggle/input/jigsaw-compressed-bert-tokens/valid-df-compressed-targets.npz'\n\n# test data paths\nTEST_INPUT_ID_PATH = '/kaggle/input/jigsaw-compressed-bert-tokens/test-df-compressed-input-ids.npz'\nTEST_TOKEN_ID_PATH = '/kaggle/input/jigsaw-compressed-bert-tokens/test-df-compressed-token-type-ids.npz'\nTEST_ATTENTION_PATH = '/kaggle/input/jigsaw-compressed-bert-tokens/test-df-compressed-attention-mask.npz'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# torch modules\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader\nfrom torch.utils.data.distributed import DistributedSampler\n\n\n# transformer modules\nfrom transformers import BertModel\nfrom transformers import XLMRobertaModel\nfrom transformers import AdamW\nfrom transformers import get_linear_schedule_with_warmup\n\n\n# tpu-specific modules\nimport torch_xla\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.data_parallel as dp\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.distributed.xla_multiprocessing as xmp\n\n\nfrom sklearn import metrics","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class JigsawDataset(object):\n    def __init__(self, input_ids, token_type_ids, attention_mask, targets=None):\n        self.input_ids = input_ids\n        self.token_type_ids = token_type_ids\n        self.attention_mask = attention_mask\n        self.targets = targets\n        \n    def __len__(self):\n        return len(self.input_ids)\n    \n    def __getitem__(self, item):\n        input_ids = self.input_ids[item]\n        token_type_ids = self.token_type_ids[item]\n        attention_mask = self.attention_mask[item]\n\n        if self.targets is None:\n            return {\n                'input_ids': torch.tensor(input_ids, dtype=torch.long),\n                'token_type_ids': torch.tensor(token_type_ids, dtype=torch.long),\n                'attention_mask': torch.tensor(attention_mask, dtype=torch.long)}\n        \n        else:\n            targets = self.targets[item]\n        \n            return {\n                'input_ids': torch.tensor(input_ids, dtype=torch.long),\n                'token_type_ids': torch.tensor(token_type_ids, dtype=torch.long),\n                'attention_mask': torch.tensor(attention_mask, dtype=torch.long),\n                'targets': torch.tensor(targets, dtype=torch.float)}\n            ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# loading compressed numpy training data\nload_train_input_ids = np.load(TRAIN_INPUT_ID_PATH, mmap_mode='r')\nload_train_token_type_ids = np.load(TRAIN_TOKEN_ID_PATH, mmap_mode='r')\nload_train_attention_mask = np.load(TRAIN_ATTENTION_PATH, mmap_mode='r')\nload_train_targets = np.load(TRAIN_TARGETS_PATH, mmap_mode='r')\n\n# training data\ntrain_input_ids = load_train_input_ids.f.arr_0\ntrain_token_type_ids = load_train_token_type_ids.f.arr_0\ntrain_attention_mask = load_train_attention_mask.f.arr_0\ntrain_targets = load_train_targets.f.arr_0\n\n# loading compressed numpy validation data\nload_valid_input_ids = np.load(VALID_INPUT_ID_PATH, mmap_mode='r')\nload_valid_token_type_ids = np.load(VALID_TOKEN_ID_PATH, mmap_mode='r')\nload_valid_attention_mask = np.load(VALID_ATTENTION_PATH, mmap_mode='r')\nload_valid_targets = np.load(VALID_TARGETS_PATH, mmap_mode='r')\n\n# validation data\nvalid_input_ids = load_valid_input_ids.f.arr_0\nvalid_token_type_ids = load_valid_token_type_ids.f.arr_0\nvalid_attention_mask = load_valid_attention_mask.f.arr_0\nvalid_targets = load_valid_targets.f.arr_0\n\n# loading compressed numpy test data\nload_test_input_ids = np.load(TEST_INPUT_ID_PATH, mmap_mode='r')\nload_test_token_type_ids = np.load(TEST_TOKEN_ID_PATH, mmap_mode='r')\nload_test_attention_mask = np.load(TEST_ATTENTION_PATH, mmap_mode='r')\n\n# test data\ntest_input_ids = load_test_input_ids.f.arr_0\ntest_token_type_ids = load_test_token_type_ids.f.arr_0\ntest_attention_mask = load_test_attention_mask.f.arr_0\n\n\n# sanity check for the sizes\nassert train_input_ids.shape[0] == train_token_type_ids.shape[0] \\\n        == train_attention_mask.shape[0] == train_targets.shape[0]\n\nassert valid_input_ids.shape[0] == valid_token_type_ids.shape[0] \\\n        == valid_attention_mask.shape[0] == valid_targets.shape[0]\n\nassert test_input_ids.shape[0] == test_token_type_ids.shape[0] == test_attention_mask.shape[0]\n\n\n# uncomment the lines below to check the sizes of the data rows\nprint(train_input_ids.shape[0], train_token_type_ids.shape[0], train_attention_mask.shape[0], train_targets.shape[0])\nprint(valid_input_ids.shape[0], valid_token_type_ids.shape[0], valid_attention_mask.shape[0], valid_targets.shape[0])\nprint(test_input_ids.shape[0], test_token_type_ids.shape[0], test_attention_mask.shape[0])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"del load_train_input_ids, load_train_token_type_ids, load_train_attention_mask, load_train_targets\ndel load_valid_input_ids, load_valid_token_type_ids, load_valid_attention_mask, load_valid_targets\ndel load_test_input_ids, load_test_token_type_ids, load_test_attention_mask\nimport gc; gc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"gc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class BertBaseUncased(nn.Module):\n    def __init__(self, bert_model, dropout):\n        super(BertBaseUncased, self).__init__()\n        self.bert_model = bert_model\n        self.fc1 = nn.Linear(768, 64)\n        self.fc2 = nn.Linear(64, 1)\n        self.dropout = nn.Dropout(dropout)\n        self.relu = nn.ReLU(True)\n        \n    def forward(self, input_ids, token_type_ids, attention_mask):\n        _, out2 = self.bert_model(\n            input_ids=input_ids,\n            token_type_ids=token_type_ids,\n            attention_mask=attention_mask)\n        \n        bert_out = self.dropout(out2)\n        out = self.relu(self.fc1(bert_out))\n        out = self.fc2(out)\n        return out\n    \n\nclass XLMRoberta(nn.Module):\n    def __init__(self, model, dropout):\n        super(XLMRoberta, self).__init__()\n        self.roberta = model\n        self.fc1 = nn.Linear(1024, 64)\n        self.fc2 = nn.Linear(64, 1)\n        self.dropout = nn.Dropout(dropout)\n        self.relu = nn.ReLU(True)\n        \n    def forward(self, input_ids, token_type_ids, attention_mask):\n        _, out2 = self.roberta(\n            input_ids=input_ids,\n            token_type_ids=token_type_ids,\n            attention_mask=attention_mask)\n        \n        out2 = self.dropout(out2)\n        out = self.relu(self.fc1(out2))\n        out = self.fc2(out)\n        return out","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def loss_fn(output, target):\n    return nn.BCEWithLogitsLoss()(output, target)\n\n\ndef train_fn(model, dataloader, optimizer, device, scheduler=None):\n    model.train()\n    train_loss = []\n    \n    for i, data in enumerate(dataloader):\n        input_ids = data['input_ids']\n        token_type_ids = data['token_type_ids']\n        attention_mask = data['attention_mask']\n        targets = data['targets']\n        \n        input_ids = input_ids.to(device)\n        token_type_ids = token_type_ids.to(device)\n        attention_mask = attention_mask.to(device)\n        targets = targets.to(device)\n        \n        output = model(input_ids, token_type_ids, attention_mask)\n        loss = loss_fn(output, targets.unsqueeze(1))\n        train_loss.append(loss.item())\n        \n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n        \n        if i % 100 == 0:\n            xm.master_print(f\"iteration: {i}, train loss: {loss:.4f}\")\n        \n        if scheduler is not None:\n            scheduler.step()\n            \n    return train_loss\n\n\ndef valid_fn(dataloader, model, device):\n    valid_loss = []\n    outputs = []\n    targets = []\n    \n    with torch.no_grad():\n      \n        for i, data in enumerate(dataloader):\n            input_ids = data['input_ids']\n            token_type_ids = data['token_type_ids']\n            attention_mask = data['attention_mask']\n            target = data['targets']\n\n            input_ids = input_ids.to(device)\n            token_type_ids = token_type_ids.to(device)\n            attention_mask = attention_mask.to(device)\n            target = target.to(device)\n\n            output = model(input_ids, token_type_ids, attention_mask)\n            \n            output_np = output.cpu().detach().numpy().tolist()\n            target_np = target.cpu().detach().numpy().tolist()\n            \n            outputs.extend(output_np)\n            targets.extend(target_np)\n            \n    return outputs, targets\n\n\ndef run(\n    epochs, \n    batch_size, \n    num_workers, \n    learning_rate, \n    warmup_steps,\n    pretrained_model,\n    dropout):\n    \n    \n    # datasets, samplers and dataloaders\n    trainset = JigsawDataset(\n        input_ids=train_input_ids[:512000],\n        token_type_ids=train_token_type_ids[:512000],\n        attention_mask=train_attention_mask[:512000],\n        targets=train_targets[:512000])\n    \n    validset = JigsawDataset(\n        input_ids=valid_input_ids,\n        token_type_ids=valid_token_type_ids,\n        attention_mask=valid_attention_mask,\n        targets=valid_targets)\n\n    \n    # samplers\n    trainsampler = DistributedSampler(\n        dataset=trainset,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=True)\n    \n    validsampler = DistributedSampler(\n        dataset=validset,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=False)\n    \n    \n    # dataloaders\n    trainloader = DataLoader(\n        dataset=trainset,\n        batch_size=batch_size,\n        sampler=trainsampler,\n        num_workers=num_workers,\n        drop_last=True,)\n    \n    validloader = DataLoader(\n        dataset=validset,\n        batch_size=batch_size,\n        sampler=validsampler,\n        num_workers=num_workers,\n        drop_last=True)\n\n    \n    xm.master_print(f\"Loading datasets....Complete!\")\n    \n    # model\n    device = xm.xla_device()\n    model = BertBaseUncased(pretrained_model, dropout)\n    model = model.to(device)\n    xm.master_print(f\"Loading model....Complete!\")\n    \n    # training_parameters, optimizers and schedulers\n    not_decay = ['LayerNorm.weight', 'LayerNorm.bias', 'bias']\n    \n    parameters = list(model.named_parameters())\n    \n    train_parameters = [\n        {'params': [p for n, p in parameters if not any(nd in n for nd in not_decay)], \n         'weight_decay': 0.001},\n        \n        {'params': [p for n, p in parameters if any(nd in n for nd in not_decay)], \n         'weight_decay': 0.001 }]\n    \n    \n    num_training_steps = int(len(trainset) / batch_size / xm.xrt_world_size())\n    xm.master_print(f\"Iterations per epoch: {num_training_steps}, world_size: {xm.xrt_world_size()}\")\n    \n    optimizer = AdamW(train_parameters, lr=learning_rate)\n    \n    scheduler = get_linear_schedule_with_warmup(\n        optimizer,\n        num_warmup_steps=warmup_steps,\n        num_training_steps=num_training_steps)\n    \n    AUC_SCORE = []\n    \n    # training and evaluation\n    for epoch in range(epochs):\n        gc.collect()\n        \n        # train\n        para_loader = pl.ParallelLoader(trainloader, [device])\n        \n        start_time = time.time()\n        \n        train_loss = train_fn(\n            model, \n            para_loader.per_device_loader(device), \n            optimizer, \n            device, \n            scheduler=scheduler)\n        \n        end_time = time.time()\n        \n        del para_loader\n        gc.collect()\n        \n        time_per_epoch = end_time - start_time\n        \n        xm.master_print(f\"Time taken: {(time_per_epoch/60):.2f} mins\")\n        \n        xm.master_print(f\"epoch: {epoch+1}/{epochs}, train loss: {np.mean(train_loss):.4f}\")\n        \n        # eval\n        para_loader = pl.ParallelLoader(validloader, [device])\n        outputs, targets = valid_fn(\n            para_loader.per_device_loader(device),\n            model,\n            device)\n        \n        del para_loader\n        gc.collect()\n        \n        auc = metrics.roc_auc_score(np.array(targets) >= 0.5, outputs)\n            \n        xm.master_print(f\"auc_score: {auc:.4f}\")\n        #xm.master_print(AUC_SCORE)\n        \n        # save model\n        #if epoch > 13  and auc > max(AUC_SCORE):\n            #AUC_SCORE.append(auc)\n            #xm.master_print(f\"Saving model {epoch+1}\")\n    xm.save(model.state_dict(), f\"bert_multilingual_{batch_size}_e{epochs}_auc_{auc:.4f}.bin\")\n        \n        #AUC_SCORE.append(auc)\n        ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# hyper parameters\nMODEL_PATH = 'bert-base-multilingual-uncased'\nBATCH_SIZE = 128\nNUM_WORKERS = 8\nDROPOUT = 0.3\nLR = 1e-5\nEPOCHS = 10\nWARMUP_STEPS = 0\n\nMODEL = BertModel.from_pretrained(MODEL_PATH)\n\ndef _mp_fn(rank, flags):\n    \n    a = run(\n        epochs=EPOCHS,\n        batch_size=BATCH_SIZE,\n        num_workers=NUM_WORKERS,\n        learning_rate=LR,\n        warmup_steps=WARMUP_STEPS,\n        pretrained_model=MODEL,\n        dropout=DROPOUT)\n    \n    \n\nFLAGS = {}\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\ndef predictions(dataloader, model, device):\n    outputs = []\n\n    for i, data in enumerate(dataloader):\n        input_ids = data['input_ids']\n        token_type_ids = data['token_type_ids']\n        attention_mask = data['attention_mask']\n        \n        input_ids = input_ids.to(device)\n        token_type_ids = token_type_ids.to(device)\n        attention_mask = attention_mask.to(device)\n        \n        output = model(\n            input_ids=input_ids, \n            token_type_ids=token_type_ids, \n            attention_mask=attention_mask)\n        \n        output_np = output.cpu().detach().numpy().tolist()\n        outputs.extend(output_np)\n        xm.master_print(len(outputs))\n    return outputs\n\n\n\ndef pred_run():\n    gc.collect()\n    # load model\n    model = BertBaseUncased(MODEL, DROPOUT)\n    device = torch.device('cpu')\n    #device = xm.xla_device()\n    model = model.to(device)\n    model.eval()\n    model.load_state_dict(torch.load(\"bert_multilingual_32_model.bin\"))\n\n    # dataset, sampler and dataloader\n    testset = JigsawDataset(\n        input_ids=test_input_ids[:4000],\n        token_type_ids=test_token_type_ids[:4000],\n        attention_mask=test_attention_mask[:4000])\n    \n    #xm.master_print(f\"size of testset: {len(testset)}\")\n    print(len(testset))\n    \n    \n    testsampler = DistributedSampler(\n        dataset=testset,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=False)\n    \n    \n    print(\"Sampler Done\")\n    \n    testloader = DataLoader(\n        dataset=testset,\n        batch_size=32,\n        shuffle=False,\n        drop_last=True)\n\n    # predictions\n    pred = predictions(dataloader=testloader, model=model, device=device)\n    del para_loader\n    gc.collect()\n    return pred\n\npred = pred_run()\n\n#FLAGS = {}\n#pred = xmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')\n\"\"\"\nprint(\"Finished\")","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}