CorefSeq.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5"
   },
   "outputs": [],
   "source": [
    "!pip install pytorch_pretrained_bert"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_cell_guid": "79c7e3d0-c299-4dcb-8224-4455121ee9b0",
    "_uuid": "d629ff2d2480ee46fbb7e2d37f6b5fab8052498a"
   },
   "outputs": [],
   "source": [
    "# This Python 3 environment comes with many helpful analytics libraries installed\n",
    "# It is defined by the kaggle/python docker image: https://github.com/kaggle/docker-python\n",
    "# For example, here's several helpful packages to load in \n",
    "\n",
    "import numpy as np # linear algebra\n",
    "import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n",
    "\n",
    "# Input data files are available in the \"../input/\" directory.\n",
    "# For example, running this (by clicking run or pressing Shift+Enter) will list the files in the input directory\n",
    "\n",
    "import os\n",
    "print(os.listdir(\"../input\"))\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "from bson import ObjectId\n",
    "\n",
    "import pdb\n",
    "\n",
    "from __future__ import absolute_import, division, print_function\n",
    "\n",
    "import collections\n",
    "import json\n",
    "import logging\n",
    "import math\n",
    "import os\n",
    "import random\n",
    "import sys\n",
    "from io import open\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler,\n",
    "                              TensorDataset, Dataset)\n",
    "\n",
    "from torch.utils.data.distributed import DistributedSampler\n",
    "from tqdm import tqdm, trange\n",
    "\n",
    "from pytorch_pretrained_bert.file_utils import PYTORCH_PRETRAINED_BERT_CACHE\n",
    "from pytorch_pretrained_bert.modeling import BertForQuestionAnswering, BertForMultipleChoice, BertForPreTraining, BertConfig, BertModel, WEIGHTS_NAME, CONFIG_NAME\n",
    "from pytorch_pretrained_bert.optimization import BertAdam\n",
    "from pytorch_pretrained_bert.tokenization import (BasicTokenizer,\n",
    "                                                  BertTokenizer,\n",
    "                                                  whitespace_tokenize)\n",
    "\n",
    "if sys.version_info[0] == 2:\n",
    "    import cPickle as pickle\n",
    "else:\n",
    "    import pickle\n",
    "\n",
    "logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',\n",
    "                    datefmt = '%m/%d/%Y %H:%M:%S',\n",
    "                    level = logging.INFO)\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "# Any results you write to the current directory are saved as output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import StratifiedKFold\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.metrics import log_loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_path = \"https://raw.githubusercontent.com/google-research-datasets/gap-coreference/master/gap-test.tsv\"\n",
    "dev_path = \"https://raw.githubusercontent.com/google-research-datasets/gap-coreference/master/gap-development.tsv\"\n",
    "val_path = \"https://raw.githubusercontent.com/google-research-datasets/gap-coreference/master/gap-validation.tsv\"\n",
    "\n",
    "dev_df = pd.read_csv(test_path, delimiter=\"\\t\")#pd.read_csv(dev_path, delimiter=\"\\t\")\n",
    "test_df = pd.read_csv(dev_path, delimiter=\"\\t\")#pd.read_csv(test_path, delimiter=\"\\t\")\n",
    "val_df = pd.read_csv(val_path, delimiter=\"\\t\")\n",
    "\n",
    "test_df_prod = test_df.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_class_label(a_coref, b_coref):\n",
    "    if a_coref:\n",
    "        return 0\n",
    "    elif b_coref:\n",
    "        return 1\n",
    "    return 2\n",
    "\n",
    "def get_gender(pronoun):\n",
    "    gender_mapping = {'he': 0, 'his': 0, 'him': 0, \n",
    "                      'she': 1, 'her': 1, 'hers': 1}\n",
    "    return gender_mapping.get(pronoun.lower(), 1)\n",
    "\n",
    "test_class_labels = [get_class_label(aco, bco) for aco, bco in zip(test_df['A-coref'], test_df['B-coref'])]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!conda remove -y greenlet\n",
    "!pip install allennlp\n",
    "!pip install https://github.com/ceshine/pytorch_helper_bot/archive/0.0.5.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from allennlp.modules.span_extractors import SelfAttentiveSpanExtractor, EndpointSpanExtractor\n",
    "\n",
    "from helperbot import (\n",
    "    TriangularLR, BaseBot, WeightDecayOptimizerWrapper,\n",
    "    GradualWarmupScheduler\n",
    ")\n",
    "\n",
    "import gc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Head(nn.Module):\n",
    "    \"\"\"The MLP submodule\"\"\"\n",
    "    def __init__(self, bert_hidden_size: int):\n",
    "        super().__init__()\n",
    "        self.bert_hidden_size = bert_hidden_size\n",
    "        self.span_extractor = SelfAttentiveSpanExtractor(bert_hidden_size * 3)\n",
    "        # EndpointSpanExtractor below also gives similar results\n",
    "#         self.span_extractor = EndpointSpanExtractor(\n",
    "#             bert_hidden_size * 3, \"x,y,x*y\"\n",
    "#         )\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Dropout(0.1),\n",
    "            nn.Linear(bert_hidden_size * 3 * 3, 512),           \n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 3)\n",
    "        )\n",
    "                \n",
    "    def forward(self, bert_outputs, offsets):\n",
    "        assert bert_outputs.size(2) == self.bert_hidden_size * 3\n",
    "        spans_contexts = self.span_extractor(\n",
    "            bert_outputs, \n",
    "            offsets[:, :4].reshape(-1, 2, 2)\n",
    "        )\n",
    "        spans_contexts = spans_contexts.reshape(offsets.size()[0], -1)\n",
    "        return self.fc(torch.cat([\n",
    "            spans_contexts,\n",
    "            torch.gather(\n",
    "                bert_outputs, 1,\n",
    "                offsets[:, [4]].unsqueeze(2).expand(-1, -1, self.bert_hidden_size * 3)\n",
    "            ).squeeze(1)\n",
    "        ], dim=1))\n",
    "\n",
    "class GAPDataset(Dataset):\n",
    "    \"\"\"Custom GAP Dataset class\"\"\"\n",
    "    def __init__(self, df, tokenizer, tokenize_fn, labeled=True):\n",
    "        self.labeled = labeled\n",
    "        if labeled:\n",
    "            self.y = df.target.values.astype(\"uint8\")\n",
    "        \n",
    "        self.offsets, self.tokens = [], []\n",
    "        for _, row in df.iterrows():\n",
    "            tokens, offsets = tokenize_fn(row, tokenizer)\n",
    "            self.offsets.append(offsets)\n",
    "            self.tokens.append(tokenizer.convert_tokens_to_ids(\n",
    "                [\"[CLS]\"] + tokens + [\"[SEP]\"]))\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.tokens)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        if self.labeled:\n",
    "            return self.tokens[idx], self.offsets[idx], self.y[idx]\n",
    "        return self.tokens[idx], self.offsets[idx], None\n",
    "\n",
    "class GAPModel(nn.Module):\n",
    "    \"\"\"The main model.\"\"\"\n",
    "    def __init__(self, bert_model: str, device: torch.device, use_layer: int = -1):\n",
    "        super().__init__()\n",
    "        self.device = device\n",
    "        self.use_layer = use_layer\n",
    "        if bert_model in (\"bert-base-uncased\", \"bert-base-cased\"):\n",
    "            self.bert_hidden_size = 768\n",
    "        elif bert_model in (\"bert-large-uncased\", \"bert-large-cased\"):\n",
    "            self.bert_hidden_size = 1024\n",
    "        else:\n",
    "            raise ValueError(\"Unsupported BERT model.\")\n",
    "        self.bert = BertModel.from_pretrained(bert_model).to(device)\n",
    "        self.head = Head(self.bert_hidden_size).to(device)\n",
    "    \n",
    "    def forward(self, token_tensor, offsets):\n",
    "        token_tensor = token_tensor.to(self.device)\n",
    "        bert_outputs, _ =  self.bert(\n",
    "            token_tensor, attention_mask=(token_tensor > 0).long(), \n",
    "            token_type_ids=None, output_all_encoded_layers=True)\n",
    "        concat_bert = torch.cat((bert_outputs[-1],bert_outputs[-2],bert_outputs[-3]),dim=-1)\n",
    "        #head_outputs = self.head(bert_outputs[self.use_layer], offsets.to(self.device))\n",
    "        head_outputs = self.head(concat_bert, offsets.to(self.device))\n",
    "        return head_outputs            \n",
    "    \n",
    "class GAPBot(BaseBot):\n",
    "    def __init__(self, model, train_loader, val_loader, *, optimizer, clip_grad=0,\n",
    "        avg_window=100, log_dir=\"./cache/logs/\", log_level=logging.INFO,\n",
    "        checkpoint_dir=\"./cache/model_cache/\", batch_idx=0, echo=False,\n",
    "        device=\"cuda:0\", use_tensorboard=False):\n",
    "        super().__init__(\n",
    "            model, train_loader, val_loader, \n",
    "            optimizer=optimizer, clip_grad=clip_grad,\n",
    "            log_dir=log_dir, checkpoint_dir=checkpoint_dir, \n",
    "            batch_idx=batch_idx, echo=echo,\n",
    "            device=device, use_tensorboard=use_tensorboard\n",
    "        )\n",
    "        self.criterion = torch.nn.CrossEntropyLoss()\n",
    "        self.loss_format = \"%.6f\"\n",
    "        \n",
    "    def extract_prediction(self, tensor):\n",
    "        return tensor\n",
    "    \n",
    "    def snapshot(self):\n",
    "        \"\"\"Override the snapshot method because Kaggle kernel has limited local disk space.\"\"\"\n",
    "        loss = self.eval(self.val_loader)\n",
    "        loss_str = self.loss_format % loss\n",
    "        self.logger.info(\"Snapshot loss %s\", loss_str)\n",
    "        self.logger.tb_scalars(\n",
    "            \"losses\", {\"val\": loss},  self.step)\n",
    "        target_path = (\n",
    "            self.checkpoint_dir / \"best.pth\")        \n",
    "        if not self.best_performers or (self.best_performers[0][0] > loss):\n",
    "            torch.save(self.model.state_dict(), target_path)\n",
    "            self.best_performers = [(loss, target_path, self.step)]\n",
    "        self.logger.info(\"Saving checkpoint %s...\", target_path)\n",
    "        assert Path(target_path).exists()\n",
    "        return loss\n",
    "\n",
    "class BERTSpanExtractor:\n",
    "    def __init__(self, dev_df, val_df, test_df, bert_model = 'bert-large-uncased', do_lower_case=True, learning_rate=1e-5, n_epochs=30,\n",
    "                train_batch_size=10, predict_batch_size=32):\n",
    "        self.dev_df = self.extract_target(dev_df)\n",
    "        self.val_df = self.extract_target(val_df)\n",
    "        self.test_df = self.extract_target(test_df)\n",
    "        self.do_lower_case = do_lower_case\n",
    "        self.train_batch_size = train_batch_size\n",
    "        self.predict_batch_size = predict_batch_size\n",
    "        self.bert_model = bert_model\n",
    "        self.tokenizer = BertTokenizer.from_pretrained(self.bert_model, do_lower_case=self.do_lower_case, \n",
    "                                                       never_split = (\"[UNK]\", \"[SEP]\", \"[PAD]\", \"[CLS]\", \"[MASK]\"))\n",
    "        self.learning_rate = learning_rate\n",
    "        self.n_epochs = n_epochs\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "        \n",
    "    def extract_target(self, df):\n",
    "        df['target'] = [get_class_label(aco, bco) for aco, bco in zip(df['A-coref'], df['B-coref'])]\n",
    "        df['gender'] = df['Pronoun'].transform(get_gender)\n",
    "\n",
    "        return df\n",
    "        \n",
    "    def tokenize(self, row, tokenizer):\n",
    "        break_points = sorted(\n",
    "            [\n",
    "                (\"A\", row[\"A-offset\"], row[\"A\"]),\n",
    "                (\"B\", row[\"B-offset\"], row[\"B\"]),\n",
    "                (\"P\", row[\"Pronoun-offset\"], row[\"Pronoun\"]),\n",
    "            ], key=lambda x: x[0]\n",
    "        )\n",
    "        tokens, spans, current_pos = [], {}, 0\n",
    "        for name, offset, text in break_points:\n",
    "            tokens.extend(tokenizer.tokenize(row[\"Text\"][current_pos:offset]))\n",
    "            # Make sure we do not get it wrong\n",
    "            assert row[\"Text\"][offset:offset+len(text)] == text\n",
    "            # Tokenize the target\n",
    "            tmp_tokens = tokenizer.tokenize(row[\"Text\"][offset:offset+len(text)])\n",
    "            spans[name] = [len(tokens), len(tokens) + len(tmp_tokens) - 1] # inclusive\n",
    "            tokens.extend(tmp_tokens)\n",
    "            current_pos = offset + len(text)\n",
    "        tokens.extend(tokenizer.tokenize(row[\"Text\"][current_pos:offset]))\n",
    "        assert spans[\"P\"][0] == spans[\"P\"][1]\n",
    "        return tokens, (spans[\"A\"] + spans[\"B\"] + [spans[\"P\"][0]])\n",
    "\n",
    "    def collate_examples(self, batch, truncate_len=490):\n",
    "        \"\"\"Batch preparation.\n",
    "\n",
    "        1. Pad the sequences\n",
    "        2. Transform the target.\n",
    "        \"\"\"    \n",
    "        transposed = list(zip(*batch))\n",
    "        max_len = min(\n",
    "            max((len(x) for x in transposed[0])),\n",
    "            truncate_len\n",
    "        )\n",
    "        tokens = np.zeros((len(batch), max_len), dtype=np.int64)\n",
    "        for i, row in enumerate(transposed[0]):\n",
    "            row = np.array(row[:truncate_len])\n",
    "            tokens[i, :len(row)] = row\n",
    "        token_tensor = torch.from_numpy(tokens)\n",
    "        # Offsets\n",
    "        offsets = torch.stack([\n",
    "            torch.LongTensor(x) for x in transposed[1]\n",
    "        ], dim=0) + 1 # Account for the [CLS] token\n",
    "        # Labels\n",
    "        if len(transposed) == 2 or transposed[2][0] is None:\n",
    "            return token_tensor, offsets, None\n",
    "        labels = torch.LongTensor(transposed[2])\n",
    "        return token_tensor, offsets, labels\n",
    "    \n",
    "    def run_k_fold(self):\n",
    "        test_ds = GAPDataset(self.test_df, self.tokenizer, self.tokenize, labeled=True) #not great, but it's a hack needed so this \"bot\" thing doesn't crash\n",
    "        test_loader = DataLoader(\n",
    "            test_ds,\n",
    "            collate_fn = self.collate_examples,\n",
    "            batch_size=self.predict_batch_size,\n",
    "            num_workers=2,\n",
    "            pin_memory=True,\n",
    "            shuffle=False\n",
    "        )\n",
    "\n",
    "        kfold_data = pd.concat([self.dev_df, self.val_df])\n",
    "        kf = StratifiedKFold(n_splits=5, shuffle=False, random_state=42)\n",
    "\n",
    "        val_preds, test_preds, val_ys, val_losses = [], [], [], []\n",
    "        for train_index, valid_index in kf.split(kfold_data, kfold_data[\"gender\"]):\n",
    "            print(\"=\" * 20)\n",
    "            print(f\"Fold {len(val_preds) + 1}\")\n",
    "            print(\"=\" * 20)\n",
    "            train_ds = GAPDataset(kfold_data.iloc[train_index], self.tokenizer, self.tokenize)\n",
    "            val_ds = GAPDataset(kfold_data.iloc[valid_index], self.tokenizer, self.tokenize)\n",
    "            train_loader = DataLoader(\n",
    "                train_ds,\n",
    "                collate_fn = self.collate_examples,\n",
    "                batch_size=self.train_batch_size,\n",
    "                num_workers=2,\n",
    "                pin_memory=True,\n",
    "                shuffle=True,\n",
    "                drop_last=False #True\n",
    "            )\n",
    "            val_loader = DataLoader(\n",
    "                val_ds,\n",
    "                collate_fn = self.collate_examples,\n",
    "                batch_size=self.predict_batch_size,\n",
    "                num_workers=2,\n",
    "                pin_memory=True,\n",
    "                shuffle=False\n",
    "            )\n",
    "            model = GAPModel(self.bert_model, self.device)\n",
    "            \n",
    "            def children(m):\n",
    "                return m if isinstance(m, (list, tuple)) else list(m.children())\n",
    "\n",
    "            def set_trainable_attr(m, b):\n",
    "                m.trainable = b\n",
    "                for p in m.parameters():\n",
    "                    p.requires_grad = b\n",
    "\n",
    "\n",
    "            def apply_leaf(m, f):\n",
    "                c = children(m)\n",
    "                if isinstance(m, nn.Module):\n",
    "                    f(m)\n",
    "                if len(c) > 0:\n",
    "                    for l in c:\n",
    "                        apply_leaf(l, f)\n",
    "\n",
    "\n",
    "            def set_trainable(l, b):\n",
    "                apply_leaf(l, lambda m: set_trainable_attr(m, b))\n",
    "                \n",
    "            # You can unfreeze the last layer of bert by calling set_trainable(model.bert.encoder.layer[23], True)\n",
    "            set_trainable(model.bert, False)\n",
    "            set_trainable(model.head, True)\n",
    "            for i in range(12,24):\n",
    "                set_trainable(model.bert.encoder.layer[i], True)\n",
    "\n",
    "            optimizer = torch.optim.Adam(model.parameters(), lr=self.learning_rate)\n",
    "\n",
    "            bot = GAPBot(\n",
    "                model, train_loader, val_loader,\n",
    "                optimizer=optimizer, echo=True,\n",
    "                avg_window=25\n",
    "            )\n",
    "            gc.collect()\n",
    "            steps_per_epoch = len(train_loader) \n",
    "            n_steps = steps_per_epoch * self.n_epochs\n",
    "            bot.train(\n",
    "                n_steps,\n",
    "                log_interval=steps_per_epoch // 2,\n",
    "                snapshot_interval=steps_per_epoch,\n",
    "                scheduler=TriangularLR(\n",
    "                    optimizer, 20, ratio=2, steps_per_cycle=steps_per_epoch * 100)\n",
    "            )\n",
    "            # Load the best checkpoint\n",
    "            bot.load_model(bot.best_performers[0][1])\n",
    "            bot.remove_checkpoints(keep=0)    \n",
    "            val_preds.append(torch.softmax(bot.predict(val_loader), -1).clamp(1e-4, 1-1e-4).cpu().numpy())\n",
    "            val_ys.append(kfold_data.iloc[valid_index].target.astype(\"uint8\").values)\n",
    "            val_losses.append(log_loss(val_ys[-1], val_preds[-1]))\n",
    "            bot.logger.info(\"Confirm val loss: %.4f\", val_losses[-1])\n",
    "            test_preds.append(torch.softmax(bot.predict(test_loader), -1).clamp(1e-4, 1-1e-4).cpu().numpy())\n",
    "            del model\n",
    "            \n",
    "        return val_preds, test_preds, val_losses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bert_span_extractor = BERTSpanExtractor(dev_df, val_df, test_df_prod, train_batch_size=10, n_epochs=15, bert_model='bert-large-uncased')\n",
    "bert_span_val_preds, bert_span_test_preds, bert_span_val_losses = bert_span_extractor.run_k_fold()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "span_test_probas = np.mean(bert_span_test_preds, axis=0)\n",
    "submission_df = pd.DataFrame([test_df_prod.ID, span_test_probas[:,0], span_test_probas[:,1], span_test_probas[:,2]], index=['ID', 'A', 'B', 'NEITHER']).transpose()\n",
    "\n",
    "submission_df.to_csv('stage1_span_only.csv', index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}