esim_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!shuf -n 80000000 /home/kesci/input/bytedance/train_final.csv > nn_train_8kw.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Parameters(object):\n",
    "    word2vec_path = \"nn_word2vec.model\"\n",
    "    fasttext_path = \"fasttext.model\"\n",
    "    \n",
    "    word_index_path = \"nn_words_index.txt\"\n",
    "    train_data_path = \"nn_train_8kw.csv\"\n",
    "    valid_data_path = \"nn_valid_10w.csv\"\n",
    "    test_data_path = \"/home/kesci/input/bytedance/bytedance_contest.final_2.csv\"\n",
    "    save_path = \"./model_file/esim_final_8_9_model.pkl\"\n",
    "    \n",
    "    BATCH_SIZE = 256\n",
    "    EMBEDDING_DIM = 300\n",
    "    MAX_QUERY_LENGTH = 10\n",
    "    MAX_TITLE_LENGTH = 20\n",
    "    HIDDEN_SIZE = 256\n",
    "    DENSE_SIZE = 256\n",
    "    DROPOUT_RATE = 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "from time import time\n",
    "\n",
    "from sklearn.metrics import roc_auc_score\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable\n",
    "\n",
    "mask_num=-2**32+1.0\n",
    "class ESIM(nn.Module):\n",
    "    \"\"\"\n",
    "    Implementation of the ESIM model\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, vocab_size, embedding_matrix):\n",
    "        super(ESIM, self).__init__()\n",
    "        self.verbose = True\n",
    "        self.use_cuda = True\n",
    "        self.vocab_size = vocab_size\n",
    "        self.embedding_dim = Parameters.EMBEDDING_DIM\n",
    "        self.hidden_size = Parameters.HIDDEN_SIZE\n",
    "        self.linear_size = Parameters.DENSE_SIZE\n",
    "        self.drop_out = Parameters.DROPOUT_RATE\n",
    "        self.learning_rate = 0.0001\n",
    "        self.optimizer = \"adam\"\n",
    "        self.eval_metric = roc_auc_score\n",
    "        self.n_epochs = 1\n",
    "        self.batch_size = Parameters.BATCH_SIZE\n",
    "\n",
    "        self.embed = nn.Embedding(self.vocab_size + 1, self.embedding_dim, padding_idx=0)\n",
    "        self.embed.weight.data.copy_(torch.from_numpy(embedding_matrix))\n",
    "        \n",
    "        self.bat_nor_embed = nn.BatchNorm1d(self.embedding_dim)\n",
    "        self.lstm1 = nn.LSTM(self.embedding_dim, self.hidden_size, batch_first=True, bidirectional=True)\n",
    "        self.lstm2 = nn.LSTM(self.embedding_dim*2*4, self.hidden_size, batch_first=True, bidirectional=True)\n",
    "        \n",
    "        self.fc = nn.Sequential(\n",
    "            nn.BatchNorm1d(self.hidden_size*4*2),\n",
    "            nn.Linear(self.hidden_size*2*4, self.linear_size),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm1d(self.linear_size),\n",
    "            nn.Dropout(self.drop_out),\n",
    "            nn.Linear(self.linear_size, 1)\n",
    "            )\n",
    "\n",
    "    def sub_mul_work(self, d1, d2):\n",
    "        mul = d1 * d2\n",
    "        sub = d1 - d2\n",
    "        \n",
    "        return torch.cat([sub, mul], dim=-1)\n",
    "    \n",
    "    def pool_work(self, d):\n",
    "        p1 = F.avg_pool1d(d.transpose(1, 2), d.size(1)).squeeze(-1)\n",
    "        p2 = F.max_pool1d(d.transpose(1, 2), d.size(1)).squeeze(-1)\n",
    "        \n",
    "        return torch.cat([p1, p2], 1)\n",
    "        \n",
    "    def soft_attention(self, d1, d2, mask1, mask2):  # mask, ignore padding data in the sequences during the computation of the attention\n",
    "        attention = torch.matmul(d1, d2.transpose(1,2))\n",
    "        mask1 = mask1.float().masked_fill(mask1, mask_num)\n",
    "        mask2 = mask2.float().masked_fill_(mask2, mask_num)\n",
    "        \n",
    "        atten_1 = F.softmax(attention + mask2.unsqueeze(1), dim=-1)\n",
    "        d1_align = torch.matmul(atten_1, d2)\n",
    "        atten_2 = F.softmax(attention.transpose(1, 2) + mask1.unsqueeze(1), dim=-1)\n",
    "        d2_align = torch.matmul(atten_2, d1)\n",
    "        \n",
    "        return d1_align, d2_align\n",
    "    \n",
    "    def forward(self, *input):\n",
    "        q, t = input[0], input[1]\n",
    "        mask1, mask2 = q.eq(0), t.eq(0)\n",
    "        \n",
    "        q = self.bat_nor_embed(self.embed(q).transpose(1, 2).contiguous()).transpose(1, 2)\n",
    "        t = self.bat_nor_embed(self.embed(t).transpose(1, 2).contiguous()).transpose(1, 2)\n",
    "        \n",
    "        q_lstm, _ = self.lstm1(q)\n",
    "        t_lstm, _ = self.lstm1(t)\n",
    "        \n",
    "        q_attn, t_attn = self.soft_attention(q_lstm, t_lstm, mask1, mask2)\n",
    "        \n",
    "        q_enhanced = torch.cat([q_lstm, q_attn, self.sub_mul_work(q_lstm, q_attn)], -1)\n",
    "        t_enhanced = torch.cat([t_lstm, t_attn, self.sub_mul_work(t_lstm, t_attn)], -1)\n",
    "        \n",
    "        v_q, _ = self.lstm2(q_enhanced)\n",
    "        v_t, _ = self.lstm2(t_enhanced)\n",
    "        \n",
    "        merged_q = self.pool_work(v_q)\n",
    "        merged_t = self.pool_work(v_t)\n",
    "        \n",
    "        merged = torch.cat([merged_q, merged_t], -1)\n",
    "        output = self.fc(merged)\n",
    "        \n",
    "        return output"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}