{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Introduction\n\nThis kernel continues the work in https://www.kaggle.com/theoviel/bert-pytorch-huggingface-starter but adds TPU multiprocessing.\n\nThis model is based on HuggingFace's library. I use the already processed data with multilingual Bert.\n\nThe approach is to finetune the model with last year's data, and validate on the validation data.\n\n\nThis work is in progress and I will spend some time improving it.\n\n### Notes\n\n- The dataset I use for the pretrained Bert is here : https://www.kaggle.com/theoviel/bertconfigs\n- The TPU implementation is adapted from https://www.kaggle.com/dhananjay3/pytorch-xla-for-tpu-with-multiprocessing/#Using-all-8-cores-of-a-v3-TPU-using-pytorch/XLA.\n- Multiprocessing makes the computations twice faster.\n\nFeel free to correct me if I made any mistakes in my implementation."},{"metadata":{},"cell_type":"markdown","source":"# Intialization"},{"metadata":{},"cell_type":"markdown","source":"### Switching to PyTorch Lightning"},{"metadata":{"trusted":true,"_kg_hide-output":true,"_kg_hide-input":true},"cell_type":"code","source":"import os\nimport collections\nfrom datetime import datetime, timedelta\n\nos.environ[\"XRT_TPU_CONFIG\"] = \"tpu_worker;0;10.0.0.2:8470\"\n\n_VersionConfig = collections.namedtuple('_VersionConfig', 'wheels,server')\nVERSION = \"torch_xla==nightly\"\nCONFIG = {\n    'torch_xla==nightly': _VersionConfig('nightly', 'XRT-dev{}'.format(\n        (datetime.today() - timedelta(1)).strftime('%Y%m%d')))}[VERSION]\n\nDIST_BUCKET = 'gs://tpu-pytorch/wheels'\nTORCH_WHEEL = 'torch-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels)\nTORCH_XLA_WHEEL = 'torch_xla-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels)\nTORCHVISION_WHEEL = 'torchvision-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels)\n\n!export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH\n!apt-get install libomp5 -y\n!apt-get install libopenblas-dev -y\n\n!pip uninstall -y torch torchvision\n!gsutil cp \"$DIST_BUCKET/$TORCH_WHEEL\" .\n!gsutil cp \"$DIST_BUCKET/$TORCH_XLA_WHEEL\" .\n!gsutil cp \"$DIST_BUCKET/$TORCHVISION_WHEEL\" .\n!pip install \"$TORCH_WHEEL\"\n!pip install \"$TORCH_XLA_WHEEL\"\n!pip install \"$TORCHVISION_WHEEL\"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Imports"},{"metadata":{"trusted":true,"_kg_hide-input":true,"_kg_hide-output":true},"cell_type":"code","source":"import gc\nimport os\nimport time\nimport math\nimport random\nimport warnings\nimport numpy as np\nimport pandas as pd\nimport seaborn as sns\nimport matplotlib.pyplot as plt\n\nfrom datetime import date\nfrom transformers import *\nfrom sklearn.metrics import *\nfrom tqdm.notebook import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.utils.data\nimport torch.nn.functional as F\n\nfrom torch import Tensor\nfrom torch.optim import *\nfrom torch.nn.modules.loss import *\nfrom torch.optim.lr_scheduler import * \nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.utils.data.sampler import RandomSampler\nfrom torch.utils.data.distributed import DistributedSampler\n\nimport torch_xla\nimport torch_xla.utils.utils as xu\nimport torch_xla.core.xla_model as xm\nimport torch_xla.debug.metrics as met\nimport torch_xla.distributed.data_parallel as dp\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.distributed.xla_multiprocessing as xmp","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Seeding"},{"metadata":{"trusted":true},"cell_type":"code","source":"def seed_everything(seed):\n    \"\"\"\n    Seeds basic parameters for reproductibility of results\n    \n    Arguments:\n        seed {int} -- Number of the seed\n    \"\"\"\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 = False","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"seed = 2020\nseed_everything(seed)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data"},{"metadata":{"trusted":true},"cell_type":"code","source":"MODEL_PATHS = {\n    'bert-multi-cased': '../input/bertconfigs/multi_cased_L-12_H-768_A-12/multi_cased_L-12_H-768_A-12/',\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DATA_PATH = '../input/jigsaw-multilingual-toxic-comment-classification/'\n\ndf_val = pd.read_csv(DATA_PATH + 'validation-processed-seqlen128.csv')\ndf_test =  pd.read_csv(DATA_PATH + 'test-processed-seqlen128.csv')\ndf_train = pd.read_csv(DATA_PATH + 'jigsaw-toxic-comment-train-processed-seqlen128.csv')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.figure(figsize=(10, 5))\nplt.subplot(1, 2, 1)\n\nsns.countplot(df_train['toxic'])\nplt.title('Target repartition on training data')\n\nplt.subplot(1, 2, 2)\nsns.countplot(df_val['toxic'])\nplt.title('Target repartition on validation data')\n\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class JigsawDataset(Dataset):\n    \"\"\"\n    Torch dataset for the competition.\n    \"\"\"\n    def __init__(self, df):\n        \"\"\"\n        Constructor\n        \n        Arguments:\n            df {pandas dataframe} -- Dataframe where the data is. Expects to be one of the []-processed-seqlen128.csv files\n        \"\"\"\n            \n        super().__init__()\n        self.df = df \n        self.word_ids = np.array([word_ids[1:-1].split(', ') for word_ids in df['input_word_ids']]).astype(int)\n        \n        try:\n            self.y = df['toxic'].values\n        except KeyError: # test data\n            self.y = np.zeros(len(df))\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        return torch.tensor(self.word_ids[idx]), torch.tensor(self.y[idx])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"TRANSFORMERS = {\n    \"bert-multi-cased\": (BertModel, BertTokenizer, \"bert-base-uncased\"),\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Transformer(nn.Module):\n    def __init__(self, model, num_classes=1):\n        \"\"\"\n        Constructor\n        \n        Arguments:\n            model {string} -- Transformer to build the model on. Expects \"camembert-base\".\n            num_classes {int} -- Number of classes (default: {1})\n        \"\"\"\n        super().__init__()\n        self.name = model\n\n        model_class, tokenizer_class, pretrained_weights = TRANSFORMERS[model]\n\n        bert_config = BertConfig.from_json_file(MODEL_PATHS[model] + 'bert_config.json')\n        bert_config.output_hidden_states = True\n        \n        self.transformer = BertModel(bert_config)\n\n        self.nb_features = self.transformer.pooler.dense.out_features\n\n        self.pooler = nn.Sequential(\n            nn.Linear(self.nb_features, self.nb_features), \n            nn.Tanh(),\n        )\n\n        self.logit = nn.Linear(self.nb_features, num_classes)\n\n    def forward(self, tokens):\n        \"\"\"\n        Usual torch forward function\n        \n        Arguments:\n            tokens {torch tensor} -- Sentence tokens\n        \n        Returns:\n            torch tensor -- Class logits\n        \"\"\"\n        _, _, hidden_states = self.transformer(\n            tokens, attention_mask=(tokens > 0).long()\n        )\n\n        hidden_states = hidden_states[-1][:, 0] # Use the representation of the first token of the last layer\n\n        ft = self.pooler(hidden_states)\n\n        return self.logit(ft)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"def fit(model, train_dataset, val_dataset, epochs=1, batch_size=32, warmup_prop=0, lr=5e-5):\n    device = xm.xla_device()\n    model = model.to(device)\n    \n    train_sampler = DistributedSampler(\n        train_dataset,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=True\n    )\n    \n    train_loader = torch.utils.data.DataLoader(\n        train_dataset,\n        batch_size=batch_size,\n        sampler=train_sampler,\n        num_workers=0,\n        drop_last=True\n    )\n    \n    valid_loader = torch.utils.data.DataLoader(\n        val_dataset,\n        batch_size=batch_size,\n        shuffle=False,\n        num_workers=0,\n        drop_last=True\n    )\n    \n    \n    lr = lr * xm.xrt_world_size()\n    batch_size = batch_size # * xm.xrt_world_size() # not sure about that\n    \n    optimizer = AdamW(model.parameters(), lr=lr)\n    \n    num_warmup_steps = int(warmup_prop * epochs * len(train_loader))\n    num_training_steps = epochs * len(train_loader)\n    \n    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps)\n    \n    loss_fct = nn.BCEWithLogitsLoss(reduction='sum').to(device)\n    \n    def train_loop(train_loader):\n        optimizer.zero_grad()\n        avg_loss = 0\n        \n        for step, (x, y_batch) in enumerate(train_loader): \n            y_pred = model(x.to(device))\n            \n            loss = loss_fct(y_pred.view(-1).float(), y_batch.float().to(device))\n            loss.backward()\n            avg_loss += loss.item()\n\n            xm.optimizer_step(optimizer, barrier=True)\n            scheduler.step()\n            model.zero_grad()\n            optimizer.zero_grad()        \n            \n        return avg_loss\n            \n    def val_loop(val_loader):\n        preds = []\n        truths = []\n        avg_val_loss = 0.\n\n        with torch.no_grad():\n            for x, y_batch in val_loader:                \n                y_pred = model(x.to(device))\n                loss = loss_fct(y_pred.detach().view(-1).float(), y_batch.float().to(device))\n                avg_val_loss += loss.item()\n                \n                probs = torch.sigmoid(y_pred).detach().cpu().numpy()\n                preds += list(probs.flatten())\n                truths += list(y_batch.cpu().numpy().flatten())\n                \n        return (preds, truths, avg_val_loss)\n        \n    for epoch in range(epochs):\n        avg_losses = []\n        val_loop_out = []\n        start_time = time.time()\n        \n        # Train\n        model.train()\n        para_loader = pl.ParallelLoader(train_loader, [device])\n        avg_losses.append(train_loop(para_loader.per_device_loader(device)))\n        \n        # Eval\n        model.eval()\n        para_loader = pl.ParallelLoader(valid_loader, [device])\n        val_loop_out.append(val_loop(para_loader.per_device_loader(device)))\n        \n        # Retrieve outputs\n        preds = np.array(np.concatenate([v[0] for v in val_loop_out])).flatten()\n        truths = np.array(np.concatenate([v[1] for v in val_loop_out])).flatten()\n        avg_val_losses = [v[2] for v in val_loop_out]\n        \n        # Metrics\n        avg_val_loss = np.sum(avg_val_losses) / len(val_dataset)\n        avg_loss = np.sum(avg_val_losses) / len(val_dataset)\n        val_auc = roc_auc_score(truths, preds)\n        \n        dt = time.time() - start_time\n        lr = scheduler.get_last_lr()[0]\n        \n        xm.master_print(f'Epoch {epoch + 1}/{epochs} \\t lr={lr:.1e} \\t t={dt:.0f}s \\t loss={avg_loss:.4f} \\t val_loss={avg_val_loss:.4f} \\t val_auc={val_auc:.4f}')\n        xm.save(model.state_dict(), \"model.pt\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = Transformer(\"bert-multi-cased\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"epochs = 1 # 1 epoch seems to be enough\nbatch_size = 32\nwarmup_prop = 0.1\nlr = 2e-5  # Important parameter to tweak","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = JigsawDataset(df_train)\nval_dataset = JigsawDataset(df_val)\ntest_dataset = JigsawDataset(df_test)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def fit_multiprocessing(rank, flags):\n    fit(model, train_dataset, val_dataset, epochs=epochs, batch_size=batch_size, warmup_prop=warmup_prop, lr=lr)\n    \nFLAGS = {}\nxmp.spawn(fit_multiprocessing, args=(FLAGS,), nprocs=8, start_method='fork')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Predicting\n- This does not use multi-threading (yet)"},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_model_weights(model, filename, verbose=1, strict=True):\n    print(f'\\n -> Loading weights from {filename}\\n')\n    model.load_state_dict(torch.load(filename, map_location='cpu'), strict=strict)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# def predict(model, dataset, batch_size=64):\n#     \"\"\"\n#     Usual predict torch function\n    \n#     Arguments:\n#         model {torch model} -- Model to predict with\n#         dataset {torch dataset} -- Dataset to get predictions from\n    \n#     Keyword Arguments:\n#         batch_size {int} -- Batch size (default: {32})\n    \n#     Returns:\n#         numpy array -- Predictions\n#     \"\"\"\n#     device = xm.xla_device()\n#     model = model.to(device)\n#     model.eval()\n    \n#     preds = []\n#     loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0)\n    \n#     def pred_loop(val_loader):\n#         preds = []\n\n#         with torch.no_grad():\n#             for x, y_batch in val_loader:                \n#                 y_pred = model(x.to(device))\n#                 preds += list(probs.flatten())\n#         return preds\n\n#     preds.append(val_loop(para_loader.per_device_loader(device)))\n    \n\n    \n#     with torch.no_grad():\n#         for x, _ in tqdm(loader):\n#             probs = torch.sigmoid(model(x.to(device))).detach().cpu().numpy()\n#             preds = np.concatenate([preds, probs])\n            \n#     return preds","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def predict(model, dataset, batch_size=64):\n    \"\"\"\n    Usual predict torch function\n    \n    Arguments:\n        model {torch model} -- Model to predict with\n        dataset {torch dataset} -- Dataset to get predictions from\n    \n    Keyword Arguments:\n        batch_size {int} -- Batch size (default: {32})\n    \n    Returns:\n        numpy array -- Predictions\n    \"\"\"\n    \n    device = xm.xla_device()\n    model = model.to(device)\n    model.eval()\n    \n    preds = np.empty((0, 1))\n    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n\n    with torch.no_grad():\n        for x, _ in tqdm(loader):\n            probs = torch.sigmoid(model(x.to(device))).detach().cpu().numpy()\n            preds = np.concatenate([preds, probs])\n            \n    return preds","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = load_model_weights(model, 'model.pt')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pred_val = predict(model, val_dataset)\ndf_val['pred'] = pred_val","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for language in df_val['lang'].unique():\n    lang_score = roc_auc_score(\n        df_val[df_val['lang'] == language]['toxic'], \n        df_val[df_val['lang']  == language]['pred']\n    )\n    print(f'AUC for language {language}: {lang_score:.4f}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"score = roc_auc_score(df_val['toxic'], pred_val)\nprint(f'Scored {score:.4f} on validation data')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pred_test = predict(model, test_dataset)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sub = pd.read_csv(DATA_PATH + 'sample_submission.csv')\nsub['toxic'] = pred_test\nsub.to_csv('submission.csv', index=False)\nsub.head()","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}