{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Configuration"},{"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\n!pip install torchsummary\n","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import torch\nimport torch_xla\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.utils.utils as xu\nimport torch_xla.distributed.xla_multiprocessing as xmp\nimport os,re,gc,pickle,random,sys,collections\nimport numpy as np \nimport pandas as pd \nfrom datetime import datetime, timedelta\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom tqdm.notebook import tqdm\nfrom transformers import XLMRobertaModel\nimport torch.nn.functional as F","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# devices = (\n#     xm.get_xla_supported_devices(\n#         max_devices=8))\n# print(\"Devices: {}\".format(devices))\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_size=32\nlearning_rate=1e-5\nMAX_LEN = 192\nnum_epochs=10\ninput1 = \"/kaggle/input/jigsaw-multilingual-toxic-comment-classification/\"\ninpath = \"../input/jwtc-xlmroberta-encoding-192-pickle/datain/\"\nform = \"training/encode_{}.pkl\"\nlangs = [\"en\",\"en2\",\"es\",\"fr\",\"it\",\"pt\",\"ru\",\"tr\"]\nused_data = [\"en\",\"en2\"]\n\n# os.environ['XLA_USE_BF16']=\"1\"\n# os.environ['XLA_TENSOR_ALLOCATOR_MAXSIZE'] = '100000000'","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Load Data\nSome code is borrowed from https://www.kaggle.com/mint101/basic-xlm-r-lb-9442-intro the tensorflow implementation of XLM-R fine tune"},{"metadata":{"trusted":true},"cell_type":"code","source":"def pick_load_format(path):\n    with open(inpath+path,\"rb\") as f:\n        return pickle.load(f)\n\ndef load(path):\n    return pick_load_format(form.format(path))\n\n# def load_data():\n#     train = []\n#     for i in used_data:\n#         df1 = load(i+\"_l1\")\n#         df0 = load(i+\"_l0\")\n#         train += [df1,df0]\n#     train = pd.concat(train)\n\n#     train = np.stack(train.comment_text.values, axis=0).astype(\"int32\"),train.toxic.values\n\n#     valid = pick_load_format(\"valid.pkl\")\n#     x_valid = np.stack(valid.comment_text.values, axis=0).astype(\"int32\")\n#     y_valid = valid.toxic.values\n\n#     test = pick_load_format(\"test.pkl\")\n#     x_test = np.stack(test.content.values, axis=0).astype(\"int32\")\n    \n#     return train,(x_valid,y_valid),x_test\n\ndef get_cong(n,verb=True):\n    tot = round(1+(n*2)/10_000)*10_000\n    if verb: print(\"Pos: {}, Sample neg: {}, Total: {}\".format(n,tot-n,tot))\n    return tot,tot-n\n\ndef load_data(seed=1214):\n    train = []\n    for i in used_data:\n        df1 = load(i+\"_l1\")\n        size, sample_size = get_cong(df1.shape[0])\n        df0 = load(i+\"_l0\").sample(n=sample_size, random_state=seed)\n        train += [df1,df0]\n    train = pd.concat(train)\n\n    train = np.stack(train.comment_text.values, axis=0).astype(\"int32\"),train.toxic.values\n\n    valid = pick_load_format(\"valid.pkl\")\n    x_valid = np.stack(valid.comment_text.values, axis=0).astype(\"int32\")\n    y_valid = valid.toxic.values\n\n    test = pick_load_format(\"test.pkl\")\n    x_test = np.stack(test.content.values, axis=0).astype(\"int32\")\n    \n    return train,(x_valid,y_valid),x_test","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\ntrain,valid,test = load_data()\ngc.collect()\n\nvalid_size = len(valid[1])\ntrain_size = len(train[1])\nprint(train_size,valid_size)\n# only for debug\ntrain=(train[0],train[1])\nvalid=(valid[0],valid[1])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Pytorch dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nfrom torch.utils.data import TensorDataset, DataLoader\nfrom torch import nn\n\n# train_dataloader=DataLoader(TensorDataset(torch.Tensor(train[0]).long(),torch.Tensor(train[0]).long()),\n#                                                 shuffle=True,batch_size=batch_size,num_workers=1)\n# valid_dataloader=DataLoader(TensorDataset(torch.Tensor(valid[0]).long(),torch.Tensor(valid[0]).long()),\n#                                                 shuffle=True,batch_size=batch_size,num_workers=1)\n# gc.collect()\ntrain_dataset=TensorDataset(torch.Tensor(train[0]).long(),torch.Tensor(train[1]).long())\nvalid_dataset=TensorDataset(torch.Tensor(valid[0]).long(),torch.Tensor(valid[1]).long())\ntest_dataset=TensorDataset(torch.Tensor(test).long())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Pretrained model"},{"metadata":{"trusted":true},"cell_type":"code","source":"from transformers import XLMRobertaModel\npretrained_XLM=XLMRobertaModel.from_pretrained('xlm-roberta-large')\n\nclass Model(nn.Module):\n    def __init__(self,pretrained,out_dim=768):\n        super().__init__()\n        self.pretrained_model=pretrained\n        self.fc=nn.Linear(out_dim,1)\n        torch.nn.init.kaiming_uniform_(self.fc.weight)\n        self.fc.bias.data.fill_(0.)\n        self.out_dim=out_dim\n    \n    def forward(self,word_id):\n        # pos id is optional\n        hidden_states=self.pretrained_model(word_id)[0]\n        avg=F.adaptive_avg_pool2d(hidden_states,(1,self.out_dim)).squeeze(1)\n        prob=F.sigmoid(self.fc(avg).squeeze(1))\n        return prob\n\nmodel=Model(pretrained_XLM)\nprint(model)\ngc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from sklearn.metrics import roc_auc_score\n\n# context=dict()\n# context['device']=xm.xla_device()\n# model.to(context['device'])\n# train_dataloader_para=pl.ParallelLoader(train_data_loader,context['device'])\n# valid_dataloader_para=pl.ParallelLoader(valid_data_loader,context['device'])\n# gc.collect()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def reduce_fn(vals):\n    # take average\n    return sum(vals) / len(vals)\n\ndef _train_fn(index):\n    # in parrallel TPU programs, all device dependent code should be written in xmp spawned function\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    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    train_dataloader=DataLoader(\n        train_dataset,\n        batch_size=batch_size,\n        sampler=train_sampler,\n        num_workers=0\n    )\n    valid_dataloader =DataLoader(\n        valid_dataset,\n        batch_size=batch_size,\n        sampler=valid_sampler,\n        num_workers=0\n    )\n    device=xm.xla_device()\n    print(device)\n    model_=model.to(device)\n    nepoch=num_epochs\n    optimizer=torch.optim.Adam(model_.parameters())\n    scheduler=torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,mode='max',factor=0.1,patience=3,threshold=0.0001,threshold_mode='rel',cooldown=0)\n    earlyStoppingPatience=8\n    loss_fn=nn.BCELoss()\n    eval_metric=roc_auc_score\n    \n    \n    best_score=0\n    no_improve=0\n    for i in range(1,nepoch+1):\n        xm.master_print('{}/{} epochs'.format(i,nepoch))\n        # train\n        model.train()\n        para_train_loader=pl.ParallelLoader(train_dataloader,[device])\n        para_train_loader=para_train_loader.per_device_loader(device)\n        xm.master_print('parallel loader created... training now')\n        pbar=tqdm(enumerate(para_train_loader),total=len(para_train_loader))\n        avg_loss=0\n        for i,(data, target) in pbar:\n#         for i,(data,target) in enumerate(para_train_loader):\n            data, target=data.to(device),target.double().to(device)\n            optimizer.zero_grad()\n            output = model_(data)\n            loss = loss_fn(output, target)\n            loss.backward()\n            xm.optimizer_step(optimizer)\n            loss_reduced = xm.mesh_reduce('loss_reduce',loss,reduce_fn) \n#             master_print will only print once (not from all 8 cores)\n            pbar.set_postfix({'loss':loss_reduced})\n        # eval\n        model.eval()\n        para_valid_loader=pl.ParallelLoader(valid_dataloader,[device])\n        para_valid_loader=para_valid_loader.per_device_loader(device)\n        xm.master_print('parallel loader created... evaluating now')\n        pbar=tqdm(enumerate(para_valid_loader),total=len(para_valid_loader))\n        outputs=[]\n        targets=[]\n        for i,(data,target) in pbar:\n#         for i,(data,target) in enumerate(para_valid_loader):\n            data, target=data.to(device),target.to(device)\n            outputs.extend(model_(data).detach().numpy().tolist())\n            targets.extend(target.detach().numpy().tolist())\n        score=eval_metric(targets,outputs)\n        xm.master_print('score={}'.format(score))\n        if score>best_score:\n            xm.save(mode.state_dict(),'robertaCKPT')\n            best_score=score\n            no_improve=0\n        else:\n            no_improve+=1\n            if no_improve>earlyStoppingPatience:\n                xm.master_print('EarlyStoppinpg while best score is {}'.format(best_score))\n            break\n        scheduler.step(score)\n\ndef _mp_fn(index):\n    _ = _train_fn()\n\n\nxmp.spawn(_train_fn,args=(),start_method='fork',nprocs=1)","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}