{"cells":[{"metadata":{"trusted":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 --version nightly --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\nimport pandas as pd\nfrom pathlib import Path\nfrom sklearn.utils import shuffle\nfrom sklearn import metrics\n# imports pytorch\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data.distributed import DistributedSampler\nfrom torch.utils.data import DataLoader\n# imports the torch_xla package\nimport torch_xla\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.xla_multiprocessing as xmp\n\n\n#transformers\nimport transformers\nfrom transformers import AdamW, get_linear_schedule_with_warmup, get_constant_schedule\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\nimport warnings\n\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# # Check is TPU is working with Pytorch\n# dev = xm.xla_device()\n# t1 = torch.ones(8, 3, device = dev)\n# print(t1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"BERT_ROOT = \"../input/bertbasemultilingualuncased/\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DATA_ROOT = Path(\"..\")/\"input\"/ \"jigsaw-multilingual-toxic-comment-classification/\"\n\ndf1,df2,df3,test,sample = [pd.read_csv(DATA_ROOT / fname) for fname in [\"jigsaw-toxic-comment-train.csv\",\n                                                                        \"jigsaw-unintended-bias-train.csv\",\n                                                                        \"validation.csv\",\n                                                                        \"test.csv\",\n                                                                        \"sample_submission.csv\"\n                                                                       ]]\ndf2.toxic = df2.toxic.round().astype(int)\ntrain = pd.concat([\n    df1[['comment_text', 'toxic']],\n    df3[['comment_text', 'toxic']],\n    df2[['comment_text', 'toxic']].query('toxic==1'),\n    df2[['comment_text', 'toxic']].query('toxic==0').sample(n=200000, random_state=0)\n])\n\nvalid = df3","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train = train.dropna()\n\ntrain = shuffle(train, random_state=22)\ntrain.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"valid = valid.dropna()\n\nvalid = shuffle(valid, random_state=22)\nvalid.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DatasetClass:\n    def __init__(self, text, labels, tokenizer, max_length):\n        self.text = text\n        self.tokenizer = tokenizer\n        self.max_length = max_length\n        self.labels = labels\n\n    def __len__(self):\n        return len(self.text)\n\n    def __getitem__(self, item):\n        text = str(self.text[item])\n        text = \" \".join(text.split())\n\n        inputs = self.tokenizer.encode_plus(\n            text,\n            None,\n            add_special_tokens=True,\n            max_length=self.max_length,\n        )\n        ids = inputs[\"input_ids\"]\n        token_type_ids = inputs[\"token_type_ids\"]\n        mask = inputs[\"attention_mask\"]\n\n        padding_length = self.max_length - len(ids)\n\n        ids = ids + ([0] * padding_length)\n        mask = mask + ([0] * padding_length)\n        token_type_ids = token_type_ids + ([0] * padding_length)\n\n        return {\n            'ids': torch.tensor(ids, dtype=torch.long),\n            'mask': torch.tensor(mask, dtype=torch.long),\n            'token_type_ids': torch.tensor(token_type_ids, dtype=torch.long),\n            'labels': torch.tensor(self.labels[item], dtype=torch.float)\n        }\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nclass Bert_fn(nn.Module):\n    def __init__(self, bert_path):\n        super(Bert_fn, self).__init__()\n        self.bert_path = bert_path\n        self.bert = transformers.BertModel.from_pretrained(self.bert_path)\n        self.bert_drop = nn.Dropout(0.3)\n        self.out = nn.Linear(768 * 2, 1)\n\n    def forward(\n            self,\n            ids,\n            mask,\n            token_type_ids\n    ):\n        o1, o2 = self.bert(\n            ids,\n            attention_mask=mask,\n            token_type_ids=token_type_ids)\n\n        apool = torch.mean(o1, 1)\n        mpool, _ = torch.max(o1, 1)\n        cat = torch.cat((apool, mpool), 1)\n\n        bo = self.bert_drop(cat)\n        p2 = self.out(bo)\n        return p2\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"BERT_MODEL = Bert_fn(bert_path=\"../input/bertbasemultilingualuncased/\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def main_fn():\n    def loss_fn(outputs, labels):\n        return nn.BCEWithLogitsLoss()(outputs, labels.view(-1, 1))\n\n    def train_fn(data_loader, model, optimizer, device, scheduler=None):\n        model.train()\n        for bi, d in enumerate(data_loader):\n            ids = d[\"ids\"]\n            mask = d[\"mask\"]\n            token_type_ids = d[\"token_type_ids\"]\n            labels = d[\"labels\"]\n\n            ids = ids.to(device, dtype=torch.long)\n            mask = mask.to(device, dtype=torch.long)\n            token_type_ids = token_type_ids.to(device, dtype=torch.long)\n            labels = labels.to(device, dtype=torch.float)\n\n            optimizer.zero_grad()\n            outputs = model(\n                ids=ids,\n                mask=mask,\n                token_type_ids=token_type_ids\n            )\n\n            loss = loss_fn(outputs, labels)\n            if bi % 10 == 0:\n                xm.master_print(f'bi={bi}, loss={loss}')\n\n            loss.backward()\n            xm.optimizer_step(optimizer)\n            if scheduler is not None:\n                scheduler.step()\n\n    def eval_fn(data_loader, model, device):\n        model.eval()\n        fin_targets = []\n        fin_outputs = []\n        for bi, d in enumerate(data_loader):\n            ids = d[\"ids\"]\n            mask = d[\"mask\"]\n            token_type_ids = d[\"token_type_ids\"]\n            labels = d[\"labels\"]\n\n            ids = ids.to(device, dtype=torch.long)\n            mask = mask.to(device, dtype=torch.long)\n            token_type_ids = token_type_ids.to(device, dtype=torch.long)\n            labels = labels.to(device, dtype=torch.float)\n\n            outputs = model(\n                ids=ids,\n                mask=mask,\n                token_type_ids=token_type_ids\n            )\n\n            targets_np = labels.cpu().detach().numpy().tolist()\n            outputs_np = outputs.cpu().detach().numpy().tolist()\n            fin_targets.extend(targets_np)\n            fin_outputs.extend(outputs_np)\n\n        return fin_outputs, fin_targets\n\n    MAX_LEN = 192\n    BATCH_SIZE = 64\n    EPOCHS = 2\n\n    tokenizer = transformers.BertTokenizer.from_pretrained(BERT_ROOT, do_lower_case=True)\n\n    train_labels = train.toxic.values\n    valid_labels = valid.toxic.values\n\n    train_dataset = DatasetClass(\n        text=train.comment_text.values,\n        labels=train_labels,\n        tokenizer=tokenizer,\n        max_length=MAX_LEN\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_loader = torch.utils.data.DataLoader(\n        train_dataset,\n        batch_size=BATCH_SIZE,\n        sampler=train_sampler,\n        drop_last=True,\n        num_workers=4\n    )\n\n    valid_dataset = DatasetClass(\n        text=valid.comment_text.values,\n        labels=valid_labels,\n        tokenizer=tokenizer,\n        max_length=MAX_LEN\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_loader = torch.utils.data.DataLoader(\n        valid_dataset,\n        batch_size=64,\n        sampler=valid_sampler,\n        drop_last=False,\n        num_workers=4\n    )\n\n    device = xm.xla_device()\n    model = BERT_MODEL.to(device)\n\n    param_optimizer = list(model.named_parameters())\n    no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']\n    optimizer_grouped_parameters = [\n        {'params': [p for n, p in param_optimizer if not any(\n            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.4 * 1e-5 * xm.xrt_world_size()\n    num_train_steps = int(len(train_dataset) /\n                          BATCH_SIZE / xm.xrt_world_size() * EPOCHS)\n    xm.master_print(\n        f'num_train_steps = {num_train_steps}, world_size={xm.xrt_world_size()}')\n\n    optimizer = AdamW(optimizer_grouped_parameters, lr=lr)\n    scheduler = get_linear_schedule_with_warmup(\n        optimizer,\n        num_warmup_steps=0,\n        num_training_steps=num_train_steps\n    )\n\n    for epoch in range(EPOCHS):\n        para_loader = pl.ParallelLoader(train_loader, [device])\n        train_fn(para_loader.per_device_loader(device), model,\n                 optimizer, device, scheduler=scheduler)\n\n        para_loader = pl.ParallelLoader(valid_loader, [device])\n        o, t = eval_fn(para_loader.per_device_loader(device), model, device)\n        xm.save(model.state_dict(), \"model.bin\")\n        auc = metrics.roc_auc_score(np.array(t) >= 0.5, o)\n        xm.master_print(f'AUC = {auc}')\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _mp_fn(rank, flags):\n    torch.set_default_tensor_type('torch.FloatTensor')\n    a = main_fn()\n\n\nFLAGS = {}\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')\n\n","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}