{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Training Pipeline by [@shonenkov](https://www.kaggle.com/shonenkov) using multi TPU on PyTorch/XLA\n\nHi everyone!\n\nMy name is Alex Shonenkov, I am researcher, in Love with NLP and DL.\n\nRecently I have published my ideas about this competition:\n\n- [[TPU-Inference] Super Fast XLMRoberta](https://www.kaggle.com/shonenkov/tpu-inference-super-fast-xlmroberta)\n- [NLP Albumentations](https://www.kaggle.com/shonenkov/nlp-albumentations)\n- [Hack with Parallel Corpus](https://www.kaggle.com/shonenkov/hack-with-parallel-corpus)\n- [Class Balance with PyTorch/XLA](https://www.kaggle.com/shonenkov/class-balance-with-pytorch-xla)\n- [open-subtitles-toxic-pseudo-labeling](https://www.kaggle.com/shonenkov/open-subtitles-toxic-pseudo-labeling)\n\nif you didn't see this kernels and datasets, I recommend to read all of them because it may help you for better understand this kernel and achieve success in competition :)"},{"metadata":{},"cell_type":"markdown","source":"## MAIN IDEA\n\nI spent a lot of time for create working kernel on the kaggle, I have tried to optimize it for 16GB RAM. But I was not able to do it for distributed MULTI TPU here, because of my datasets is too big for this aims.\n\nHere I would like to demonstrate my training pipeline without running and also I would like to provide you, my firends, prepared Colab notebook with kaggle structure!\n\nSo lets start!"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true,"_kg_hide-input":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 20200420 --apt-packages libomp5 libopenblas-dev > /dev/null\n!pip install transformers==2.5.1 > /dev/null\n!pip install pandarallel > /dev/null\n!pip install catalyst==20.4.2 > /dev/null","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\n\nimport os\nos.environ['XLA_USE_BF16'] = \"1\"\n\nfrom glob import glob\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset,DataLoader\nfrom torch.autograd import Variable\nfrom torch.utils.data.sampler import SequentialSampler, RandomSampler\nimport sklearn\n\nimport time\nimport random\nfrom datetime import datetime\nfrom tqdm import tqdm\ntqdm.pandas()\n\nfrom transformers import BertModel, BertTokenizer\nfrom transformers import XLMRobertaModel, XLMRobertaTokenizer\nfrom transformers import AdamW, get_linear_schedule_with_warmup, get_constant_schedule\nfrom catalyst.data.sampler import DistributedSamplerWrapper, BalanceClassSampler\n\nimport gc\nimport re\n\n# !pip install nltk > /dev/null\nimport nltk\nnltk.download('punkt')\n\nfrom nltk import sent_tokenize\n\nfrom pandarallel import pandarallel\n\npandarallel.initialize(nb_workers=4, progress_bar=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"SEED = 42\n\nMAX_LENGTH = 224\nBACKBONE_PATH = 'xlm-roberta-large'\nROOT_PATH = f'..'\n# ROOT_PATH = f'/content/drive/My Drive/jigsaw2020-kaggle-public-baseline' # for colab\n\n\ndef seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = True\n\nseed_everything(SEED)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### [NLP Albumentations](https://www.kaggle.com/shonenkov/nlp-albumentations)"},{"metadata":{"trusted":true},"cell_type":"code","source":"from nltk import sent_tokenize\nfrom random import shuffle\nimport random\nimport albumentations\nfrom albumentations.core.transforms_interface import DualTransform, BasicTransform\n\n\nLANGS = {\n    'en': 'english',\n    'it': 'italian', \n    'fr': 'french', \n    'es': 'spanish',\n    'tr': 'turkish', \n    'ru': 'russian',\n    'pt': 'portuguese'\n}\n\ndef get_sentences(text, lang='en'):\n    return sent_tokenize(text, LANGS.get(lang, 'english'))\n\ndef exclude_duplicate_sentences(text, lang='en'):\n    sentences = []\n    for sentence in get_sentences(text, lang):\n        sentence = sentence.strip()\n        if sentence not in sentences:\n            sentences.append(sentence)\n    return ' '.join(sentences)\n\ndef clean_text(text, lang='en'):\n    text = str(text)\n    text = re.sub(r'[0-9\"]', '', text)\n    text = re.sub(r'#[\\S]+\\b', '', text)\n    text = re.sub(r'@[\\S]+\\b', '', text)\n    text = re.sub(r'https?\\S+', '', text)\n    text = re.sub(r'\\s+', ' ', text)\n    text = exclude_duplicate_sentences(text, lang)\n    return text.strip()\n\n\nclass NLPTransform(BasicTransform):\n    \"\"\" Transform for nlp task.\"\"\"\n\n    @property\n    def targets(self):\n        return {\"data\": self.apply}\n    \n    def update_params(self, params, **kwargs):\n        if hasattr(self, \"interpolation\"):\n            params[\"interpolation\"] = self.interpolation\n        if hasattr(self, \"fill_value\"):\n            params[\"fill_value\"] = self.fill_value\n        return params\n\n    def get_sentences(self, text, lang='en'):\n        return sent_tokenize(text, LANGS.get(lang, 'english'))\n\nclass ShuffleSentencesTransform(NLPTransform):\n    \"\"\" Do shuffle by sentence \"\"\"\n    def __init__(self, always_apply=False, p=0.5):\n        super(ShuffleSentencesTransform, self).__init__(always_apply, p)\n\n    def apply(self, data, **params):\n        text, lang = data\n        sentences = self.get_sentences(text, lang)\n        random.shuffle(sentences)\n        return ' '.join(sentences), lang\n\nclass ExcludeDuplicateSentencesTransform(NLPTransform):\n    \"\"\" Exclude equal sentences \"\"\"\n    def __init__(self, always_apply=False, p=0.5):\n        super(ExcludeDuplicateSentencesTransform, self).__init__(always_apply, p)\n\n    def apply(self, data, **params):\n        text, lang = data\n        sentences = []\n        for sentence in self.get_sentences(text, lang):\n            sentence = sentence.strip()\n            if sentence not in sentences:\n                sentences.append(sentence)\n        return ' '.join(sentences), lang\n\nclass ExcludeNumbersTransform(NLPTransform):\n    \"\"\" exclude any numbers \"\"\"\n    def __init__(self, always_apply=False, p=0.5):\n        super(ExcludeNumbersTransform, self).__init__(always_apply, p)\n\n    def apply(self, data, **params):\n        text, lang = data\n        text = re.sub(r'[0-9]', '', text)\n        text = re.sub(r'\\s+', ' ', text)\n        return text, lang\n\nclass ExcludeHashtagsTransform(NLPTransform):\n    \"\"\" Exclude any hashtags with # \"\"\"\n    def __init__(self, always_apply=False, p=0.5):\n        super(ExcludeHashtagsTransform, self).__init__(always_apply, p)\n\n    def apply(self, data, **params):\n        text, lang = data\n        text = re.sub(r'#[\\S]+\\b', '', text)\n        text = re.sub(r'\\s+', ' ', text)\n        return text, lang\n\nclass ExcludeUsersMentionedTransform(NLPTransform):\n    \"\"\" Exclude @users \"\"\"\n    def __init__(self, always_apply=False, p=0.5):\n        super(ExcludeUsersMentionedTransform, self).__init__(always_apply, p)\n\n    def apply(self, data, **params):\n        text, lang = data\n        text = re.sub(r'@[\\S]+\\b', '', text)\n        text = re.sub(r'\\s+', ' ', text)\n        return text, lang\n\nclass ExcludeUrlsTransform(NLPTransform):\n    \"\"\" Exclude urls \"\"\"\n    def __init__(self, always_apply=False, p=0.5):\n        super(ExcludeUrlsTransform, self).__init__(always_apply, p)\n\n    def apply(self, data, **params):\n        text, lang = data\n        text = re.sub(r'https?\\S+', '', text)\n        text = re.sub(r'\\s+', ' ', text)\n        return text, lang","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### [Pseudo-labeling with open-subtitles](https://www.kaggle.com/shonenkov/hack-with-parallel-corpus)\n\nMore noise with mix of languages can help. I have used [pseudo-labeled open-subtitles dataset](https://www.kaggle.com/shonenkov/open-subtitles-toxic-pseudo-labeling) for this approach. \n\nIt is some analogue for Cutmix in Computer Vision:\n\n<img src='https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F2605845%2Ff29492171d83dfa6b6fcae2af414fcf8%2FCutmix_exmaple.png?generation=1579343294489994&alt=media' align=\"left\"> "},{"metadata":{"trusted":true},"cell_type":"code","source":"class SynthesicOpenSubtitlesTransform(NLPTransform):\n    def __init__(self, always_apply=False, p=0.5):\n        super(SynthesicOpenSubtitlesTransform, self).__init__(always_apply, p)\n        df = pd.read_csv(f'{ROOT_PATH}/input/open-subtitles-toxic-pseudo-labeling/open-subtitles-synthesic.csv', index_col='id')[['comment_text', 'toxic', 'lang']]\n        df = df[~df['comment_text'].isna()]\n        df['comment_text'] = df.parallel_apply(lambda x: clean_text(x['comment_text'], x['lang']), axis=1)\n        df = df.drop_duplicates(subset='comment_text')\n        df['toxic'] = df['toxic'].round().astype(np.int)\n\n        self.synthesic_toxic = df[df['toxic'] == 1].comment_text.values\n        self.synthesic_non_toxic = df[df['toxic'] == 0].comment_text.values\n\n        del df\n        gc.collect();\n\n    def generate_synthesic_sample(self, text, toxic):\n        texts = [text]\n        if toxic == 0:\n            for i in range(random.randint(1,5)):\n                texts.append(random.choice(self.synthesic_non_toxic))\n        else:\n            for i in range(random.randint(0,2)):\n                texts.append(random.choice(self.synthesic_non_toxic))\n            \n            for i in range(random.randint(1,3)):\n                texts.append(random.choice(self.synthesic_toxic))\n        random.shuffle(texts)\n        return ' '.join(texts)\n\n    def apply(self, data, **params):\n        text, toxic = data\n        text = self.generate_synthesic_sample(text, toxic)\n        return text, toxic","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_train_transforms():\n    return albumentations.Compose([\n        ExcludeUsersMentionedTransform(p=0.95),\n        ExcludeUrlsTransform(p=0.95),\n        ExcludeNumbersTransform(p=0.95),\n        ExcludeHashtagsTransform(p=0.95),\n        ExcludeDuplicateSentencesTransform(p=0.95),\n    ], p=1.0)\n\ndef get_synthesic_transforms():\n    return SynthesicOpenSubtitlesTransform(p=0.5)\n\n\ntrain_transforms = get_train_transforms();\nsynthesic_transforms = get_synthesic_transforms()\ntokenizer = XLMRobertaTokenizer.from_pretrained(BACKBONE_PATH)\nshuffle_transforms = ShuffleSentencesTransform(always_apply=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def onehot(size, target):\n    vec = torch.zeros(size, dtype=torch.float32)\n    vec[target] = 1.\n    return vec\n\nclass DatasetRetriever(Dataset):\n\n    def __init__(self, labels_or_ids, comment_texts, langs, use_train_transforms=False, test=False):\n        self.test = test\n        self.labels_or_ids = labels_or_ids\n        self.comment_texts = comment_texts\n        self.langs = langs\n        self.use_train_transforms = use_train_transforms\n        \n    def get_tokens(self, text):\n        encoded = tokenizer.encode_plus(\n            text, \n            add_special_tokens=True, \n            max_length=MAX_LENGTH, \n            pad_to_max_length=True\n        )\n        return encoded['input_ids'], encoded['attention_mask']\n\n    def __len__(self):\n        return self.comment_texts.shape[0]\n\n    def __getitem__(self, idx):\n        text = self.comment_texts[idx]\n        lang = self.langs[idx]\n        if self.test is False:\n            label = self.labels_or_ids[idx]\n            target = onehot(2, label)\n\n        if self.use_train_transforms:\n            text, _ = train_transforms(data=(text, lang))['data']\n            tokens, attention_mask = self.get_tokens(str(text))\n            token_length = sum(attention_mask)\n            if token_length > 0.8*MAX_LENGTH:\n                text, _ = shuffle_transforms(data=(text, lang))['data']\n            elif token_length < 60:\n                text, _ = synthesic_transforms(data=(text, label))['data']\n            else:\n                tokens, attention_mask = torch.tensor(tokens), torch.tensor(attention_mask)\n                return target, tokens, attention_mask\n\n        tokens, attention_mask = self.get_tokens(str(text))\n        tokens, attention_mask = torch.tensor(tokens), torch.tensor(attention_mask)\n\n        if self.test is False:\n            return target, tokens, attention_mask\n        return self.labels_or_ids[idx], tokens, attention_mask\n\n    def get_labels(self):\n        return list(np.char.add(self.labels_or_ids.astype(str), self.langs))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Here I have used [this kernel](https://www.kaggle.com/shonenkov/prepare-training-data) for merging all train data "},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\ndf_train = pd.read_csv(f'{ROOT_PATH}/input/jigsaw-public-baseline-train-data/train_data.csv')\n\n\ntrain_dataset = DatasetRetriever(\n    labels_or_ids=df_train['toxic'].values, \n    comment_texts=df_train['comment_text'].values, \n    langs=df_train['lang'].values,\n    use_train_transforms=True,\n)\n\ndel df_train\ngc.collect();\n\nfor targets, tokens, attention_masks in train_dataset:\n    break\n    \nprint(targets)\nprint(tokens.shape)\nprint(attention_masks.shape)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Class Balance\n\nAfter some experiments I have decided that [class balance](https://www.kaggle.com/shonenkov/class-balance-with-pytorch-xla) in this competition is very important. Also I noticed impact if use balancing dataset by languages.\n\nHere you can see unique values for get_labels method:"},{"metadata":{"trusted":true},"cell_type":"code","source":"np.unique(train_dataset.get_labels())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df_val = pd.read_csv(f'{ROOT_PATH}/input/jigsaw-multilingual-toxic-comment-classification/validation.csv', index_col='id')\n\nvalidation_tune_dataset = DatasetRetriever(\n    labels_or_ids=df_val['toxic'].values, \n    comment_texts=df_val['comment_text'].values, \n    langs=df_val['lang'].values,\n    use_train_transforms=True,\n)\n\ndf_val['comment_text'] = df_val.parallel_apply(lambda x: clean_text(x['comment_text'], x['lang']), axis=1)\n\nvalidation_dataset = DatasetRetriever(\n    labels_or_ids=df_val['toxic'].values, \n    comment_texts=df_val['comment_text'].values, \n    langs=df_val['lang'].values,\n    use_train_transforms=False,\n)\n\ndel df_val\ngc.collect();\n\nfor targets, tokens, attention_masks in validation_dataset:\n    break\n\nprint(targets)\nprint(tokens.shape)\nprint(attention_masks.shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df_test = pd.read_csv(f'{ROOT_PATH}/input/jigsaw-multilingual-toxic-comment-classification/test.csv', index_col='id')\ndf_test['comment_text'] = df_test.parallel_apply(lambda x: clean_text(x['content'], x['lang']), axis=1)\n\ntest_dataset = DatasetRetriever(\n    labels_or_ids=df_test.index.values, \n    comment_texts=df_test['comment_text'].values, \n    langs=df_test['lang'].values,\n    use_train_transforms=False,\n    test=True\n)\n\ndel df_test\ngc.collect();\n\nfor ids, tokens, attention_masks in test_dataset:\n    break\n\nprint(ids)\nprint(tokens.shape)\nprint(attention_masks.shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class RocAucMeter(object):\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.y_true = np.array([0,1])\n        self.y_pred = np.array([0.5,0.5])\n        self.score = 0\n\n    def update(self, y_true, y_pred):\n        y_true = y_true.cpu().numpy().argmax(axis=1)\n        y_pred = nn.functional.softmax(y_pred, dim=1).data.cpu().numpy()[:,1]\n        self.y_true = np.hstack((self.y_true, y_true))\n        self.y_pred = np.hstack((self.y_pred, y_pred))\n        self.score = sklearn.metrics.roc_auc_score(self.y_true, self.y_pred, labels=np.array([0, 1]))\n    \n    @property\n    def avg(self):\n        return self.score\n\nclass AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Label Smoothing is all you need\nNow we can use translating and augmenting data for training with this Loss: "},{"metadata":{"trusted":true},"cell_type":"code","source":"class LabelSmoothing(nn.Module):\n    def __init__(self, smoothing = 0.1):\n        super(LabelSmoothing, self).__init__()\n        self.confidence = 1.0 - smoothing\n        self.smoothing = smoothing\n\n    def forward(self, x, target):\n        if self.training:\n            x = x.float()\n            target = target.float()\n            logprobs = torch.nn.functional.log_softmax(x, dim = -1)\n            nll_loss = -logprobs * target\n            nll_loss = nll_loss.sum(-1)\n            smooth_loss = -logprobs.mean(dim=-1)\n            loss = self.confidence * nll_loss + self.smoothing * smooth_loss\n            return loss.mean()\n        else:\n            return torch.nn.functional.cross_entropy(x, target)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Custom TPU Fitter\n\n<img src='https://image.made-in-china.com/202f0j10dPkYMNLhhabE/Children-Bicycle-Baby-Kids-BMX-Bike.jpg' width=250 align=\"left\"> \n\nP.S. Lets go to do contributing [Catalyst](https://github.com/catalyst-team/catalyst) with TPU backend :)\n\n<img src='https://raw.githubusercontent.com/catalyst-team/catalyst-pics/master/pics/catalyst_logo.png' width=100 align=\"center\">\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"import warnings\n\nwarnings.filterwarnings(\"ignore\")\n\nimport 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\nfrom catalyst.data.sampler import DistributedSamplerWrapper, BalanceClassSampler\n\nclass TPUFitter:\n    \n    def __init__(self, model, device, config):\n        if not os.path.exists('node_submissions'):\n            os.makedirs('node_submissions')\n\n        self.config = config\n        self.epoch = 0\n        self.log_path = 'log.txt'\n\n        self.model = model\n        self.device = device\n\n        param_optimizer = list(self.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(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\n        self.optimizer = AdamW(optimizer_grouped_parameters, lr=config.lr*xm.xrt_world_size())\n        self.scheduler = config.SchedulerClass(self.optimizer, **config.scheduler_params)\n\n        self.criterion = config.criterion\n        xm.master_print(f'Fitter prepared. Device is {self.device}')\n\n    def fit(self, train_loader, validation_loader):\n        for e in range(self.config.n_epochs):\n            if self.config.verbose:\n                lr = self.optimizer.param_groups[0]['lr']\n                timestamp = datetime.utcnow().isoformat()\n                self.log(f'\\n{timestamp}\\nLR: {lr}')\n\n            t = time.time()\n            para_loader = pl.ParallelLoader(train_loader, [self.device])\n            losses, final_scores = self.train_one_epoch(para_loader.per_device_loader(self.device))\n            \n            self.log(f'[RESULT]: Train. Epoch: {self.epoch}, loss: {losses.avg:.5f}, final_score: {final_scores.avg:.5f}, time: {(time.time() - t):.5f}')\n\n            t = time.time()\n            para_loader = pl.ParallelLoader(validation_loader, [self.device])\n            losses, final_scores = self.validation(para_loader.per_device_loader(self.device))\n\n            self.log(f'[RESULT]: Validation. Epoch: {self.epoch}, loss: {losses.avg:.5f}, final_score: {final_scores.avg:.5f}, time: {(time.time() - t):.5f}')\n\n            if self.config.validation_scheduler:\n                self.scheduler.step(metrics=final_scores.avg)\n\n            self.epoch += 1\n    \n    def run_tuning_and_inference(self, test_loader, validation_tune_loader):\n        for e in range(2):\n            self.optimizer.param_groups[0]['lr'] = self.config.lr*xm.xrt_world_size() / (e + 1)\n            para_loader = pl.ParallelLoader(validation_tune_loader, [self.device])\n            losses, final_scores = self.train_one_epoch(para_loader.per_device_loader(self.device))\n            para_loader = pl.ParallelLoader(test_loader, [self.device])\n            self.run_inference(para_loader.per_device_loader(self.device))\n\n    def validation(self, val_loader):\n        self.model.eval()\n        losses = AverageMeter()\n        final_scores = RocAucMeter()\n\n        t = time.time()\n        for step, (targets, inputs, attention_masks) in enumerate(val_loader):\n            if self.config.verbose:\n                if step % self.config.verbose_step == 0:\n                    xm.master_print(\n                        f'Valid Step {step}, loss: ' + \\\n                        f'{losses.avg:.5f}, final_score: {final_scores.avg:.5f}, ' + \\\n                        f'time: {(time.time() - t):.5f}'\n                    )\n            with torch.no_grad():\n                inputs = inputs.to(self.device, dtype=torch.long) \n                attention_masks = attention_masks.to(self.device, dtype=torch.long) \n                targets = targets.to(self.device, dtype=torch.float) \n\n                outputs = self.model(inputs, attention_masks)\n                loss = self.criterion(outputs, targets)\n                \n                batch_size = inputs.size(0)\n\n                final_scores.update(targets, outputs)\n                losses.update(loss.detach().item(), batch_size)\n                \n        return losses, final_scores\n         \n    def train_one_epoch(self, train_loader):\n        self.model.train()\n\n        losses = AverageMeter()\n        final_scores = RocAucMeter()\n        t = time.time()\n        for step, (targets, inputs, attention_masks) in enumerate(train_loader):   \n            if self.config.verbose:\n                if step % self.config.verbose_step == 0:\n                    self.log(\n                        f'Train Step {step}, loss: ' + \\\n                        f'{losses.avg:.5f}, final_score: {final_scores.avg:.5f}, ' + \\\n                        f'time: {(time.time() - t):.5f}'\n                    )\n\n            inputs = inputs.to(self.device, dtype=torch.long)\n            attention_masks = attention_masks.to(self.device, dtype=torch.long)\n            targets = targets.to(self.device, dtype=torch.float)\n\n            self.optimizer.zero_grad()\n\n            outputs = self.model(inputs, attention_masks)\n            loss = self.criterion(outputs, targets)\n\n            batch_size = inputs.size(0)\n            \n            final_scores.update(targets, outputs)\n            \n            losses.update(loss.detach().item(), batch_size)\n\n            loss.backward()\n            xm.optimizer_step(self.optimizer)\n\n            if self.config.step_scheduler:\n                self.scheduler.step()\n        \n        self.model.eval()\n        self.save('last-checkpoint.bin')\n        return losses, final_scores\n\n    def run_inference(self, test_loader):\n        self.model.eval()\n        result = {'id': [], 'toxic': []}\n        t = time.time()\n        for step, (ids, inputs, attention_masks) in enumerate(test_loader):\n            if self.config.verbose:\n                if step % self.config.verbose_step == 0:\n                    xm.master_print(f'Prediction Step {step}, time: {(time.time() - t):.5f}')\n\n            with torch.no_grad():\n                inputs = inputs.to(self.device, dtype=torch.long) \n                attention_masks = attention_masks.to(self.device, dtype=torch.long)\n                outputs = self.model(inputs, attention_masks)\n                toxics = nn.functional.softmax(outputs, dim=1).data.cpu().numpy()[:,1]\n\n            result['id'].extend(ids.cpu().numpy())\n            result['toxic'].extend(toxics)\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}_{random.random()}.csv', index=False)\n\n    def save(self, path):        \n        xm.save(self.model.state_dict(), path)\n\n    def log(self, message):\n        if self.config.verbose:\n            xm.master_print(message)\n        with open(self.log_path, 'a+') as logger:\n            xm.master_print(f'{message}', logger)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"from transformers import XLMRobertaModel\n\nclass ToxicSimpleNNModel(nn.Module):\n\n    def __init__(self):\n        super(ToxicSimpleNNModel, self).__init__()\n        self.backbone = XLMRobertaModel.from_pretrained(BACKBONE_PATH)\n        self.dropout = nn.Dropout(0.3)\n        self.linear = nn.Linear(\n            in_features=self.backbone.pooler.dense.out_features*2,\n            out_features=2,\n        )\n\n    def forward(self, input_ids, attention_masks):\n        bs, seq_length = input_ids.shape\n        seq_x, _ = self.backbone(input_ids=input_ids, attention_mask=attention_masks)\n        apool = torch.mean(seq_x, 1)\n        mpool, _ = torch.max(seq_x, 1)\n        x = torch.cat((apool, mpool), 1)\n        x = self.dropout(x)\n        return self.linear(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"net = ToxicSimpleNNModel()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Custom Config"},{"metadata":{"trusted":true},"cell_type":"code","source":"class TrainGlobalConfig:\n    num_workers = 0 \n    batch_size = 16 \n    n_epochs = 3\n    lr = 0.5 * 1e-5\n\n    # -------------------\n    verbose = True\n    verbose_step = 50\n    # -------------------\n\n    # --------------------\n    step_scheduler = False  # do scheduler.step after optimizer.step\n    validation_scheduler = True  # do scheduler.step after validation stage loss\n    SchedulerClass = torch.optim.lr_scheduler.ReduceLROnPlateau\n    scheduler_params = dict(\n        mode='max',\n        factor=0.7,\n        patience=0,\n        verbose=False, \n        threshold=0.0001,\n        threshold_mode='abs',\n        cooldown=0, \n        min_lr=1e-8,\n        eps=1e-08\n    )\n    # --------------------\n\n    # -------------------\n    criterion = LabelSmoothing()\n    # -------------------","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Main method"},{"metadata":{"trusted":true},"cell_type":"code","source":"def _mp_fn(rank, flags):\n    device = xm.xla_device()\n    net.to(device)\n\n    train_sampler = DistributedSamplerWrapper(\n        sampler=BalanceClassSampler(labels=train_dataset.get_labels(), mode=\"downsampling\"),\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=TrainGlobalConfig.batch_size,\n        sampler=train_sampler,\n        pin_memory=False,\n        drop_last=True,\n        num_workers=TrainGlobalConfig.num_workers,\n    )\n    validation_sampler = torch.utils.data.distributed.DistributedSampler(\n        validation_dataset,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=False\n    )\n    validation_loader = torch.utils.data.DataLoader(\n        validation_dataset,\n        batch_size=TrainGlobalConfig.batch_size,\n        sampler=validation_sampler,\n        pin_memory=False,\n        drop_last=False,\n        num_workers=TrainGlobalConfig.num_workers\n    )\n    validation_tune_sampler = torch.utils.data.distributed.DistributedSampler(\n        validation_tune_dataset,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=True\n    )\n    validation_tune_loader = torch.utils.data.DataLoader(\n        validation_tune_dataset,\n        batch_size=TrainGlobalConfig.batch_size,\n        sampler=validation_tune_sampler,\n        pin_memory=False,\n        drop_last=False,\n        num_workers=TrainGlobalConfig.num_workers\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=TrainGlobalConfig.batch_size,\n        sampler=test_sampler,\n        pin_memory=False,\n        drop_last=False,\n        num_workers=TrainGlobalConfig.num_workers\n    )\n\n    if rank == 0:\n        time.sleep(1)\n    \n    fitter = TPUFitter(model=net, device=device, config=TrainGlobalConfig)\n    fitter.fit(train_loader, validation_loader)\n    fitter.run_tuning_and_inference(test_loader, validation_tune_loader)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Colab Notebook\n\nI hope Kaggle Team will increase RAM memory for tpu notebook as soon as possible. But now I recommend you use colab pro with HIGH RAM mode :)\n\n[Here](https://drive.google.com/drive/folders/1hbcSRfvtTTlERs7remsRST2amIWAFVry?usp=sharing) I have created public read-only google drive with colab notebook! You can save copy and start training right now!\n\nAlso you can run this code here with nprocs=1, if you need. It works! But it is very slow (~1.5 P100)."},{"metadata":{"trusted":true},"cell_type":"code","source":"# FLAGS={}\n# xmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# submission = pd.concat([pd.read_csv(path) for path in glob('node_submissions/*.csv')]).groupby('id').mean()\n# submission['toxic'].hist(bins=100)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's imagine that this logs have got using Kaggle:"},{"metadata":{"trusted":true},"cell_type":"code","source":"file = open('../input/jigsaw-public-baseline-results/log.txt', 'r')\nfor line in file.readlines():\n    print(line[:-1])\nfile.close()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"This model should be trained ~10 epoch, I have run only 3 epoch for this kernel."},{"metadata":{},"cell_type":"markdown","source":"## Submission\n\nIf you want to get high score ~0.945-0.946 such as [[TPU-Inference] Super Fast XLMRoberta](https://www.kaggle.com/shonenkov/tpu-inference-super-fast-xlmroberta) you should do blend such as [here](https://www.kaggle.com/hamditarek/ensemble), but I would like to make submission with only this kernel"},{"metadata":{"trusted":true},"cell_type":"code","source":"submission = pd.read_csv('../input/jigsaw-public-baseline-results/submission.csv', index_col='id')\nsubmission.hist(bins=100)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"submission.to_csv('submission.csv')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Thank you for reading my kernel!\n\n[Here](https://drive.google.com/drive/folders/1hbcSRfvtTTlERs7remsRST2amIWAFVry?usp=sharing) I have created public read-only google drive with colab notebook! You can save copy and start training right now!\n\n\nIf you like this format of notebooks I would like continue to make kernels with realizations of my ideas."}],"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}