roberta-base-acc(0.715).ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "_cell_guid": "79c7e3d0-c299-4dcb-8224-4455121ee9b0",
    "_uuid": "d629ff2d2480ee46fbb7e2d37f6b5fab8052498a"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "import warnings\n",
    "import random\n",
    "import torch \n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from sklearn.model_selection import StratifiedKFold\n",
    "import tokenizers\n",
    "from transformers import RobertaModel, RobertaConfig\n",
    "from tqdm import tqdm\n",
    "\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "seed=42"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TweetDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, df, max_len=96):\n",
    "        self.df = df\n",
    "        self.max_len = max_len\n",
    "        self.labeled = 'selected_text' in df\n",
    "        self.tokenizer = tokenizers.ByteLevelBPETokenizer(\n",
    "            vocab_file='../input/roberta-base/vocab.json', \n",
    "            merges_file='../input/roberta-base/merges.txt', \n",
    "            lowercase=True,\n",
    "            add_prefix_space=True)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        data = {}\n",
    "        row = self.df.iloc[index]\n",
    "        \n",
    "        ids, masks, tweet, offsets = self.get_input_data(row)\n",
    "        data['ids'] = ids\n",
    "        data['masks'] = masks\n",
    "        data['tweet'] = tweet\n",
    "        data['offsets'] = offsets\n",
    "        \n",
    "        if self.labeled:\n",
    "            start_idx, end_idx = self.get_target_idx(row, tweet, offsets)\n",
    "            data['start_idx'] = start_idx\n",
    "            data['end_idx'] = end_idx\n",
    "        \n",
    "        return data\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.df)\n",
    "    \n",
    "    def get_input_data(self, row):\n",
    "        tweet = \" \" + \" \".join(row.text.lower().split())\n",
    "        encoding = self.tokenizer.encode(tweet)\n",
    "        sentiment_id = self.tokenizer.encode(row.sentiment).ids\n",
    "        ids = [0] + sentiment_id + [2, 2] + encoding.ids + [2]\n",
    "        offsets = [(0, 0)] * 4 + encoding.offsets + [(0, 0)]\n",
    "                \n",
    "        pad_len = self.max_len - len(ids)\n",
    "        if pad_len > 0:\n",
    "            ids += [1] * pad_len\n",
    "            offsets += [(0, 0)] * pad_len\n",
    "        \n",
    "        ids = torch.tensor(ids)\n",
    "        masks = torch.where(ids != 1, torch.tensor(1), torch.tensor(0))\n",
    "        offsets = torch.tensor(offsets)\n",
    "        \n",
    "        return ids, masks, tweet, offsets\n",
    "        \n",
    "    def get_target_idx(self, row, tweet, offsets):\n",
    "        selected_text = \" \" +  \" \".join(row.selected_text.lower().split())\n",
    "\n",
    "        len_st = len(selected_text) - 1\n",
    "        idx0 = None\n",
    "        idx1 = None\n",
    "\n",
    "        for ind in (i for i, e in enumerate(tweet) if e == selected_text[1]):\n",
    "            if \" \" + tweet[ind: ind+len_st] == selected_text:\n",
    "                idx0 = ind\n",
    "                idx1 = ind + len_st - 1\n",
    "                break\n",
    "\n",
    "        char_targets = [0] * len(tweet)\n",
    "        if idx0 != None and idx1 != None:\n",
    "            for ct in range(idx0, idx1 + 1):\n",
    "                char_targets[ct] = 1\n",
    "\n",
    "        target_idx = []\n",
    "        for j, (offset1, offset2) in enumerate(offsets):\n",
    "            if sum(char_targets[offset1: offset2]) > 0:\n",
    "                target_idx.append(j)\n",
    "\n",
    "        start_idx = target_idx[0]\n",
    "        end_idx = target_idx[-1]\n",
    "        \n",
    "        return start_idx, end_idx\n",
    "        \n",
    "def get_train_val_loaders(df, train_idx, val_idx, batch_size=8):\n",
    "    train_df = df.iloc[train_idx]\n",
    "    val_df = df.iloc[val_idx]\n",
    "\n",
    "    train_loader = torch.utils.data.DataLoader(\n",
    "        TweetDataset(train_df), \n",
    "        batch_size=batch_size, \n",
    "        shuffle=True, \n",
    "        num_workers=2,\n",
    "        drop_last=True)\n",
    "\n",
    "    val_loader = torch.utils.data.DataLoader(\n",
    "        TweetDataset(val_df), \n",
    "        batch_size=batch_size, \n",
    "        shuffle=False, \n",
    "        num_workers=2)\n",
    "\n",
    "    dataloaders_dict = {\"train\": train_loader, \"val\": val_loader}\n",
    "\n",
    "    return dataloaders_dict\n",
    "\n",
    "def get_test_loader(df, batch_size=32):\n",
    "    loader = torch.utils.data.DataLoader(\n",
    "        TweetDataset(df), \n",
    "        batch_size=batch_size, \n",
    "        shuffle=False, \n",
    "        num_workers=2)    \n",
    "    return loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TweetModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(TweetModel, self).__init__()\n",
    "        \n",
    "        config = RobertaConfig.from_pretrained(\n",
    "            '../input/roberta-base/config.json', output_hidden_states=True)    \n",
    "        self.roberta = RobertaModel.from_pretrained(\n",
    "            '../input/roberta-base/pytorch_model.bin', config=config)\n",
    "        self.dropout = nn.Dropout(0.5)\n",
    "        self.conv1 = nn.Conv1d(in_channels=768, out_channels=768, kernel_size=3, padding=1)\n",
    "        self.conv2 = nn.Conv1d(in_channels=768, out_channels=64, kernel_size=3, padding=1)\n",
    "        self.conv3 = nn.Conv1d(in_channels=768, out_channels=768, kernel_size=3, padding=1)\n",
    "        self.conv4 = nn.Conv1d(in_channels=768, out_channels=64, kernel_size=3, padding=1)\n",
    "        self.fc1 = nn.Linear(64, 1)\n",
    "        self.fc2 = nn.Linear(64, 1)\n",
    "        nn.init.normal_(self.fc1.weight, std=0.02)\n",
    "        nn.init.normal_(self.fc1.bias, 0)\n",
    "        nn.init.normal_(self.fc2.weight, std=0.02)\n",
    "        nn.init.normal_(self.fc2.bias, 0)\n",
    "        nn.init.xavier_normal_(self.conv1.weight)\n",
    "        nn.init.xavier_normal_(self.conv2.weight)\n",
    "        nn.init.xavier_normal_(self.conv3.weight)\n",
    "        nn.init.xavier_normal_(self.conv4.weight)\n",
    "\n",
    "    def forward(self, input_ids, attention_mask):\n",
    "        _, _, hs = self.roberta(input_ids, attention_mask)\n",
    "         \n",
    "        x = torch.stack([hs[-1], hs[-2], hs[-3], hs[-4]])\n",
    "        x = torch.mean(x, 0)\n",
    "        x = self.dropout(x)\n",
    "        \n",
    "        x1=x.permute(0,2,1)\n",
    "        x1 = self.conv1(x1)\n",
    "        x1 = self.conv2(F.leaky_relu(x1))\n",
    "        x1 = x1.permute(0,2,1)\n",
    "        x1 = self.fc1(x1)\n",
    "        start_logits = x1.squeeze(-1)\n",
    "        \n",
    "        x2 = x.permute(0,2,1)\n",
    "        x2 = self.conv3(x2)\n",
    "        x2 = self.conv4(F.leaky_relu(x2))\n",
    "        x2 = x2.permute(0,2,1)\n",
    "        x2 = self.fc2(x2)\n",
    "        end_logits = x2.squeeze(-1)\n",
    "                \n",
    "        return start_logits, end_logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_fn(start_logits, end_logits, start_positions, end_positions):\n",
    "    ce_loss = nn.CrossEntropyLoss()\n",
    "    start_loss = ce_loss(start_logits, start_positions)\n",
    "    end_loss = ce_loss(end_logits, end_positions)    \n",
    "    total_loss = start_loss + end_loss\n",
    "    return total_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_selected_text(text, start_idx, end_idx, offsets):\n",
    "    selected_text = \"\"\n",
    "    for ix in range(start_idx, end_idx + 1):\n",
    "        selected_text += text[offsets[ix][0]: offsets[ix][1]]\n",
    "        if (ix + 1) < len(offsets) and offsets[ix][1] < offsets[ix + 1][0]:\n",
    "            selected_text += \" \"\n",
    "    return selected_text\n",
    "\n",
    "def jaccard(str1, str2): \n",
    "    a = set(str1.lower().split()) \n",
    "    b = set(str2.lower().split())\n",
    "    c = a.intersection(b)\n",
    "    return float(len(c)) / (len(a) + len(b) - len(c))\n",
    "\n",
    "def compute_jaccard_score(text, start_idx, end_idx, start_logits, end_logits, offsets):\n",
    "    start_pred = np.argmax(start_logits)\n",
    "    end_pred = np.argmax(end_logits)\n",
    "    if start_pred > end_pred:\n",
    "        pred = text\n",
    "    else:\n",
    "        pred = get_selected_text(text, start_pred, end_pred, offsets)\n",
    "        \n",
    "    true = get_selected_text(text, start_idx, end_idx, offsets)\n",
    "    \n",
    "    return jaccard(true, pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, dataloaders_dict, criterion, optimizer, num_epochs, filename):\n",
    "    model.cuda()\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        for phase in ['train', 'val']:\n",
    "            if phase == 'train':\n",
    "                model.train()\n",
    "            else:\n",
    "                model.eval()\n",
    "\n",
    "            epoch_loss = 0.0\n",
    "            epoch_jaccard = 0.0\n",
    "            \n",
    "            for data in (dataloaders_dict[phase]):\n",
    "                ids = data['ids'].cuda()\n",
    "                masks = data['masks'].cuda()\n",
    "                tweet = data['tweet']\n",
    "                offsets = data['offsets'].numpy()\n",
    "                start_idx = data['start_idx'].cuda()\n",
    "                end_idx = data['end_idx'].cuda()\n",
    "\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "                with torch.set_grad_enabled(phase == 'train'):\n",
    "\n",
    "                    start_logits, end_logits = model(ids, masks)\n",
    "\n",
    "                    loss = criterion(start_logits, end_logits, start_idx, end_idx)\n",
    "                    \n",
    "                    if phase == 'train':\n",
    "                        loss.backward()\n",
    "                        optimizer.step()\n",
    "\n",
    "                    epoch_loss += loss.item() * len(ids)\n",
    "                    \n",
    "                    start_idx = start_idx.cpu().detach().numpy()\n",
    "                    end_idx = end_idx.cpu().detach().numpy()\n",
    "                    start_logits = torch.softmax(start_logits, dim=1).cpu().detach().numpy()\n",
    "                    end_logits = torch.softmax(end_logits, dim=1).cpu().detach().numpy()\n",
    "                    \n",
    "                    for i in range(len(ids)):                        \n",
    "                        jaccard_score = compute_jaccard_score(\n",
    "                            tweet[i],\n",
    "                            start_idx[i],\n",
    "                            end_idx[i],\n",
    "                            start_logits[i], \n",
    "                            end_logits[i], \n",
    "                            offsets[i])\n",
    "                        epoch_jaccard += jaccard_score\n",
    "                    \n",
    "            epoch_loss = epoch_loss / len(dataloaders_dict[phase].dataset)\n",
    "            epoch_jaccard = epoch_jaccard / len(dataloaders_dict[phase].dataset)\n",
    "            \n",
    "            print('Epoch {}/{} | {:^5} | Loss: {:.4f} | Jaccard: {:.4f}'.format(\n",
    "                epoch + 1, num_epochs, phase, epoch_loss, epoch_jaccard))\n",
    "        \n",
    "        if epoch<2:\n",
    "            jaccard_check=epoch_jaccard\n",
    "            torch.save(model.state_dict(), filename)\n",
    "            print('Saving model')\n",
    "        elif epoch_jaccard>jaccard_check:\n",
    "            jaccard_check=epoch_jaccard\n",
    "            torch.save(model.state_dict(), filename)\n",
    "            print('Saving model')\n",
    "        else:\n",
    "            print('Stopping iteration')\n",
    "            break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epochs = 10\n",
    "batch_size = 32\n",
    "skf = StratifiedKFold(n_splits=10, shuffle=True, random_state=seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fold: 1\n",
      "Epoch 1/10 | train | Loss: 2.0403 | Jaccard: 0.6811\n",
      "Epoch 1/10 |  val  | Loss: 1.6423 | Jaccard: 0.7140\n",
      "Saving model\n",
      "Epoch 2/10 | train | Loss: 1.5329 | Jaccard: 0.7277\n",
      "Epoch 2/10 |  val  | Loss: 1.6114 | Jaccard: 0.7183\n",
      "Saving model\n",
      "Epoch 3/10 | train | Loss: 1.3671 | Jaccard: 0.7507\n",
      "Epoch 3/10 |  val  | Loss: 1.6579 | Jaccard: 0.7092\n",
      "Stopping iteration\n",
      "Fold: 2\n",
      "Epoch 1/10 | train | Loss: 2.0941 | Jaccard: 0.6729\n",
      "Epoch 1/10 |  val  | Loss: 1.7050 | Jaccard: 0.7088\n",
      "Saving model\n",
      "Epoch 2/10 | train | Loss: 1.5713 | Jaccard: 0.7229\n",
      "Epoch 2/10 |  val  | Loss: 1.6644 | Jaccard: 0.7146\n",
      "Saving model\n",
      "Epoch 3/10 | train | Loss: 1.4246 | Jaccard: 0.7416\n",
      "Epoch 3/10 |  val  | Loss: 1.6677 | Jaccard: 0.7111\n",
      "Stopping iteration\n",
      "Fold: 3\n",
      "Epoch 1/10 | train | Loss: 2.0559 | Jaccard: 0.6805\n",
      "Epoch 1/10 |  val  | Loss: 1.6905 | Jaccard: 0.7135\n",
      "Saving model\n",
      "Epoch 2/10 | train | Loss: 1.5615 | Jaccard: 0.7223\n",
      "Epoch 2/10 |  val  | Loss: 1.6464 | Jaccard: 0.7052\n",
      "Saving model\n",
      "Epoch 3/10 | train | Loss: 1.3946 | Jaccard: 0.7460\n",
      "Epoch 3/10 |  val  | Loss: 1.6786 | Jaccard: 0.7116\n",
      "Saving model\n",
      "Epoch 4/10 | train | Loss: 1.2205 | Jaccard: 0.7708\n",
      "Epoch 4/10 |  val  | Loss: 1.8014 | Jaccard: 0.7101\n",
      "Stopping iteration\n",
      "Fold: 4\n",
      "Epoch 1/10 | train | Loss: 2.1433 | Jaccard: 0.6661\n",
      "Epoch 1/10 |  val  | Loss: 1.5540 | Jaccard: 0.7218\n",
      "Saving model\n",
      "Epoch 2/10 | train | Loss: 1.5867 | Jaccard: 0.7187\n",
      "Epoch 2/10 |  val  | Loss: 1.5694 | Jaccard: 0.7226\n",
      "Saving model\n",
      "Epoch 3/10 | train | Loss: 1.4407 | Jaccard: 0.7410\n",
      "Epoch 3/10 |  val  | Loss: 1.6114 | Jaccard: 0.7283\n",
      "Saving model\n",
      "Epoch 4/10 | train | Loss: 1.3157 | Jaccard: 0.7561\n",
      "Epoch 4/10 |  val  | Loss: 1.6296 | Jaccard: 0.7156\n",
      "Stopping iteration\n",
      "Fold: 5\n",
      "Epoch 1/10 | train | Loss: 2.0814 | Jaccard: 0.6710\n",
      "Epoch 1/10 |  val  | Loss: 1.6346 | Jaccard: 0.7096\n",
      "Saving model\n",
      "Epoch 2/10 | train | Loss: 1.5547 | Jaccard: 0.7254\n",
      "Epoch 2/10 |  val  | Loss: 1.6109 | Jaccard: 0.7126\n",
      "Saving model\n",
      "Epoch 3/10 | train | Loss: 1.3970 | Jaccard: 0.7480\n",
      "Epoch 3/10 |  val  | Loss: 1.6099 | Jaccard: 0.7155\n",
      "Saving model\n",
      "Epoch 4/10 | train | Loss: 1.2343 | Jaccard: 0.7722\n",
      "Epoch 4/10 |  val  | Loss: 1.6818 | Jaccard: 0.7099\n",
      "Stopping iteration\n",
      "Fold: 6\n",
      "Epoch 1/10 | train | Loss: 2.0348 | Jaccard: 0.6824\n",
      "Epoch 1/10 |  val  | Loss: 1.6342 | Jaccard: 0.7159\n",
      "Saving model\n",
      "Epoch 2/10 | train | Loss: 1.5574 | Jaccard: 0.7239\n",
      "Epoch 2/10 |  val  | Loss: 1.6357 | Jaccard: 0.7213\n",
      "Saving model\n",
      "Epoch 3/10 | train | Loss: 1.3900 | Jaccard: 0.7494\n",
      "Epoch 3/10 |  val  | Loss: 1.6861 | Jaccard: 0.7141\n",
      "Stopping iteration\n",
      "Fold: 7\n",
      "Epoch 1/10 | train | Loss: 2.1133 | Jaccard: 0.6677\n",
      "Epoch 1/10 |  val  | Loss: 1.6116 | Jaccard: 0.7134\n",
      "Saving model\n",
      "Epoch 2/10 | train | Loss: 1.5687 | Jaccard: 0.7243\n",
      "Epoch 2/10 |  val  | Loss: 1.5774 | Jaccard: 0.7145\n",
      "Saving model\n",
      "Epoch 3/10 | train | Loss: 1.4159 | Jaccard: 0.7447\n",
      "Epoch 3/10 |  val  | Loss: 1.6661 | Jaccard: 0.7172\n",
      "Saving model\n",
      "Epoch 4/10 | train | Loss: 1.2662 | Jaccard: 0.7687\n",
      "Epoch 4/10 |  val  | Loss: 1.7146 | Jaccard: 0.7138\n",
      "Stopping iteration\n",
      "Fold: 8\n",
      "Epoch 1/10 | train | Loss: 2.0499 | Jaccard: 0.6786\n",
      "Epoch 1/10 |  val  | Loss: 1.6384 | Jaccard: 0.7095\n",
      "Saving model\n",
      "Epoch 2/10 | train | Loss: 1.5621 | Jaccard: 0.7226\n",
      "Epoch 2/10 |  val  | Loss: 1.6087 | Jaccard: 0.7211\n",
      "Saving model\n",
      "Epoch 3/10 | train | Loss: 1.4162 | Jaccard: 0.7414\n",
      "Epoch 3/10 |  val  | Loss: 1.6383 | Jaccard: 0.7257\n",
      "Saving model\n",
      "Epoch 4/10 | train | Loss: 1.2616 | Jaccard: 0.7675\n",
      "Epoch 4/10 |  val  | Loss: 1.7547 | Jaccard: 0.7158\n",
      "Stopping iteration\n",
      "Fold: 9\n",
      "Epoch 1/10 | train | Loss: 2.0530 | Jaccard: 0.6792\n",
      "Epoch 1/10 |  val  | Loss: 1.7135 | Jaccard: 0.7170\n",
      "Saving model\n",
      "Epoch 2/10 | train | Loss: 1.5526 | Jaccard: 0.7243\n",
      "Epoch 2/10 |  val  | Loss: 1.6778 | Jaccard: 0.7152\n",
      "Saving model\n",
      "Epoch 3/10 | train | Loss: 1.4223 | Jaccard: 0.7437\n",
      "Epoch 3/10 |  val  | Loss: 1.6283 | Jaccard: 0.7195\n",
      "Saving model\n",
      "Epoch 4/10 | train | Loss: 1.2628 | Jaccard: 0.7680\n",
      "Epoch 4/10 |  val  | Loss: 1.7715 | Jaccard: 0.7142\n",
      "Stopping iteration\n",
      "Fold: 10\n",
      "Epoch 1/10 | train | Loss: 2.0549 | Jaccard: 0.6759\n",
      "Epoch 1/10 |  val  | Loss: 1.6515 | Jaccard: 0.7041\n",
      "Saving model\n",
      "Epoch 2/10 | train | Loss: 1.5510 | Jaccard: 0.7240\n",
      "Epoch 2/10 |  val  | Loss: 1.6182 | Jaccard: 0.7116\n",
      "Saving model\n",
      "Epoch 3/10 | train | Loss: 1.3933 | Jaccard: 0.7469\n",
      "Epoch 3/10 |  val  | Loss: 1.6195 | Jaccard: 0.7181\n",
      "Saving model\n",
      "Epoch 4/10 | train | Loss: 1.2194 | Jaccard: 0.7759\n",
      "Epoch 4/10 |  val  | Loss: 1.8749 | Jaccard: 0.7041\n",
      "Stopping iteration\n",
      "CPU times: user 1h 49min 41s, sys: 45min 41s, total: 2h 35min 23s\n",
      "Wall time: 2h 41min 39s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "train_df = pd.read_csv('../input/tweet-sentiment-extraction/train.csv')\n",
    "train_df['text'] = train_df['text'].astype(str)\n",
    "train_df['selected_text'] = train_df['selected_text'].astype(str)\n",
    "\n",
    "for fold, (train_idx, val_idx) in enumerate(skf.split(train_df, train_df.sentiment), start=1): \n",
    "    print(f'Fold: {fold}')\n",
    "\n",
    "    model = TweetModel()\n",
    "    optimizer = optim.AdamW(model.parameters(), lr=4e-5, betas=(0.9, 0.999))\n",
    "    criterion = loss_fn    \n",
    "    dataloaders_dict = get_train_val_loaders(train_df, train_idx, val_idx, batch_size)\n",
    "\n",
    "    train_model(\n",
    "        model, \n",
    "        dataloaders_dict,\n",
    "        criterion, \n",
    "        optimizer, \n",
    "        num_epochs,\n",
    "        f'roberta_fold{fold}.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 56s, sys: 42.6 s, total: 2min 38s\n",
      "Wall time: 2min 48s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "test_df = pd.read_csv('../input/tweet-sentiment-extraction/test.csv')\n",
    "test_df['text'] = test_df['text'].astype(str)\n",
    "test_loader = get_test_loader(test_df)\n",
    "predictions = []\n",
    "models = []\n",
    "for fold in range(skf.n_splits):\n",
    "    model = TweetModel()\n",
    "    model.cuda()\n",
    "    model.load_state_dict(torch.load(f'roberta_fold{fold+1}.pth'))\n",
    "    model.eval()\n",
    "    models.append(model)\n",
    "\n",
    "for data in test_loader:\n",
    "    ids = data['ids'].cuda()\n",
    "    masks = data['masks'].cuda()\n",
    "    tweet = data['tweet']\n",
    "    offsets = data['offsets'].numpy()\n",
    "\n",
    "    start_logits = []\n",
    "    end_logits = []\n",
    "    for model in models:\n",
    "        with torch.no_grad():\n",
    "            output = model(ids, masks)\n",
    "            start_logits.append(torch.softmax(output[0], dim=1).cpu().detach().numpy())\n",
    "            end_logits.append(torch.softmax(output[1], dim=1).cpu().detach().numpy())\n",
    "\n",
    "    start_logits = np.mean(start_logits, axis=0)\n",
    "    end_logits = np.mean(end_logits, axis=0)\n",
    "    for i in range(len(ids)):    \n",
    "        start_pred = np.argmax(start_logits[i])\n",
    "        end_pred = np.argmax(end_logits[i])\n",
    "        if start_pred > end_pred:\n",
    "            pred = tweet[i]\n",
    "        else:\n",
    "            pred = get_selected_text(tweet[i], start_pred, end_pred, offsets[i])\n",
    "        predictions.append(pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>textID</th>\n",
       "      <th>selected_text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>f87dea47db</td>\n",
       "      <td>last session of the day http://twitpic.com/67ezh</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>96d74cb729</td>\n",
       "      <td>exciting</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>eee518ae67</td>\n",
       "      <td>such a shame!</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>01082688c6</td>\n",
       "      <td>happy bday!</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>33987a8ee5</td>\n",
       "      <td>i like it!!</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       textID                                      selected_text\n",
       "0  f87dea47db   last session of the day http://twitpic.com/67ezh\n",
       "1  96d74cb729                                           exciting\n",
       "2  eee518ae67                                      such a shame!\n",
       "3  01082688c6                                        happy bday!\n",
       "4  33987a8ee5                                        i like it!!"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sub_df = pd.read_csv('../input/tweet-sentiment-extraction/sample_submission.csv')\n",
    "sub_df['selected_text'] = predictions\n",
    "sub_df['selected_text'] = sub_df['selected_text'].apply(lambda x: x.replace('!!!!', '!') if len(x.split())==1 else x)\n",
    "sub_df['selected_text'] = sub_df['selected_text'].apply(lambda x: x.replace('..', '.') if len(x.split())==1 else x)\n",
    "sub_df['selected_text'] = sub_df['selected_text'].apply(lambda x: x.replace('...', '.') if len(x.split())==1 else x)\n",
    "sub_df.to_csv('submission.csv', index=False)\n",
    "sub_df.head()"
   ]
  }
 ],
 "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}