Visual-BERT_pair_box_tfidf-neg_focal_all_predict-all.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### imports\n",
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/u1451673/.local/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:516: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint8 = np.dtype([(\"qint8\", np.int8, 1)])\n",
      "/home/u1451673/.local/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:517: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint8 = np.dtype([(\"quint8\", np.uint8, 1)])\n",
      "/home/u1451673/.local/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:518: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint16 = np.dtype([(\"qint16\", np.int16, 1)])\n",
      "/home/u1451673/.local/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:519: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint16 = np.dtype([(\"quint16\", np.uint16, 1)])\n",
      "/home/u1451673/.local/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:520: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint32 = np.dtype([(\"qint32\", np.int32, 1)])\n",
      "/home/u1451673/.local/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:525: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  np_resource = np.dtype([(\"resource\", np.ubyte, 1)])\n",
      "/home/u1451673/.local/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:541: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint8 = np.dtype([(\"qint8\", np.int8, 1)])\n",
      "/home/u1451673/.local/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:542: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint8 = np.dtype([(\"quint8\", np.uint8, 1)])\n",
      "/home/u1451673/.local/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:543: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint16 = np.dtype([(\"qint16\", np.int16, 1)])\n",
      "/home/u1451673/.local/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:544: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint16 = np.dtype([(\"quint16\", np.uint16, 1)])\n",
      "/home/u1451673/.local/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:545: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint32 = np.dtype([(\"qint32\", np.int32, 1)])\n",
      "/home/u1451673/.local/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:550: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  np_resource = np.dtype([(\"resource\", np.ubyte, 1)])\n"
     ]
    }
   ],
   "source": [
    "import torch, math\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils import data\n",
    "from torch import optim\n",
    "from torch.optim.lr_scheduler import LambdaLR\n",
    "\n",
    "from transformers import AdamW, get_linear_schedule_with_warmup\n",
    "from transformers import AutoTokenizer, AutoModel\n",
    "\n",
    "from gensim import corpora, similarities, models\n",
    "from gensim.matutils import corpus2csc\n",
    "\n",
    "import pandas as pd \n",
    "import numpy as np\n",
    "from collections import Counter\n",
    "import base64\n",
    "from tqdm import tqdm\n",
    "import pickle, random\n",
    "import matplotlib.pyplot as plt\n",
    "from multiprocessing import Pool\n",
    "import sys, csv, json, os, gc, time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### parameters\n",
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda:1\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7ffb157914d0>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "no = '1'\n",
    "device = torch.device('cuda:'+no) if torch.cuda.is_available() else torch.device('cpu')\n",
    "print(device)\n",
    "\n",
    "k = 10\n",
    "lr = 1e-5\n",
    "# true batch size = batch_size * grad_step\n",
    "batch_size = 64\n",
    "margin = 8\n",
    "grad_step = 1\n",
    "max_lang_len = 15\n",
    "max_img_len = 30\n",
    "epochs = 1\n",
    "MOD = 20000\n",
    "shuffle_fold = True\n",
    "seed = 9527\n",
    "random.seed(seed)\n",
    "torch.manual_seed(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class params:\n",
    "    IMG_FEAT_SIZE = 2048+6\n",
    "    WORD_EMBED_SIZE = 1024\n",
    "    LAYER = 6\n",
    "    HIDDEN_SIZE = 1024\n",
    "    MULTI_HEAD = 8\n",
    "    DROPOUT_R = 0.1\n",
    "    FLAT_MLP_SIZE = 512\n",
    "    FLAT_GLIMPSES = 1\n",
    "    FLAT_OUT_SIZE = 2048\n",
    "    FF_SIZE = HIDDEN_SIZE*4\n",
    "    HIDDEN_SIZE_HEAD = HIDDEN_SIZE // MULTI_HEAD\n",
    "    OPT_BETAS = (0.9, 0.98)\n",
    "    OPT_EPS = 1e-9\n",
    "    TRAIN_SIZE = 3000000\n",
    "\n",
    "__C = params()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### load data\n",
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "trash = {'!', '$', \"'ll\", \"'s\", ',', '&', ':', 'and', 'cut', 'is', 'are', 'was'}\n",
    "trash_replace = ['\"hey siri, play some', 'however, ', 'yin and yang, ',\n",
    "                 'shopping mall/']\n",
    "\n",
    "def process(x):\n",
    "    tmp = x.split()\n",
    "    if tmp[0] in trash: x = ' '.join(tmp[1:])\n",
    "    if tmp[0][0] == '-': x = x[1:]\n",
    "    for tr in trash_replace:\n",
    "        x = x.replace(tr, '')\n",
    "    return x\n",
    "\n",
    "def normalize(x):\n",
    "    ret = x['boxes'].copy()\n",
    "    ret[:,0] /= x['image_h']\n",
    "    ret[:,1] /= x['image_w']\n",
    "    ret[:,2] /= x['image_h']\n",
    "    ret[:,3] /= x['image_w']\n",
    "    wh = (ret[:,2]-ret[:,0]) * (ret[:,3]-ret[:,1])\n",
    "    wh2 = (ret[:,2]-ret[:,0]) / (ret[:,3]-ret[:,1]+1e-6)\n",
    "    ret = np.hstack((ret, wh.reshape(-1,1), wh2.reshape(-1,1)))\n",
    "    return ret\n",
    "\n",
    "def sort_by_area(x):\n",
    "    return np.array(sorted(x.tolist(), key=lambda x: x[-1], reverse=True))\n",
    "\n",
    "def load_data(file_name, reset=False, decode=True):\n",
    "    ret = pd.read_csv(file_name, sep='\\t')\n",
    "    if decode:\n",
    "        ret['boxes'] = ret['boxes'].apply(lambda x: np.frombuffer(base64.b64decode(x), dtype=np.float32).reshape(-1, 4))\n",
    "        ret['features'] = ret['features'].apply(lambda x: np.frombuffer(base64.b64decode(x), dtype=np.float32).reshape(-1, 2048))\n",
    "        ret['class_labels'] = ret['class_labels'].apply(lambda x: np.frombuffer(base64.b64decode(x), dtype=np.int64).reshape(-1, 1))\n",
    "        ret['boxes'] = ret.apply(lambda x: normalize(x), axis=1)\n",
    "        ret['features'] = ret.apply(lambda x: np.concatenate((x['features'], x['boxes']), axis=1)[:max_img_len], axis=1)\n",
    "    ret['query'] = ret['query'].apply(lambda x: process(x))\n",
    "    # reset query_id\n",
    "    if reset:\n",
    "        query2qid = {query: qid for qid, (query, _) in enumerate(tqdm(ret.groupby('query')))}\n",
    "        ret['query_id'] = ret.apply(lambda x: query2qid[x['query']], axis=1)\n",
    "    return ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = './'\n",
    "test = load_data(path+'valid.tsv')\n",
    "testA = load_data(path+'testB.tsv')\n",
    "answers = json.loads(open(path+'valid_answer.json', 'r').read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### preprocess\n",
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 496/496 [00:00<00:00, 2725.40it/s]\n",
      "100%|██████████| 994/994 [00:00<00:00, 2738.13it/s]\n"
     ]
    }
   ],
   "source": [
    "# load pre-trained model\n",
    "take = 'bert-large-uncased'\n",
    "emb_size = __C.WORD_EMBED_SIZE\n",
    "tokenizer = AutoTokenizer.from_pretrained(take)\n",
    "pretrained_emb = AutoModel.from_pretrained(take)\n",
    "pad_id = tokenizer.pad_token_id\n",
    "sep_id = tokenizer.sep_token_id\n",
    "\n",
    "qid2token = {qid: tokenizer.encode(group['query'].values[0]) for qid, group in tqdm(test.groupby('query_id'))}\n",
    "test['token'] = test['query_id'].apply(lambda x: qid2token[x])\n",
    "test['token'] = test['token'].apply(lambda x: x[:max_lang_len-1]+[sep_id] if len(x) > max_lang_len else x)\n",
    "qid2token = {qid: tokenizer.encode(group['query'].values[0]) for qid, group in tqdm(testA.groupby('query_id'))}\n",
    "testA['token'] = testA['query_id'].apply(lambda x: qid2token[x])\n",
    "testA['token'] = testA['token'].apply(lambda x: x[:max_lang_len-1]+[sep_id] if len(x) > max_lang_len else x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### model\n",
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FC(nn.Module):\n",
    "    def __init__(self, in_size, out_size, dropout_r=0., use_relu=True):\n",
    "        super(FC, self).__init__()\n",
    "        self.dropout_r = dropout_r\n",
    "        self.use_relu = use_relu\n",
    "\n",
    "        self.linear = nn.Linear(in_size, out_size)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.dropout = nn.Dropout(dropout_r)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.linear(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.dropout(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class MLP(nn.Module):\n",
    "    def __init__(self, in_size, mid_size, out_size, dropout_r=0., use_relu=True):\n",
    "        super(MLP, self).__init__()\n",
    "\n",
    "        self.fc = FC(in_size, mid_size, dropout_r=dropout_r, use_relu=use_relu)\n",
    "        self.linear = nn.Linear(mid_size, out_size)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.linear(self.fc(x))\n",
    "\n",
    "\n",
    "class VisualBERT(nn.Module):\n",
    "    def __init__(self, __C, bert):\n",
    "        super(VisualBERT, self).__init__()\n",
    "        \n",
    "        self.bert = bert\n",
    "        self.linear_img = nn.Linear(__C.IMG_FEAT_SIZE, __C.WORD_EMBED_SIZE)\n",
    "        self.out = MLP(__C.WORD_EMBED_SIZE, __C.WORD_EMBED_SIZE//2, 1)\n",
    "        \n",
    "    def forward(self, ques_ix, img_feats):\n",
    "        proj_feats = []\n",
    "        for img_feat in img_feats:\n",
    "            # Make mask & token type ids\n",
    "            mask = self.make_mask(ques_ix.unsqueeze(2), img_feat, pad_id)\n",
    "            token = self.get_token_type(ques_ix, img_feat)\n",
    "            # Preprocess features\n",
    "            lang_feat = self.bert.embeddings.word_embeddings(ques_ix)\n",
    "            img_feat = self.linear_img(img_feat)\n",
    "            combine_feat = torch.cat((lang_feat, img_feat), dim=1)\n",
    "            # Token embeddings & position embeddings\n",
    "            position_ids = torch.arange(token.size(1), dtype=torch.long, device=device)\n",
    "            position_ids = position_ids.unsqueeze(0).expand(token.size())\n",
    "            position_embeddings = self.bert.embeddings.position_embeddings(position_ids)\n",
    "            token_type_embeddings = self.bert.embeddings.token_type_embeddings(token)\n",
    "            # Add all\n",
    "            embeddings = combine_feat+position_embeddings+token_type_embeddings\n",
    "            embeddings = self.bert.embeddings.dropout(self.bert.embeddings.LayerNorm(embeddings))\n",
    "            # Go through the rest of BERT\n",
    "            head_mask = self.bert.get_head_mask(None, self.bert.config.num_hidden_layers)\n",
    "            extended_attention_mask = self.bert.get_extended_attention_mask(mask, mask.size(), device)\n",
    "            encoder_outputs = self.bert.encoder(embeddings,\n",
    "                                                attention_mask=extended_attention_mask,\n",
    "                                                head_mask=head_mask,\n",
    "                                                encoder_hidden_states=None,\n",
    "                                                encoder_attention_mask=None)\n",
    "            # CLS embedding & output value\n",
    "            outputs = encoder_outputs[0][:,0,:]\n",
    "            proj_feats.append(self.out(outputs))\n",
    "        return proj_feats\n",
    "            \n",
    "    # Masking\n",
    "    def make_mask(self, lang_feat, img_feat, target):\n",
    "        # 1 for NOT masked; 0 for masked\n",
    "        # [batch, len]\n",
    "        lang_mask = (torch.sum(torch.abs(lang_feat), dim=-1) != target).long()\n",
    "        img_mask = (torch.sum(torch.abs(img_feat), dim=-1) != 0).long()\n",
    "        return torch.cat((lang_mask, img_mask), dim=1)\n",
    "    \n",
    "    # Token type ids\n",
    "    def get_token_type(self, lang_feat, img_feat):\n",
    "        #    lang      img\n",
    "        # 0 0 0 0 0 0 1 1 1 1 1\n",
    "        lang_token = torch.zeros(lang_feat.size(0), lang_feat.size(1)).to(device)\n",
    "        img_token = torch.ones(img_feat.size(0), img_feat.size(1)).to(device)\n",
    "        return torch.cat((lang_token, img_token), dim=1).long()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### train\n",
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomDataset(data.Dataset):\n",
    "    def __init__(self, train_x):\n",
    "        self.train_x = train_x\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        tokens, pos_features, neg_features = self.train_x[index][0], self.train_x[index][1], self.train_x[index][2]\n",
    "        return [tokens, pos_features, neg_features]\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.train_x)\n",
    "    \n",
    "def collate_fn(batch):\n",
    "    tokens, pos_features, neg_features = zip(*batch)\n",
    "    max_len_t, max_len_pf, max_len_nf = len(max(tokens, key=lambda x: len(x))), len(max(pos_features, key=lambda x: len(x))), len(max(neg_features, key=lambda x: len(x)))\n",
    "    tokens, pos_features, neg_features = [token+[pad_id]*(max_len_t-len(token)) for token in tokens], [np.concatenate((feature, np.zeros((max_len_pf-feature.shape[0], feature.shape[1]))), axis=0) for feature in pos_features], [np.concatenate((feature, np.zeros((max_len_nf-feature.shape[0], feature.shape[1]))), axis=0) for feature in neg_features]\n",
    "    return torch.LongTensor(tokens), torch.FloatTensor(pos_features), torch.FloatTensor(neg_features)\n",
    "\n",
    "def custom_schedule(optimizer, num_warmup_steps, num_training_steps, num_cycles=0.5, amplitude=0.1, last_epoch=-1):\n",
    "    \n",
    "    def lr_lambda(current_step):\n",
    "        if current_step < num_warmup_steps:\n",
    "            return float(current_step) / float(max(1, num_warmup_steps))\n",
    "        progress = 2.0 * math.pi * float(num_cycles) * float(current_step-num_warmup_steps) / float(max(1, num_training_steps-num_warmup_steps))\n",
    "        linear = float(num_training_steps-current_step) / float(max(1, num_training_steps-num_warmup_steps))\n",
    "        return abs(linear + math.sin(progress)*linear*amplitude)\n",
    "\n",
    "    return LambdaLR(optimizer, lr_lambda, last_epoch)\n",
    "\n",
    "def shuffle(x):\n",
    "    idxs = [i for i in range(x.shape[0])]\n",
    "    random.shuffle(idxs)\n",
    "    return x[idxs]\n",
    "\n",
    "def nDCG_score(preds, answers):\n",
    "    iDCG = sum([sum([np.log(2)/np.log(i+2) for i in range(min(len(answer), 5))]) \\\n",
    "                for answer in list(answers.values())])\n",
    "    DCG = sum([sum([np.log(2)/np.log(i+2) if preds[qid][i] in answers[str(qid)] else 0 \\\n",
    "                    for i in range(len(preds[qid]))]) for qid in list(preds.keys())])\n",
    "    return DCG/iDCG\n",
    "\n",
    "class FocalLoss(nn.Module):\n",
    "    def __init__(self, alpha=1, gamma=2, logits=False, reduce=True):\n",
    "        super(FocalLoss, self).__init__()\n",
    "        self.alpha = alpha\n",
    "        self.gamma = gamma\n",
    "        self.logits = logits\n",
    "        self.reduce = reduce\n",
    "\n",
    "    def forward(self, inputs, targets):\n",
    "        if self.logits:\n",
    "            BCE_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduce=False)\n",
    "        else:\n",
    "            BCE_loss = F.binary_cross_entropy(inputs, targets, reduce=False)\n",
    "        pt = torch.exp(-BCE_loss)\n",
    "        F_loss = self.alpha * (1-pt)**self.gamma * BCE_loss\n",
    "\n",
    "        if self.reduce:\n",
    "            return torch.mean(F_loss)\n",
    "        else:\n",
    "            return F_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "initializing model...\n"
     ]
    }
   ],
   "source": [
    "print('initializing model...')\n",
    "nDCGs = []\n",
    "best_nDCG = 0.0\n",
    "model = VisualBERT(__C, pretrained_emb).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### prediction\n",
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(model, test, pad_len):\n",
    "    model = model.eval()\n",
    "    counts = Counter(test['product_id'].values.tolist())\n",
    "    preds = {}\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for qid, group in tqdm(test.groupby('query_id')):\n",
    "            # prepare batch\n",
    "            tokens, features = group['token'].values.tolist(), group['features'].values.tolist()\n",
    "            max_len_f = len(max(features, key=lambda x: len(x)))\n",
    "            features = [np.concatenate((feature, np.zeros((max_len_f-feature.shape[0], feature.shape[1]))), axis=0) for feature in features]\n",
    "            # # to tensor\n",
    "            tokens = torch.LongTensor(tokens).to(device)\n",
    "            features = torch.FloatTensor(features).to(device)\n",
    "            # predict\n",
    "            out = model(tokens, (features,))[0].view(-1)\n",
    "            pred = [(pid, val) for pid, val in zip(group['product_id'].values.tolist(), out.tolist())]\n",
    "            pred.sort(key=lambda x: x[1], reverse=True)\n",
    "            assert len(pred) <= pad_len\n",
    "            pid, score = [p for p, s in pred], [s for p, s in pred]\n",
    "            pid, score = pid+[np.nan]*(pad_len-len(pred)), score+[np.nan]*(pad_len-len(pred))\n",
    "            preds[qid] = pid+score\n",
    "            \n",
    "    model = model.train()\n",
    "    return preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "seed: 62; fold: 30\n",
      "seed: 62; fold: 31\n",
      "seed: 62; fold: 32\n",
      "seed: 62; fold: 33\n",
      "seed: 62; fold: 34\n",
      "seed: 62; fold: 35\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 994/994 [01:22<00:00, 12.02it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "seed: 62; fold: 36\n",
      "seed: 62; fold: 37\n",
      "seed: 62; fold: 38\n",
      "seed: 62; fold: 39\n"
     ]
    }
   ],
   "source": [
    "seeds = [62]\n",
    "folds = [i for i in range(30, 40)]\n",
    "pad_len = 30\n",
    "\n",
    "for seed in seeds:\n",
    "    for fold in folds:\n",
    "        print('seed: {}; fold: {}'.format(seed, fold))\n",
    "        # load model weights\n",
    "        try: model.load_state_dict(torch.load(path+'models/model_Visual-BERT_pair_box_tfidf-neg_focal_all_{}_{}'.format(seed, fold), map_location=device))\n",
    "        except: continue\n",
    "        # predict\n",
    "        preds = predict(model, testA, pad_len)\n",
    "        # write to file\n",
    "        header = ['qid'] + ['p'+str(i) for i in range(pad_len)] + ['s'+str(i) for i in range(pad_len)]\n",
    "        with open('predictions/prediction_all_{}_{}.csv'.format(seed, fold), 'w', newline='') as f:\n",
    "            w = csv.writer(f)\n",
    "            w.writerow(header)\n",
    "            for qid in sorted(list(preds.keys())):\n",
    "                w.writerow([qid]+preds[qid])"
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}