jgs_bert_inference.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import, division, print_function\n",
    "\n",
    "import os, sys, re, gc, pickle, operator, shutil, copy\n",
    "import time, datetime\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.model_selection import KFold, StratifiedKFold\n",
    "from sklearn.metrics import accuracy_score, roc_auc_score\n",
    "from sklearn.utils import shuffle\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import TensorDataset, Dataset, DataLoader, Sampler\n",
    "\n",
    "from tqdm import tqdm, tqdm_notebook\n",
    "from IPython.core.interactiveshell import InteractiveShell\n",
    "InteractiveShell.ast_node_interactivity = \"all\"\n",
    "import warnings\n",
    "warnings.filterwarnings(action='once')\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.6/importlib/_bootstrap.py:219: ImportWarning: can't resolve package from __spec__ or __package__, falling back on __name__ and __path__\n",
      "  return f(*args, **kwds)\n",
      "/opt/conda/lib/python3.6/importlib/_bootstrap.py:219: RuntimeWarning: numpy.ufunc size changed, may indicate binary incompatibility. Expected 192 from C header, got 216 from PyObject\n",
      "  return f(*args, **kwds)\n"
     ]
    }
   ],
   "source": [
    "from pytorch_pretrained_bert import convert_tf_checkpoint_to_pytorch\n",
    "from pytorch_pretrained_bert import BertTokenizer, BertForSequenceClassification, BertAdam, BertConfig\n",
    "from pytorch_pretrained_bert.modeling import BertModel, BertPreTrainedModel\n",
    "device = torch.device('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "WORK_DIR = \"./\"\n",
    "DATA_DIR = \"../input/jigsaw-unintended-bias-in-toxicity-classification/\"\n",
    "BERT_MODEL_PATH = '../input/bert-pretrained-models/uncased_l-12_h-768_a-12/uncased_L-12_H-768_A-12/'\n",
    "CONFIG_PATH = '../input/jgs-bert-0/'\n",
    "MODEL0_PATH = '../input/jgs-bert-0/'\n",
    "MODEL1_PATH = '../input/jgs-bert-1/'\n",
    "MODEL2_PATH = '../input/jgs-bert-2/'\n",
    "MODEL3_PATH = '../input/jgs-bert-3/'\n",
    "MODEL4_PATH = '../input/jgs-bert-4/'\n",
    "\n",
    "SEED = 2019\n",
    "\n",
    "MAX_SEQUENCE_LENGTH = 220\n",
    "batch_size = 16\n",
    "\n",
    "identity_columns = [\n",
    "    'male', 'female', 'homosexual_gay_or_lesbian', 'christian', 'jewish',\n",
    "    'muslim', 'black', 'white', 'psychiatric_or_mental_illness']\n",
    "\n",
    "aux_columns = ['severe_toxicity', 'obscene', 'threat', 'insult', 'identity_attack', 'sexual_explicit']\n",
    "\n",
    "label_column = 'target'\n",
    "pred_column = 'prediction'\n",
    "text_column = 'comment_text'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Seed for randomness in pytorch\n",
    "def seed_torch(seed=SEED):\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    \n",
    "# Convert target and identity columns to booleans\n",
    "def convert_dataframe_to_bool(df):\n",
    "    def convert_to_bool(df, col_name):\n",
    "        df[col_name] = np.where(df[col_name] >= 0.5, True, False)\n",
    "\n",
    "    bool_df = df.copy()\n",
    "    for col in [label_column] + identity_columns + aux_columns:\n",
    "        convert_to_bool(bool_df, col)\n",
    "    return bool_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tokenizing the lines to BERT format\n",
    "# Thanks to https://www.kaggle.com/httpwwwfszyc/bert-in-keras-taming\n",
    "def convert_lines(texts, max_seq_length, tokenizer):\n",
    "    max_seq_length -= 2\n",
    "    all_tokens = []\n",
    "\n",
    "    for text in texts:\n",
    "        tokens = tokenizer.tokenize(text)\n",
    "        if len(tokens) > max_seq_length:\n",
    "            tokens = tokens[-max_seq_length:]\n",
    "        one_token = tokenizer.convert_tokens_to_ids([\"[CLS]\"]+tokens+[\"[SEP]\"])\n",
    "        all_tokens.append(one_token)\n",
    "\n",
    "    return np.array(all_tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare dataset and dataloader\n",
    "\n",
    "class Toxic_comments(Dataset):\n",
    "\n",
    "    def __init__(self, tokenized_comments, targets=None, split=None, maxlen=256):\n",
    "        self.comments = tokenized_comments\n",
    "        self.targets = targets\n",
    "        self.split = split\n",
    "        assert self.split in {'train', 'valid', 'test'}\n",
    "        self.maxlen = maxlen\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        comment = self.comments[index]\n",
    "        if self.targets is not None:\n",
    "            target = self.targets[index]\n",
    "            return comment, torch.FloatTensor(target)\n",
    "        else:\n",
    "            return comment\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.comments)\n",
    "\n",
    "    def get_lens(self):\n",
    "        lengths = np.fromiter(\n",
    "            ((min(self.maxlen, len(seq))) for seq in self.comments),\n",
    "            dtype=np.int32)\n",
    "        return lengths\n",
    "\n",
    "    def collate_fn(self, batch):\n",
    "        \"\"\"\n",
    "        Collate function for sequence bucketing\n",
    "        Note: this need not be defined in this Class, can be standalone.\n",
    "\n",
    "        :param batch: an iterable of N sets from __getitem__()\n",
    "        :return: a tensor of comments, and targets\n",
    "        \"\"\"\n",
    "\n",
    "        if self.split in ('train', 'valid'):\n",
    "            comments, targets = zip(*batch)\n",
    "        else:\n",
    "            comments = batch\n",
    "\n",
    "        lengths = [len(c) for c in comments]\n",
    "        maxlen = max(lengths)\n",
    "        padded_comments = []\n",
    "        for i, c in enumerate(comments):\n",
    "            padded_comments.append(c+[0]*(maxlen - lengths[i]))\n",
    "\n",
    "        if self.split in ('train', 'valid'):\n",
    "            return torch.LongTensor(padded_comments), torch.stack(targets)\n",
    "        else:\n",
    "            return torch.LongTensor(padded_comments)\n",
    "\n",
    "\n",
    "class BucketSampler(Sampler):\n",
    "\n",
    "    def __init__(self, data_source, sort_lens, bucket_size=None, batch_size=1024, shuffle_data=True):\n",
    "        super().__init__(data_source)\n",
    "        self.shuffle = shuffle_data\n",
    "        self.batch_size = batch_size\n",
    "        self.sort_lens = sort_lens\n",
    "        self.bucket_size = bucket_size if bucket_size is not None else len(sort_lens)\n",
    "        self.weights = None\n",
    "\n",
    "        if not shuffle_data:\n",
    "            self.index = self.prepare_buckets()\n",
    "        else:\n",
    "            self.index = None\n",
    "\n",
    "    def set_weights(self, weights):\n",
    "        assert weights >= 0\n",
    "        total = np.sum(weights)\n",
    "        if total != 1:\n",
    "            weights = weights / total\n",
    "        self.weights = weights\n",
    "\n",
    "    def __iter__(self):\n",
    "        indices = None\n",
    "        if self.weights is not None:\n",
    "            total = len(self.sort_lens)\n",
    "            indices = np.random.choice(total, (total,), p=self.weights)\n",
    "        if self.shuffle:\n",
    "            self.index = self.prepare_buckets(indices)\n",
    "        return iter(self.index)\n",
    "\n",
    "    def get_reverse_indexes(self):\n",
    "        indexes = np.zeros((len(self.index),), dtype=np.int32)\n",
    "        for i, j in enumerate(self.index):\n",
    "            indexes[j] = i\n",
    "        return indexes\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.sort_lens)\n",
    "\n",
    "    def prepare_buckets(self, indices=None):\n",
    "        lengths = - self.sort_lens\n",
    "        assert self.bucket_size % self.batch_size == 0 or self.bucket_size == len(lengths)\n",
    "\n",
    "        if indices is None:\n",
    "            if self.shuffle:\n",
    "                indices = shuffle(np.arange(len(lengths), dtype=np.int32))\n",
    "                lengths = lengths[indices]\n",
    "            else:\n",
    "                indices = np.arange(len(lengths), dtype=np.int32)\n",
    "\n",
    "        #  bucket iterator\n",
    "        def divide_chunks(l, n):\n",
    "            if n == len(l):\n",
    "                yield np.arange(len(l), dtype=np.int32), l\n",
    "            else:\n",
    "                # looping till length l\n",
    "                for i in range(0, len(l), n):\n",
    "                    data = l[i:i + n]\n",
    "                    yield np.arange(i, i + len(data), dtype=np.int32), data\n",
    "\n",
    "        new_indices = []\n",
    "        extra_batch_idx = None\n",
    "        for chunk_index, chunk in divide_chunks(lengths, self.bucket_size):\n",
    "            # sort indices in bucket by descending order of length\n",
    "            indices_sorted = chunk_index[np.argsort(chunk)]\n",
    "\n",
    "            batch_idxes = []\n",
    "            for _, batch_idx in divide_chunks(indices_sorted, self.batch_size):\n",
    "                if len(batch_idx) == self.batch_size:\n",
    "                    batch_idxes.append(batch_idx.tolist())\n",
    "                else:\n",
    "                    assert extra_batch_idx is None\n",
    "                    assert batch_idx is not None\n",
    "                    extra_batch_idx = batch_idx.tolist()\n",
    "\n",
    "            # shuffling batches within buckets\n",
    "            if self.shuffle:\n",
    "                batch_idxes = shuffle(batch_idxes)\n",
    "            for batch_idx in batch_idxes:\n",
    "                new_indices.extend(batch_idx)\n",
    "\n",
    "        if extra_batch_idx is not None:\n",
    "            new_indices.extend(extra_batch_idx)\n",
    "\n",
    "        if not self.shuffle:\n",
    "            self.original_indices = np.argsort(indices_sorted).tolist()\n",
    "        return indices[new_indices]\n",
    "\n",
    "\n",
    "def prepare_loader(x, y=None, batch_size=1024, split=None):\n",
    "    assert split in {'train', 'valid', 'test'}\n",
    "    dataset = Toxic_comments(x, y, split, MAX_SEQUENCE_LENGTH)\n",
    "    if split == 'train':\n",
    "        sampler = BucketSampler(dataset, dataset.get_lens(),\n",
    "                                bucket_size=batch_size*50, batch_size=batch_size)\n",
    "        return DataLoader(dataset, batch_size=batch_size, sampler=sampler,\n",
    "                          collate_fn=dataset.collate_fn)\n",
    "    else:\n",
    "        sampler = BucketSampler(dataset, dataset.get_lens(),\n",
    "                                batch_size=batch_size, shuffle_data=False)\n",
    "        return DataLoader(dataset, batch_size=batch_size, sampler=sampler,\n",
    "                          collate_fn=dataset.collate_fn), sampler.original_indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Metrics\n",
    "\n",
    "SUBGROUP_AUC = 'subgroup_auc'\n",
    "BPSN_AUC = 'bpsn_auc'  # stands for background positive, subgroup negative\n",
    "BNSP_AUC = 'bnsp_auc'  # stands for background negative, subgroup positive\n",
    "\n",
    "def compute_auc(y_true, y_pred):\n",
    "    try:\n",
    "        return roc_auc_score(y_true, y_pred)\n",
    "    except ValueError:\n",
    "        return np.nan\n",
    "\n",
    "def compute_subgroup_auc(df, subgroup, label, pred_column):\n",
    "    subgroup_examples = df[df[subgroup]]\n",
    "    return compute_auc(subgroup_examples[label], subgroup_examples[pred_column])\n",
    "\n",
    "def compute_bpsn_auc(df, subgroup, label, pred_column):\n",
    "    \"\"\"Computes the AUC of the within-subgroup negative examples and the background positive examples.\"\"\"\n",
    "    subgroup_negative_examples = df[df[subgroup] & ~df[label]]\n",
    "    non_subgroup_positive_examples = df[~df[subgroup] & df[label]]\n",
    "    examples = subgroup_negative_examples.append(non_subgroup_positive_examples)\n",
    "    return compute_auc(examples[label], examples[pred_column])\n",
    "\n",
    "def compute_bnsp_auc(df, subgroup, label, pred_column):\n",
    "    \"\"\"Computes the AUC of the within-subgroup positive examples and the background negative examples.\"\"\"\n",
    "    subgroup_positive_examples = df[df[subgroup] & df[label]]\n",
    "    non_subgroup_negative_examples = df[~df[subgroup] & ~df[label]]\n",
    "    examples = subgroup_positive_examples.append(non_subgroup_negative_examples)\n",
    "    return compute_auc(examples[label], examples[pred_column])\n",
    "\n",
    "def compute_bias_metrics_for_model(dataset,\n",
    "                                   subgroups,\n",
    "                                   model,\n",
    "                                   label_col,\n",
    "                                   include_asegs=False):\n",
    "    \"\"\"Computes per-subgroup metrics for all subgroups and one model.\"\"\"\n",
    "    records = []\n",
    "    for subgroup in subgroups:\n",
    "        record = {\n",
    "            'subgroup': subgroup,\n",
    "            'subgroup_size': len(dataset[dataset[subgroup]])\n",
    "        }\n",
    "        record[SUBGROUP_AUC] = compute_subgroup_auc(dataset, subgroup, label_col, model)\n",
    "        record[BPSN_AUC] = compute_bpsn_auc(dataset, subgroup, label_col, model)\n",
    "        record[BNSP_AUC] = compute_bnsp_auc(dataset, subgroup, label_col, model)\n",
    "        records.append(record)\n",
    "    return pd.DataFrame(records).sort_values('subgroup_auc', ascending=True)\n",
    "\n",
    "def calculate_overall_auc(df, pred_col, label_col):\n",
    "    true_labels = df[label_col]\n",
    "    predicted_labels = df[pred_col]\n",
    "    return roc_auc_score(true_labels, predicted_labels)\n",
    "\n",
    "def power_mean(series, p):\n",
    "    total = sum(np.power(series, p))\n",
    "    return np.power(total / len(series), 1 / p)\n",
    "\n",
    "def get_final_metric(bias_df, overall_auc, POWER=-5, OVERALL_MODEL_WEIGHT=0.25):\n",
    "    bias_score = np.average([\n",
    "        power_mean(bias_df[SUBGROUP_AUC], POWER),\n",
    "        power_mean(bias_df[BPSN_AUC], POWER),\n",
    "        power_mean(bias_df[BNSP_AUC], POWER)\n",
    "    ])\n",
    "    return (OVERALL_MODEL_WEIGHT * overall_auc) + ((1 - OVERALL_MODEL_WEIGHT) * bias_score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Functions for the validation\n",
    "\n",
    "def validate(val_loader, model, val_df, val_original_indices):\n",
    "    model.eval()\n",
    "    targets, scores, losses = [], [], []\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for x, y in val_loader:\n",
    "            x, y = x.to(device=device, dtype=torch.long), y.to(device=device, dtype=torch.float)\n",
    "            score = model(x, attention_mask=(x>0))\n",
    "            targets.append((y[:,0].cpu().numpy()>=0.5).astype(int))\n",
    "            scores.append(torch.sigmoid(score[:,0]).cpu().numpy())\n",
    "\n",
    "    targets = np.concatenate(targets)\n",
    "    scores = np.concatenate(scores)\n",
    "    val_auc = roc_auc_score(targets, scores)\n",
    "    print('{\"metric\": \"Val. AUC\", \"value\": %.4f}' % (val_auc, ))\n",
    "    \n",
    "    val_scores = scores[val_original_indices]\n",
    "    val_unbias_auc = check_unbias_auc(val_df, val_scores)\n",
    "    print('{\"metric\": \"Val. Unbiased AUC\", \"value\": %.4f}' % (val_unbias_auc, ))\n",
    "    \n",
    "    return val_scores\n",
    "\n",
    "def check_unbias_auc(df, scores, print_table=True):\n",
    "    df[pred_column] = scores\n",
    "    bias_metrics_df = compute_bias_metrics_for_model(df, identity_columns, pred_column, label_column)\n",
    "    unbias_auc = get_final_metric(bias_metrics_df, calculate_overall_auc(df, pred_column, label_column))\n",
    "    if print_table:\n",
    "        print(bias_metrics_df)\n",
    "    return unbias_auc\n",
    "\n",
    "def eval_model(model, test_loader):\n",
    "    model.eval()\n",
    "    test_scores = []\n",
    "    with torch.no_grad():\n",
    "        for x in test_loader:\n",
    "            x = x.to(device=device, dtype=torch.long)\n",
    "            score = torch.sigmoid(model(x, attention_mask=(x>0))[:,0])\n",
    "            test_scores.append(score.cpu().numpy())\n",
    "    return np.concatenate(test_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_val_split(train_x, train_y):\n",
    "    kf = StratifiedKFold(n_splits=10, shuffle=True, random_state=SEED)\n",
    "    cv_indices = [(tr_idx, val_idx) for tr_idx, val_idx in kf.split(train_x, train_y)]\n",
    "    return cv_indices\n",
    "\n",
    "def load_and_preproc():\n",
    "    train_df = pd.read_csv(DATA_DIR+'train.csv')\n",
    "    test_df = pd.read_csv(DATA_DIR+'test.csv')\n",
    "    train_df[identity_columns] = train_df[identity_columns].copy().fillna(0)\n",
    "\n",
    "    sample_weights = np.ones(len(train_df))\n",
    "    sample_weights += train_df[identity_columns].values.sum(1) * 3\n",
    "    sample_weights += train_df[label_column].values * 8\n",
    "    sample_weights /= sample_weights.max()\n",
    "    train_tars = train_df[[label_column]+aux_columns+identity_columns].values\n",
    "    train_tars = np.hstack([train_tars, sample_weights[:,None]]).astype('float32')\n",
    "\n",
    "    train_df = convert_dataframe_to_bool(train_df)\n",
    "    df = train_df[[label_column]+identity_columns].copy()\n",
    "    df[label_column] = df[label_column].astype('uint8')\n",
    "\n",
    "    return train_df[text_column], test_df[text_column], train_tars, df, test_df['id']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tokenizing...\n",
      "tokenizing complete in 173 seconds.\n"
     ]
    }
   ],
   "source": [
    "train_seq, x_test, train_tars, trn_df, test_id = load_and_preproc()\n",
    "cv_indices = train_val_split(train_seq, (train_tars[:,0]>=0.5).astype(int))\n",
    "trn_idx, val_idx = cv_indices[0]\n",
    "\n",
    "print('tokenizing...')\n",
    "t0 = time.time()\n",
    "tokenizer = BertTokenizer.from_pretrained(BERT_MODEL_PATH, do_lower_case=True)\n",
    "# x_val = convert_lines(train_seq[val_idx], MAX_SEQUENCE_LENGTH, tokenizer)\n",
    "x_test = convert_lines(x_test, MAX_SEQUENCE_LENGTH, tokenizer)\n",
    "print('tokenizing complete in {:.0f} seconds.'.format(time.time()-t0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# y_val = train_tars[val_idx]\n",
    "# val_loader, val_original_indices = prepare_loader(x_val, y_val, batch_size, split='valid')\n",
    "test_loader, test_original_indices = prepare_loader(x_test, batch_size=batch_size, split='test')\n",
    "# val_df = trn_df.iloc[val_idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import gc\n",
    "del train_seq, train_tars, trn_df\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0M\n",
      "0.0M\n"
     ]
    }
   ],
   "source": [
    "print(str(torch.cuda.memory_allocated(device)/1000000 ) + 'M')\n",
    "print(str(torch.cuda.memory_cached(device)/1000000 ) + 'M')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building model...\n"
     ]
    }
   ],
   "source": [
    "# BERT configuration file for model loading\n",
    "bert_config = BertConfig(BERT_MODEL_PATH + 'bert_config.json')\n",
    "\n",
    "# model setup\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "print(\"Building model...\")\n",
    "model = BertForSequenceClassification(bert_config, num_labels=16)\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "439.111168M\n",
      "494.927872M\n"
     ]
    }
   ],
   "source": [
    "print(str(torch.cuda.memory_allocated(device)/1000000 ) + 'M')\n",
    "print(str(torch.cuda.memory_cached(device)/1000000 ) + 'M')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class WeightAvg(object):\n",
    "    def __init__(self, model):\n",
    "        self.weight_copy = {}\n",
    "        for name, param in model.named_parameters():\n",
    "            if param.requires_grad:\n",
    "                self.weight_copy[name] = param.data\n",
    "\n",
    "    def average_models(self, models):\n",
    "        nb_models = len(models)+1\n",
    "        for name in self.weight_copy.keys():\n",
    "            self.weight_copy[name] *= (1./nb_models)\n",
    "        for md in models:\n",
    "            for name, param in md.named_parameters():\n",
    "                if param.requires_grad:\n",
    "                    self.weight_copy[name] += (1./nb_models) * param.data\n",
    "\n",
    "    def set_weights(self, avg_model):\n",
    "        for name, param in avg_model.named_parameters():\n",
    "            if param.requires_grad:\n",
    "                param.data = self.weight_copy[name]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "models = torch.load(MODEL0_PATH+'epk_1_bert_models.pt')['model']\n",
    "model.load_state_dict(models[list(models.keys())[-1]])\n",
    "\n",
    "models = torch.load(MODEL1_PATH+'epk_1_1_bert_models.pt')['model']\n",
    "model_1 = copy.deepcopy(model)\n",
    "model_1.load_state_dict(models[list(models.keys())[-1]])\n",
    "\n",
    "models = torch.load(MODEL2_PATH+'epk_1_2_bert_models.pt')['model']\n",
    "model_2 = copy.deepcopy(model)\n",
    "model_2.load_state_dict(models[list(models.keys())[-1]])\n",
    "\n",
    "models = torch.load(MODEL3_PATH+'epk_1_3_bert_models.pt')['model']\n",
    "model_3 = copy.deepcopy(model)\n",
    "model_3.load_state_dict(models[list(models.keys())[-1]])\n",
    "\n",
    "models = torch.load(MODEL4_PATH+'epk_1_4_bert_models.pt')['model']\n",
    "model_4 = copy.deepcopy(model)\n",
    "model_4.load_state_dict(models[list(models.keys())[-1]])\n",
    "\n",
    "avg_model = copy.deepcopy(model)\n",
    "wa = WeightAvg(model)\n",
    "wa.average_models([model_1, model_2, model_3, model_4])\n",
    "wa.set_weights(avg_model)\n",
    "avg_model = avg_model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time = 4m 6s.\n"
     ]
    }
   ],
   "source": [
    "# test inference\n",
    "t0 = time.time()\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "wavg_test_preds = eval_model(avg_model, test_loader)[test_original_indices]\n",
    "\n",
    "time_elapsed = time.time() - t0\n",
    "print('time = {:.0f}m {:.0f}s.'.format(time_elapsed // 60, time_elapsed % 60))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# models = torch.load(MODEL_PATH+'epk_1_bert_models.pt')['model']\n",
    "# model.load_state_dict(models[list(models.keys())[-1]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # test inference\n",
    "# t0 = time.time()\n",
    "# torch.cuda.empty_cache()\n",
    "\n",
    "# # val_preds = validate(val_loader, model, val_df, val_original_indices)\n",
    "# test_preds = eval_model(model, test_loader)[test_original_indices]\n",
    "# # test_preds = []\n",
    "# # ckpt_weights = [2**e for e in range(len(models))]\n",
    "# # for state in models.values():\n",
    "# #     model.load_state_dict(state)\n",
    "# #     test_preds.append(eval_model(model, test_loader)[test_original_indices])\n",
    "\n",
    "# time_elapsed = time.time() - t0\n",
    "# print('time = {:.0f}m {:.0f}s.'.format(time_elapsed // 60, time_elapsed % 60))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# submission = pd.DataFrame.from_dict({\n",
    "#     'id': test_id,\n",
    "#     'prediction': wavg_test_preds\n",
    "# #     'prediction': np.average(test_preds, weights=ckpt_weights, axis=0)\n",
    "# })\n",
    "# submission.to_csv('submission.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# val_preds = []\n",
    "test_preds = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time = 4m 8s.\n"
     ]
    }
   ],
   "source": [
    "# validation\n",
    "t0 = time.time()\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "models = torch.load(MODEL0_PATH+'epk_1_bert_models.pt')['model']\n",
    "model.load_state_dict(models[list(models.keys())[-1]])\n",
    "# model.load_state_dict(torch.load(MODEL1_PATH+'epk_1_finetuned_bert_pytorch.bin'))\n",
    "# val_scores = validate(val_loader, model, criterion, val_df, 1, val_original_indices)\n",
    "test_scores = eval_model(model, test_loader)[test_original_indices]\n",
    "# val_preds.append(val_scores)\n",
    "test_preds.append(test_scores)\n",
    "\n",
    "time_elapsed = time.time() - t0\n",
    "print('time = {:.0f}m {:.0f}s.'.format(time_elapsed // 60, time_elapsed % 60))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2634.667008M\n",
      "4301.258752M\n"
     ]
    }
   ],
   "source": [
    "print(str(torch.cuda.memory_allocated(device)/1000000 ) + 'M')\n",
    "print(str(torch.cuda.memory_cached(device)/1000000 ) + 'M')\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time = 4m 8s.\n"
     ]
    }
   ],
   "source": [
    "# validation\n",
    "t0 = time.time()\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "models = torch.load(MODEL1_PATH+'epk_1_1_bert_models.pt')['model']\n",
    "model.load_state_dict(models[list(models.keys())[-1]])\n",
    "# model.load_state_dict(torch.load(MODEL2_PATH+'epk_2_finetuned_bert_pytorch.bin'))\n",
    "# val_scores = validate(val_loader, model, criterion, val_df, 1, val_original_indices)\n",
    "test_scores = eval_model(model, test_loader)[test_original_indices]\n",
    "# val_preds.append(val_scores)\n",
    "test_preds.append(test_scores)\n",
    "\n",
    "time_elapsed = time.time() - t0\n",
    "print('time = {:.0f}m {:.0f}s.'.format(time_elapsed // 60, time_elapsed % 60))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2634.667008M\n",
      "4301.258752M\n"
     ]
    }
   ],
   "source": [
    "print(str(torch.cuda.memory_allocated(device)/1000000 ) + 'M')\n",
    "print(str(torch.cuda.memory_cached(device)/1000000 ) + 'M')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time = 4m 8s.\n"
     ]
    }
   ],
   "source": [
    "# validation\n",
    "t0 = time.time()\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "models = torch.load(MODEL2_PATH+'epk_1_2_bert_models.pt')['model']\n",
    "model.load_state_dict(models[list(models.keys())[-1]])\n",
    "# model.load_state_dict(torch.load(MODEL3_PATH+'epk_3_finetuned_bert_pytorch.bin'))\n",
    "# val_scores = validate(val_loader, model, criterion, val_df, 1, val_original_indices)\n",
    "test_scores = eval_model(model, test_loader)[test_original_indices]\n",
    "# val_preds.append(val_scores)\n",
    "test_preds.append(test_scores)\n",
    "\n",
    "time_elapsed = time.time() - t0\n",
    "print('time = {:.0f}m {:.0f}s.'.format(time_elapsed // 60, time_elapsed % 60))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2634.667008M\n",
      "4301.258752M\n"
     ]
    }
   ],
   "source": [
    "print(str(torch.cuda.memory_allocated(device)/1000000 ) + 'M')\n",
    "print(str(torch.cuda.memory_cached(device)/1000000 ) + 'M')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time = 4m 7s.\n"
     ]
    }
   ],
   "source": [
    "# validation\n",
    "t0 = time.time()\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "models = torch.load(MODEL3_PATH+'epk_1_3_bert_models.pt')['model']\n",
    "model.load_state_dict(models[list(models.keys())[-1]])\n",
    "# model.load_state_dict(torch.load(MODEL3_PATH+'epk_3_finetuned_bert_pytorch.bin'))\n",
    "# val_scores = validate(val_loader, model, criterion, val_df, 1, val_original_indices)\n",
    "test_scores = eval_model(model, test_loader)[test_original_indices]\n",
    "# val_preds.append(val_scores)\n",
    "test_preds.append(test_scores)\n",
    "\n",
    "time_elapsed = time.time() - t0\n",
    "print('time = {:.0f}m {:.0f}s.'.format(time_elapsed // 60, time_elapsed % 60))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2634.667008M\n",
      "4301.258752M\n"
     ]
    }
   ],
   "source": [
    "print(str(torch.cuda.memory_allocated(device)/1000000 ) + 'M')\n",
    "print(str(torch.cuda.memory_cached(device)/1000000 ) + 'M')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time = 4m 7s.\n"
     ]
    }
   ],
   "source": [
    "# validation\n",
    "t0 = time.time()\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "models = torch.load(MODEL4_PATH+'epk_1_4_bert_models.pt')['model']\n",
    "model.load_state_dict(models[list(models.keys())[-1]])\n",
    "# model.load_state_dict(torch.load(MODEL3_PATH+'epk_3_finetuned_bert_pytorch.bin'))\n",
    "# val_scores = validate(val_loader, model, criterion, val_df, 1, val_original_indices)\n",
    "test_scores = eval_model(model, test_loader)[test_original_indices]\n",
    "# val_preds.append(val_scores)\n",
    "test_preds.append(test_scores)\n",
    "\n",
    "time_elapsed = time.time() - t0\n",
    "print('time = {:.0f}m {:.0f}s.'.format(time_elapsed // 60, time_elapsed % 60))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2634.667008M\n",
      "4301.258752M\n"
     ]
    }
   ],
   "source": [
    "print(str(torch.cuda.memory_allocated(device)/1000000 ) + 'M')\n",
    "print(str(torch.cuda.memory_cached(device)/1000000 ) + 'M')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# val_preds = np.mean(val_preds, 0)\n",
    "# print(check_unbias_auc(val_df.copy(), val_preds))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.6/importlib/_bootstrap.py:219: ImportWarning: can't resolve package from __spec__ or __package__, falling back on __name__ and __path__\n",
      "  return f(*args, **kwds)\n"
     ]
    }
   ],
   "source": [
    "test_preds = np.mean(test_preds, 0)\n",
    "submission = pd.DataFrame.from_dict({\n",
    "    'id': test_id,\n",
    "    'prediction': 0.85*test_preds + 0.15*wavg_test_preds\n",
    "})\n",
    "submission.to_csv('submission.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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}