{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Super Fast Inference with TPUs for BERT\n This kernel is based on the [excellent kernel](https://www.kaggle.com/shonenkov/tpu-inference-super-fast-xlmroberta) by @shonenkov. I have just adapted the code for my dataloaders and nn.module implementation so that it can be used for BERT (@abhishek).\n Usually inference for bert-base-multilingual-uncased takes around 10 mins on GPU, but this one gets done in 2 mins. Have fun!"},{"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 > /dev/null\n!python pytorch-xla-env-setup.py --version 20200416 --apt-packages libomp5 libopenblas-dev > /dev/null\n!pip install transformers > /dev/null\n!pip install pandarallel > /dev/null","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-output":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)\nimport os\nfrom datetime import datetime\n\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\nfrom glob import glob\nfor path in glob(f'../input/*'):\n    print(path)\n\nfrom nltk import sent_tokenize\nfrom pandarallel import pandarallel\npandarallel.initialize(nb_workers=2, progress_bar=True)\n# Any results you write to the current directory are saved as output.","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import os\nimport torch\nimport pandas as pd\nfrom scipy import stats\nimport numpy as np\nimport warnings\nfrom tqdm import tqdm\nfrom collections import OrderedDict, namedtuple\nimport torch.nn as nn\nfrom torch.optim import lr_scheduler\nimport joblib\nimport time\nimport logging\nimport transformers\nfrom transformers import AdamW, get_linear_schedule_with_warmup, get_constant_schedule\nimport sys\nfrom sklearn import metrics, model_selection\nimport re\nwarnings.filterwarnings(\"ignore\")\nimport nltk\nnltk.download('punkt')\nimport matplotlib.pyplot as plt","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df_test = pd.read_csv(\"/kaggle/input/jigsaw-multilingual-toxic-test-translated/jigsaw_miltilingual_test_translated.csv\")\nsample = pd.read_csv(\"/kaggle/input/jigsaw-multilingual-toxic-comment-classification/sample_submission.csv\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"MAX_LEN = 192\nBERT_PATH = \"/kaggle/input/bert-base-multilingual-uncased/\"\nMODEL_PATH = \"/kaggle/input/bert-base-multilingual-uncased/pytorch_model.bin\"\nTOKENIZER = transformers.BertTokenizer.from_pretrained(\n    BERT_PATH,\n    do_lower_case=True\n)\nclass BERTDataset:\n    def __init__(self, comment_text, tokenizer, max_len):\n        self.comment_text = comment_text\n        self.tokenizer = tokenizer\n        self.max_len = max_len\n    \n    def __len__(self):\n        return len(self.comment_text)\n    \n    def __getitem__(self, item):\n        comment_text = str(self.comment_text[item])\n        comment_text = \" \".join(comment_text.split())\n\n        inputs = self.tokenizer.encode_plus(\n            comment_text,\n            None,\n            add_special_tokens=True,\n            max_length=self.max_len,\n            pad_to_max_length=True\n        )\n\n        ids = inputs[\"input_ids\"]\n        mask = inputs[\"attention_mask\"]\n        token_type_ids = inputs[\"token_type_ids\"]\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        }\n    \nclass BERTBaseUncased(nn.Module):\n    def __init__(self):\n        super(BERTBaseUncased, self).__init__()\n        self.bert = transformers.BertModel.from_pretrained(BERT_PATH)\n        self.bert_drop = nn.Dropout(0.3)\n        self.out = nn.Linear(768*2, 1)\n\n    def forward(self, ids, mask, token_type_ids):\n        o1 , o2 = self.bert(\n            ids,\n            attention_mask=mask,\n            token_type_ids=token_type_ids\n        )\n        mean_pooling = torch.mean(o1,1)\n        max_pooling, _ = torch.max(o1,1)\n        final = torch.cat((mean_pooling, max_pooling),1)\n        final = self.bert_drop(final)\n        output = self.out(final)\n        return output","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch_xla\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\n\nclass MultiTPUPredictor:\n    \n    def __init__(self, model, device):\n        if not os.path.exists('node_submissions'):\n            os.makedirs('node_submissions')\n\n        self.model = model\n        self.device = device\n\n        xm.master_print(f'Model prepared. Device is {self.device}')\n\n\n    def run_inference(self, test_loader, verbose=True, verbose_step=50):\n        self.model.eval()\n        result = {'id': [], 'toxic': []}\n        t = time.time()\n        with torch.no_grad():\n            for bi, d in tqdm(enumerate(test_loader), total=len(test_loader)):\n                ids = d[\"ids\"]\n                token_type_ids = d[\"token_type_ids\"]\n                mask = d[\"mask\"]\n\n                ids = ids.to(self.device, dtype=torch.long)\n                inputs = token_type_ids.to(self.device, dtype=torch.long)\n                attention_masks = mask.to(self.device, dtype=torch.long)\n\n                outputs = self.model(\n                    ids=ids,\n                    mask=attention_masks,\n                    token_type_ids=inputs\n                )\n                if verbose:\n                    if bi % 50 == 0:\n                        xm.master_print(f'Prediction Step {bi}, time: {(time.time() - t):.5f}')\n\n                toxic = torch.sigmoid(outputs).cpu().detach().numpy()\n\n                result['id'].extend(ids.cpu().detach().numpy())\n                result['toxic'].extend(toxic)\n\n        result = pd.DataFrame(result)\n        node_count = len(glob('node_submissions/*.csv'))\n        result.to_csv(f'node_submissions/submission_{node_count}_{datetime.utcnow().microsecond}.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"BERT_PATH=\"bert-base-multilingual-uncased\"\nmx = BERTBaseUncased()\ntokenizer = transformers.BertTokenizer.from_pretrained(\"bert-base-multilingual-uncased\", do_lower_case=True)\n# mx.load_state_dict(torch.load(\"/content/drive/My Drive/TCC/Translated/TCC-Translated-Epoch-4-Fold-1-93.37.bin\"))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_dataset = BERTDataset(\n        comment_text=df_test.content.values,\n        tokenizer=TOKENIZER,\n        max_len=MAX_LEN\n    )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _mp_fn(rank, flags):\n    device = xm.xla_device()\n    model = mx.to(device)\n\n    test_sampler = torch.utils.data.distributed.DistributedSampler(\n        test_dataset,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=False\n    )\n    test_loader = torch.utils.data.DataLoader(\n        test_dataset,\n        batch_size=16,\n        sampler=test_sampler,\n        pin_memory=False,\n        drop_last=False,\n        num_workers=1\n    )\n\n    fitter = MultiTPUPredictor(model=model, device=device)\n    fitter.run_inference(test_loader)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\nFLAGS={}\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}