Visual-BERT_pair_box_tfidf-neg_focal_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:0\n"
     ]
    }
   ],
   "source": [
    "no = '0'\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 = 32\n",
    "margin = 8\n",
    "grad_step = 2\n",
    "max_lang_len = 15\n",
    "max_img_len = 30\n",
    "epochs = 1\n",
    "MOD = 20000\n",
    "seed = 7328\n",
    "random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "fold_idxs = [i for i in range(20)]\n",
    "random.shuffle(fold_idxs)"
   ]
  },
  {
   "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": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 496/496 [00:00<00:00, 2808.17it/s]\n",
      "100%|██████████| 994/994 [00:00<00:00, 2862.14it/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": [
    "### training data\n",
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def one(x):\n",
    "    return (x+1e-10)/(x+1e-10)\n",
    "\n",
    "def get_negs(train):\n",
    "    qid2idxs = {}\n",
    "    corpus = []\n",
    "    idx2qid = {}\n",
    "\n",
    "    for idx, (qid, group) in enumerate(train.groupby('query_id')):\n",
    "        qid2idxs[qid] = group.index\n",
    "        corpus.append(group['query'].values[0])\n",
    "        idx2qid[idx] = qid\n",
    "\n",
    "    topk = len(max(qid2idxs.values(), key=lambda x: len(x)))*3\n",
    "    corpus = [sent.split() for sent in corpus]\n",
    "    dictionary = corpora.Dictionary(corpus)\n",
    "    corpus = [dictionary.doc2bow(text) for text in corpus]\n",
    "    tfidf_model = models.TfidfModel(corpus, wlocal=one, dictionary=dictionary)\n",
    "    corpus_tfidf = corpus2csc(tfidf_model[corpus])\n",
    "    sm = corpus_tfidf.T.dot(corpus_tfidf)\n",
    "    qid2negs = {}\n",
    "    \n",
    "    for idx, (le, ri) in enumerate(tqdm(zip(sm.indptr[:-1], sm.indptr[1:]), total=sm.shape[0])):\n",
    "        n_row_pick = min(topk, ri-le)\n",
    "        top_n_idx = sm.indices[le+np.argpartition(sm.data[le:ri], -n_row_pick)[-n_row_pick:]].tolist()\n",
    "        if n_row_pick < topk: top_n_idx += random.sample(range(sm.shape[0]), topk-n_row_pick)\n",
    "        qid2negs[idx2qid[idx]] = [idx2qid[neg] for neg in top_n_idx if neg != idx]\n",
    "        \n",
    "    return qid2negs, qid2idxs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_data(train, qid2negs, qid2idxs):\n",
    "    train_x = [] # [tokens, feature1, feature2]\n",
    "    \n",
    "    for qid, group in tqdm(train.groupby('query_id')):\n",
    "        # positive\n",
    "        pos = group.index\n",
    "        token = group['token'].iloc[0]\n",
    "        # negative\n",
    "        for i in range(k):\n",
    "            neg = random.sample(qid2negs[qid], len(pos))\n",
    "            neg = [random.choice(qid2idxs[n]) for n in neg]\n",
    "            train_x += [[token, pos[i], neg[i]] for i in range(len(pos))]\n",
    "    \n",
    "    print('number of training data:', len(train_x))\n",
    "    return train_x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### model\n",
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(model):\n",
    "    model = model.eval()\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",
    "            preds[qid] = [pid for pid, _ in pred[:5]]\n",
    "            \n",
    "    model = model.train()\n",
    "    return preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "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": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_data_merge(idx):\n",
    "#     train = load_data(path+'train.sample.tsv', reset=True)\n",
    "    train = load_data(path+'data/train_{}.tsv'.format(idx), reset=True, decode=False)\n",
    "    qid2token = {qid: tokenizer.encode(group['query'].values[0]) for qid, group in train.groupby('query_id')}\n",
    "    train['token'] = train['query_id'].apply(lambda x: qid2token[x])\n",
    "    qid2negs, qid2idxs = get_negs(train)\n",
    "    pickle.dump(get_data(train, qid2negs, qid2idxs),\n",
    "                open(path+'data/train_x_{}_{}.pkl'.format(idx, seed), 'wb'))\n",
    "    return None\n",
    "\n",
    "def load_features(idx):\n",
    "#     train = load_data(path+'train.sample.tsv', reset=True)\n",
    "    train = load_data(path+'data/train_{}.tsv'.format(idx), reset=True)\n",
    "    train_x = pickle.load(open(path+'data/train_x_{}_{}.pkl'.format(idx, seed), 'rb'))\n",
    "    train_x = [[t, train['features'].iloc[p], train['features'].iloc[n]] \\\n",
    "               for t, p, n in tqdm(train_x)]\n",
    "    return train_x\n",
    "\n",
    "def get_data_chunk(start, end):\n",
    "    ret = []\n",
    "    for fold in range(start, end):\n",
    "        t0 = time.time()\n",
    "        print('reading fold:', fold)\n",
    "        ret += load_features(fold_idxs[fold])\n",
    "        t = round((time.time()-t0))\n",
    "        print('time consumed: {} min {} sec'.format(t//60, t%60))\n",
    "    return ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# t0 = time.time()\n",
    "# with Pool(5) as pool:\n",
    "#     pool.map(get_data_merge, [i for i in range(20)])\n",
    "# t = round((time.time()-t0)/60)\n",
    "# print('time consumed: {} hr {} min'.format(t//60, t%60))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "initializing model...\n",
      "reading fold: 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 114132/114132 [00:12<00:00, 8869.85it/s]\n",
      "100%|██████████| 1500000/1500000 [00:32<00:00, 46699.55it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time consumed: 2 min 46 sec\n",
      "reading fold: 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 114554/114554 [00:13<00:00, 8736.17it/s]\n",
      "100%|██████████| 1500000/1500000 [00:32<00:00, 46183.66it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time consumed: 2 min 44 sec\n",
      "reading fold: 2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 114314/114314 [00:13<00:00, 8775.52it/s]\n",
      "100%|██████████| 1500000/1500000 [00:31<00:00, 47437.66it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time consumed: 2 min 45 sec\n",
      "train all size: 750000\n",
      "start training!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/23438 [00:00<?, ?it/s]/home/u1451673/.local/lib/python3.6/site-packages/torch/nn/_reduction.py:43: UserWarning: size_average and reduce args will be deprecated, please use reduction='none' instead.\n",
      "  warnings.warn(warning.format(ret))\n",
      "  0%|          | 1/23438 [00:01<7:05:31,  1.09s/it, loss=0.347]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loss:0.3467542231082916\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/23438 [00:01<7:05:31,  1.09s/it, loss=0.344]/pytorch/torch/csrc/utils/python_arg_parser.cpp:756: UserWarning: This overload of add_ is deprecated:\n",
      "\tadd_(Number alpha, Tensor other)\n",
      "Consider using one of the following signatures instead:\n",
      "\tadd_(Tensor other, *, Number alpha)\n",
      " 12%|█▏        | 2824/23438 [14:59<2:03:43,  2.78it/s, loss=0.348]"
     ]
    }
   ],
   "source": [
    "print('initializing model...')\n",
    "nDCGs = []\n",
    "best_nDCG = 0.0\n",
    "model = VisualBERT(__C, pretrained_emb).to(device)\n",
    "num_training_steps = np.ceil(__C.TRAIN_SIZE*k / (batch_size*grad_step)) * epochs\n",
    "num_warmup_steps = int(num_training_steps*0.1)\n",
    "eval_steps = num_training_steps//40*grad_step + 100\n",
    "optimizer = AdamW(model.parameters(), lr=lr, correct_bias=False)\n",
    "scheduler = custom_schedule(optimizer,\n",
    "                            num_warmup_steps=num_warmup_steps,\n",
    "                            num_training_steps=num_training_steps,\n",
    "                            num_cycles=6,\n",
    "                            amplitude=0.3)\n",
    "\n",
    "for chunk in range(7):\n",
    "    train_x_all = get_data_chunk(chunk*3, min((chunk+1)*3, 20))\n",
    "    random.shuffle(train_x_all)\n",
    "    steps = (min((chunk+1)*3, 20) - chunk*3) * 2\n",
    "    step_size = len(train_x_all)//steps\n",
    "    \n",
    "    for step in range(steps):\n",
    "        train_loader = data.DataLoader(CustomDataset(train_x_all[step_size*step:step_size*(step+1)]),\n",
    "                                       batch_size=batch_size,\n",
    "                                       shuffle=True,\n",
    "                                       collate_fn=collate_fn,\n",
    "                                       num_workers=2)\n",
    "        print('train all size:', len(train_loader.dataset))\n",
    "\n",
    "        print('start training!')\n",
    "        model = model.train()\n",
    "        criterion = FocalLoss()\n",
    "        total_loss = 0.0\n",
    "        step = len(train_loader)\n",
    "        optimizer.zero_grad()\n",
    "        pbar = tqdm(enumerate(train_loader), total=step)\n",
    "\n",
    "        for i, batch in pbar:\n",
    "            # prepare batch\n",
    "            tokens, pos_features, neg_features = batch\n",
    "            # # to device\n",
    "            tokens = tokens.to(device)\n",
    "            pos_features = pos_features.to(device)\n",
    "            neg_features = neg_features.to(device)\n",
    "            # predict\n",
    "            pos, neg = model(tokens, (pos_features, neg_features))\n",
    "            pos = torch.sigmoid(pos).view(-1)\n",
    "            neg = torch.sigmoid(neg).view(-1)\n",
    "            l = criterion(pos, torch.ones(pos.size()).to(device))\n",
    "            l.backward()\n",
    "            total_loss += l.item()\n",
    "            l = criterion(neg, torch.zeros(neg.size()).to(device))\n",
    "            l.backward()\n",
    "            total_loss += l.item()\n",
    "            pbar.set_postfix({'loss': total_loss/(i+1)})\n",
    "            # optim step\n",
    "            if (i+1)%grad_step == 0:\n",
    "                optimizer.step()\n",
    "                scheduler.step()\n",
    "                optimizer.zero_grad()\n",
    "            # evaluate\n",
    "            if (i+1)%eval_steps == 0:\n",
    "                preds = predict(model)\n",
    "                nDCG = nDCG_score(preds, answers)\n",
    "                nDCGs.append(nDCG)\n",
    "                pickle.dump(nDCGs, open('log_nDCG_Visual-BERT_pair_box_tfidf-neg_focal_all', 'wb'))\n",
    "                print('nDCG@5:', nDCG)\n",
    "                # save models\n",
    "                if len(nDCGs) > 30:\n",
    "                    print('saving model...')\n",
    "                    torch.save(model.state_dict(), path+'models/model_Visual-BERT_pair_box_tfidf-neg_focal_all_{}_{}'.format(seed, len(nDCGs)-1))\n",
    "            # print loss\n",
    "            if i%MOD == 0: print('Loss:{}'.format(total_loss/(i+1)))\n",
    "\n",
    "        if step%grad_step:\n",
    "            optimizer.step()\n",
    "            scheduler.step()\n",
    "\n",
    "        # evaluation\n",
    "        preds = predict(model)\n",
    "        nDCG = nDCG_score(preds, answers)\n",
    "        nDCGs.append(nDCG)\n",
    "        pickle.dump(nDCGs, open('log_nDCG_Visual-BERT_pair_box_tfidf-neg_focal_all', 'wb'))\n",
    "        print('nDCG@5:', nDCG)\n",
    "        # delete garbage\n",
    "        del train_loader\n",
    "        gc.collect()\n",
    "        # save models\n",
    "        if len(nDCGs) > 30:\n",
    "            print('saving model...')\n",
    "            torch.save(model.state_dict(), path+'models/model_Visual-BERT_pair_box_tfidf-neg_focal_all_{}_{}'.format(seed, len(nDCGs)-1))\n",
    "            \n",
    "    # delete garbage\n",
    "    del train_x_all\n",
    "    gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'nDCGs' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-b0aaa7f58fc8>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mnDCGs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m: name 'nDCGs' is not defined"
     ]
    }
   ],
   "source": [
    "nDCGs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### prediction\n",
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(model, test, pad_len):\n",
    "    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",
    "    return preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "folds = [i for i in range(30, 40)]\n",
    "pad_len = 30\n",
    "\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
}