SE125Proj2_transformerEncoder.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A TransformerEncoder + RNNDecoder Implementation\n",
    "\n",
    "We provide a TransformerEncoder + RNNDecoder model for conversation modeling using deep learning.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Libraries\n",
    "In this section, we import third-party libraries to be used in this project.\n",
    "You may need to install them using `pip`:\n",
    "```\n",
    "    pip install tqdm\n",
    "    pip install cython\n",
    "    pip install tables\n",
    "    pip install tensorboardX\n",
    "    ...\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: http://mirrors.tencentyun.com/pypi/simple\n",
      "Requirement already satisfied: tqdm in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (4.46.1)\n",
      "Looking in indexes: http://mirrors.tencentyun.com/pypi/simple\n",
      "Requirement already satisfied: cython in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (0.29.21)\n",
      "Looking in indexes: http://mirrors.tencentyun.com/pypi/simple\n",
      "Requirement already satisfied: tables in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (3.6.1)\n",
      "Requirement already satisfied: numpy>=1.9.3 in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from tables) (1.18.5)\n",
      "Requirement already satisfied: numexpr>=2.6.2 in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from tables) (2.7.2)\n",
      "Looking in indexes: http://mirrors.tencentyun.com/pypi/simple\n",
      "Requirement already satisfied: tensorboardX in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (2.1)\n",
      "Requirement already satisfied: six in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from tensorboardX) (1.15.0)\n",
      "Requirement already satisfied: numpy in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from tensorboardX) (1.18.5)\n",
      "Requirement already satisfied: protobuf>=3.8.0 in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from tensorboardX) (3.12.2)\n",
      "Requirement already satisfied: setuptools in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from protobuf>=3.8.0->tensorboardX) (47.3.1.post20200616)\n",
      "Looking in indexes: http://mirrors.tencentyun.com/pypi/simple\n",
      "Collecting tensorboard\n",
      "  Downloading http://mirrors.tencentyun.com/pypi/packages/02/83/179c8f76e5716030cc3ee9433721161cfcc1d854e9ba20c9205180bb100a/tensorboard-2.4.0-py3-none-any.whl (10.6 MB)\n",
      "\u001b[K     |████████████████████████████████| 10.6 MB 11.3 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: numpy>=1.12.0 in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from tensorboard) (1.18.5)\n",
      "Collecting grpcio>=1.24.3\n",
      "  Downloading http://mirrors.tencentyun.com/pypi/packages/91/ab/8a7b37278fb59f3688af01cd069a5a4d2f3383eea2a2f78ddea4c7be047a/grpcio-1.34.0-cp36-cp36m-manylinux2014_x86_64.whl (4.0 MB)\n",
      "\u001b[K     |████████████████████████████████| 4.0 MB 12.5 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting tensorboard-plugin-wit>=1.6.0\n",
      "  Downloading http://mirrors.tencentyun.com/pypi/packages/b6/85/5c5ac0a8c5efdfab916e9c6bc18963f6a6996a8a1e19ec4ad8c9ac9c623c/tensorboard_plugin_wit-1.7.0-py3-none-any.whl (779 kB)\n",
      "\u001b[K     |████████████████████████████████| 779 kB 93.3 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: requests<3,>=2.21.0 in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from tensorboard) (2.24.0)\n",
      "Requirement already satisfied: wheel>=0.26; python_version >= \"3\" in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from tensorboard) (0.34.2)\n",
      "Collecting google-auth<2,>=1.6.3\n",
      "  Downloading http://mirrors.tencentyun.com/pypi/packages/3e/fd/df1b3a59aca9537a187413651ad63b290d165559150a868a298fa837fe7b/google_auth-1.24.0-py2.py3-none-any.whl (114 kB)\n",
      "\u001b[K     |████████████████████████████████| 114 kB 33.8 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: six>=1.10.0 in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from tensorboard) (1.15.0)\n",
      "Collecting absl-py>=0.4\n",
      "  Downloading http://mirrors.tencentyun.com/pypi/packages/bc/58/0aa6fb779dc69cfc811df3398fcbeaeefbf18561b6e36b185df0782781cc/absl_py-0.11.0-py3-none-any.whl (127 kB)\n",
      "\u001b[K     |████████████████████████████████| 127 kB 95.6 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting markdown>=2.6.8\n",
      "  Downloading http://mirrors.tencentyun.com/pypi/packages/ac/ef/24a91ca96efa0d7802dffb83ccc7a3c677027bea19ec3c9ee80be740408e/Markdown-3.3.3-py3-none-any.whl (96 kB)\n",
      "\u001b[K     |████████████████████████████████| 96 kB 81.4 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting google-auth-oauthlib<0.5,>=0.4.1\n",
      "  Downloading http://mirrors.tencentyun.com/pypi/packages/81/67/e2c34bb0628984c7ce71cce6ba6964cb29c418873847fc285f826e032e6e/google_auth_oauthlib-0.4.2-py2.py3-none-any.whl (18 kB)\n",
      "Collecting werkzeug>=0.11.15\n",
      "  Downloading http://mirrors.tencentyun.com/pypi/packages/cc/94/5f7079a0e00bd6863ef8f1da638721e9da21e5bacee597595b318f71d62e/Werkzeug-1.0.1-py2.py3-none-any.whl (298 kB)\n",
      "\u001b[K     |████████████████████████████████| 298 kB 20.6 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: setuptools>=41.0.0 in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from tensorboard) (47.3.1.post20200616)\n",
      "Requirement already satisfied: protobuf>=3.6.0 in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from tensorboard) (3.12.2)\n",
      "Requirement already satisfied: idna<3,>=2.5 in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from requests<3,>=2.21.0->tensorboard) (2.10)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from requests<3,>=2.21.0->tensorboard) (2020.6.20)\n",
      "Requirement already satisfied: chardet<4,>=3.0.2 in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from requests<3,>=2.21.0->tensorboard) (3.0.4)\n",
      "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from requests<3,>=2.21.0->tensorboard) (1.25.9)\n",
      "Collecting cachetools<5.0,>=2.0.0\n",
      "  Downloading http://mirrors.tencentyun.com/pypi/packages/92/da/d3c94fc7c72ad9298072681ec3e8cea86949acc5c4cce4290ba21f7050a8/cachetools-4.2.0-py3-none-any.whl (12 kB)\n",
      "Collecting rsa<5,>=3.1.4; python_version >= \"3.6\"\n",
      "  Downloading http://mirrors.tencentyun.com/pypi/packages/1c/df/c3587a667d6b308fadc90b99e8bc8774788d033efcc70f4ecaae7fad144b/rsa-4.6-py3-none-any.whl (47 kB)\n",
      "\u001b[K     |████████████████████████████████| 47 kB 66.5 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting pyasn1-modules>=0.2.1\n",
      "  Downloading http://mirrors.tencentyun.com/pypi/packages/95/de/214830a981892a3e286c3794f41ae67a4495df1108c3da8a9f62159b9a9d/pyasn1_modules-0.2.8-py2.py3-none-any.whl (155 kB)\n",
      "\u001b[K     |████████████████████████████████| 155 kB 82.4 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: importlib-metadata; python_version < \"3.8\" in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from markdown>=2.6.8->tensorboard) (1.7.0)\n",
      "Collecting requests-oauthlib>=0.7.0\n",
      "  Downloading http://mirrors.tencentyun.com/pypi/packages/a3/12/b92740d845ab62ea4edf04d2f4164d82532b5a0b03836d4d4e71c6f3d379/requests_oauthlib-1.3.0-py2.py3-none-any.whl (23 kB)\n",
      "Collecting pyasn1>=0.1.3\n",
      "  Downloading http://mirrors.tencentyun.com/pypi/packages/62/1e/a94a8d635fa3ce4cfc7f506003548d0a2447ae76fd5ca53932970fe3053f/pyasn1-0.4.8-py2.py3-none-any.whl (77 kB)\n",
      "\u001b[K     |████████████████████████████████| 77 kB 83.1 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: zipp>=0.5 in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from importlib-metadata; python_version < \"3.8\"->markdown>=2.6.8->tensorboard) (3.1.0)\n",
      "Collecting oauthlib>=3.0.0\n",
      "  Downloading http://mirrors.tencentyun.com/pypi/packages/05/57/ce2e7a8fa7c0afb54a0581b14a65b56e62b5759dbc98e80627142b8a3704/oauthlib-3.1.0-py2.py3-none-any.whl (147 kB)\n",
      "\u001b[K     |████████████████████████████████| 147 kB 30.6 MB/s eta 0:00:01\n",
      "\u001b[?25hInstalling collected packages: grpcio, tensorboard-plugin-wit, cachetools, pyasn1, rsa, pyasn1-modules, google-auth, absl-py, markdown, oauthlib, requests-oauthlib, google-auth-oauthlib, werkzeug, tensorboard\n",
      "Successfully installed absl-py-0.11.0 cachetools-4.2.0 google-auth-1.24.0 google-auth-oauthlib-0.4.2 grpcio-1.34.0 markdown-3.3.3 oauthlib-3.1.0 pyasn1-0.4.8 pyasn1-modules-0.2.8 requests-oauthlib-1.3.0 rsa-4.6 tensorboard-2.4.0 tensorboard-plugin-wit-1.7.0 werkzeug-1.0.1\n",
      "Looking in indexes: http://mirrors.tencentyun.com/pypi/simple\n",
      "Requirement already satisfied: nltk in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (3.5)\n",
      "Requirement already satisfied: click in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from nltk) (7.1.2)\n",
      "Requirement already satisfied: regex in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from nltk) (2020.11.13)\n",
      "Requirement already satisfied: tqdm in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from nltk) (4.46.1)\n",
      "Requirement already satisfied: joblib in /opt/conda/envs/pytorch_py3/lib/python3.6/site-packages (from nltk) (0.15.1)\n"
     ]
    }
   ],
   "source": [
    "from torch.nn.modules.transformer import _get_clones\n",
    "!pip install tqdm\n",
    "!pip install cython\n",
    "!pip install tables\n",
    "!pip install tensorboardX\n",
    "!pip install tensorboard\n",
    "!pip install nltk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import time\n",
    "import os\n",
    "import math\n",
    "import sys\n",
    "import tables\n",
    "import json\n",
    "import random\n",
    "from tqdm import tqdm\n",
    "\n",
    "import torch \n",
    "import torch.utils.data as data\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "import logging\n",
    "logger = logging.getLogger(__name__)\n",
    "logging.basicConfig(level=logging.DEBUG, format=\"%(message)s\")#,format=\"%(asctime)s: %(name)s: %(levelname)s: %(message)s\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Utilities\n",
    "\n",
    "In this section we maintain utilities for model construction and training. \n",
    "Please put your own utility modules/functions in this section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[nltk_data] Error loading punkt: <urlopen error [Errno 111] Connection\n",
      "[nltk_data]     refused>\n"
     ]
    }
   ],
   "source": [
    "PAD_ID, SOS_ID, EOS_ID, UNK_ID = [0, 1, 2, 3]\n",
    "\n",
    "def asHHMMSS(s):\n",
    "    m = math.floor(s / 60)\n",
    "    s -= m * 60\n",
    "    h = math.floor(m /60)\n",
    "    m -= h *60\n",
    "    return '%d:%d:%d'% (h, m, s)\n",
    "\n",
    "def timeSince(since, percent):\n",
    "    now = time.time()\n",
    "    s = now - since\n",
    "    es = s / (percent)\n",
    "    rs = es - s\n",
    "    return '%s<%s'%(asHHMMSS(s), asHHMMSS(rs))\n",
    "\n",
    "#######################################################################\n",
    "import nltk\n",
    "try: \n",
    "    nltk.word_tokenize(\"hello world\")\n",
    "except LookupError: \n",
    "    nltk.download('punkt')\n",
    "    \n",
    "def sent2indexes(sentence, vocab, maxlen):\n",
    "    '''sentence: a string or list of string\n",
    "       return: a numpy array of word indices\n",
    "    '''      \n",
    "    def convert_sent(sent, vocab, maxlen):\n",
    "        idxes = np.zeros(maxlen, dtype=np.int64)\n",
    "        idxes.fill(PAD_ID)\n",
    "        tokens = nltk.word_tokenize(sent.strip())\n",
    "        idx_len = min(len(tokens), maxlen)\n",
    "        for i in range(idx_len): idxes[i] = vocab.get(tokens[i], UNK_ID)\n",
    "        return idxes, idx_len\n",
    "    if type(sentence) is list:\n",
    "        inds, lens = [], []\n",
    "        for sent in sentence:\n",
    "            idxes, idx_len = convert_sent(sent, vocab, maxlen)\n",
    "            #idxes, idx_len = np.expand_dims(idxes, 0), np.array([idx_len])\n",
    "            inds.append(idxes)\n",
    "            lens.append(idx_len)\n",
    "        return np.vstack(inds), np.vstack(lens)\n",
    "    else:\n",
    "        inds, lens = sent2indexes([sentence], vocab, maxlen)\n",
    "        return inds[0], lens[0]\n",
    "\n",
    "def indexes2sent(indexes, vocab, ignore_tok=PAD_ID): \n",
    "    '''indexes: numpy array'''\n",
    "    def revert_sent(indexes, ivocab, ignore_tok=PAD_ID):\n",
    "        toks=[]\n",
    "        length=0\n",
    "        indexes=filter(lambda i: i!=ignore_tok, indexes)\n",
    "        for idx in indexes:\n",
    "            toks.append(ivocab[idx])\n",
    "            length+=1\n",
    "            if idx == EOS_ID:\n",
    "                break\n",
    "        return ' '.join(toks), length\n",
    "    \n",
    "    ivocab = {v: k for k, v in vocab.items()}\n",
    "    if indexes.ndim==1:# one sentence\n",
    "        return revert_sent(indexes, ivocab, ignore_tok)\n",
    "    else:# dim>1\n",
    "        sentences=[] # a batch of sentences\n",
    "        lens=[]\n",
    "        for inds in indexes:\n",
    "            sentence, length = revert_sent(inds, ivocab, ignore_tok)\n",
    "            sentences.append(sentence)\n",
    "            lens.append(length)\n",
    "        return sentences, lens\n",
    "    \n",
    "def save_model(model, epoch):\n",
    "    \"\"\"Save model parameters to checkpoint\"\"\"\n",
    "    ckpt_path=f'./output/checkpoint_iter{epoch}.pkl'\n",
    "    #print(f'Saving model parameters to {ckpt_path}')\n",
    "    torch.save(model.state_dict(), ckpt_path)\n",
    "        \n",
    "def load_model(model, epoch):\n",
    "    \"\"\"Load parameters from checkpoint\"\"\"\n",
    "    ckpt_path=f'./output/checkpoint_iter{epoch}.pkl'\n",
    "    #print(f'Loading model parameters from {ckpt_path}')\n",
    "    model.load_state_dict(torch.load(ckpt_path))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Configuration\n",
    "In this section, we configurate some hyperparameters for the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "def get_config():\n",
    "    conf = {\n",
    "    'maxlen':40, # maximum utterance length\n",
    "    'diaglen':10, # how many utterance kept in the context window\n",
    "\n",
    "    # Model Arguments\n",
    "    'emb_size':200, # size of word embeddings\n",
    "    'rnn_hid_utt':512, # number of rnn hidden units for utterance encoder\n",
    "    'rnn_hid_ctx':512, # number of rnn hidden units for context encoder\n",
    "    'rnn_hid_dec':512, # number of rnn hidden units for decoder\n",
    "    'n_layers':1, # number of layers\n",
    "    'dropout':0.5, # dropout applied to layers (0 = no dropout)\n",
    "    'teach_force': 0.8, # use teach force for decoder\n",
    "      \n",
    "    # Training Arguments\n",
    "    'batch_size':64,\n",
    "    'epochs':10, # maximum number of epochs\n",
    "    'lr':2e-4, # autoencoder learning rate\n",
    "    'beta1':0.9, # beta1 for adam\n",
    "    'init_w':0.05, # initial w\n",
    "    'clip':5.0,  # gradient clipping, max norm        \n",
    "    }\n",
    "    return conf "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Data Loader\n",
    "A tool to load batches from the binarized (.h5) dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "class DialogDataset(data.Dataset):\n",
    "    def __init__(self, filepath, max_ctx_len=7, max_utt_len=40):\n",
    "        # 1. Initialize file path or list of file names.\n",
    "        \"\"\"read training sentences(list of int array) from a hdf5 file\"\"\"\n",
    "        self.max_ctx_len=max_ctx_len\n",
    "        self.max_utt_len=max_utt_len\n",
    "        \n",
    "        print(\"loading data...\")\n",
    "        table = tables.open_file(filepath)\n",
    "        self.data = table.get_node('/sentences')[:].astype(np.long)\n",
    "        self.index = table.get_node('/indices')[:]\n",
    "        self.data_len = self.index.shape[0]\n",
    "        print(\"{} entries\".format(self.data_len))\n",
    "\n",
    "    def __getitem__(self, offset):\n",
    "        pos_utt, ctx_len, res_len = self.index[offset]['pos_utt'], self.index[offset]['ctx_len'], self.index[offset]['res_len']\n",
    "        ctx_arr=self.data[pos_utt-ctx_len:pos_utt]\n",
    "        res_arr=self.data[pos_utt:pos_utt+res_len]\n",
    "        ## split context array into utterances\n",
    "        context=[]\n",
    "        utt_lens=[]\n",
    "        utt=[]\n",
    "        for i, tok in enumerate(ctx_arr):\n",
    "            utt.append(ctx_arr[i])\n",
    "            if tok==EOS_ID:\n",
    "                if len(utt)<self.max_utt_len+1:\n",
    "                    utt_lens.append(len(utt)-1)# floor is not counted in the utt length\n",
    "                    utt.extend([PAD_ID]*(self.max_utt_len+1-len(utt)))  \n",
    "                else:\n",
    "                    utt=utt[:self.max_utt_len+1]\n",
    "                    utt[-1]=EOS_ID\n",
    "                    utt_lens.append(self.max_utt_len)\n",
    "                context.append(utt)                \n",
    "                utt=[]  \n",
    "        # utt=[]\n",
    "        # utt.copy(ctx_arr)\n",
    "        # utt.remove(SOS_ID)\n",
    "        # pad_size = self.max_ctx_len * self.max_utt_len - len(utt)\n",
    "        # for i in range(pad_size):\n",
    "        #     utt.append(PAD_ID)\n",
    "            \n",
    "        if len(context)>self.max_ctx_len: # trunk long context\n",
    "            context=context[-self.max_ctx_len:]\n",
    "            utt_lens=utt_lens[-self.max_ctx_len:]\n",
    "        context_len=len(context)\n",
    "        if len(context)<self.max_ctx_len: # pad short context\n",
    "            for i in range(len(context), self.max_ctx_len):\n",
    "                context.append([0, SOS_ID, EOS_ID]+[PAD_ID]*(self.max_utt_len-2)) # [floor, <sos>, <eos>, <pad>, <pad> ...]\n",
    "                utt_lens.append(2) # <s> and </s>\n",
    "        context = np.array(context)        \n",
    "        utt_lens=np.array(utt_lens)\n",
    "        floors=context[:,0]\n",
    "        context = context[:,1:]\n",
    "        \n",
    "        ## Padding ##    \n",
    "        response = res_arr[1:]\n",
    "        if len(response)<self.max_utt_len:\n",
    "            res_len=len(response)\n",
    "            response=np.append(response,[PAD_ID]*(self.max_utt_len-len(response)))\n",
    "        else:\n",
    "            response=response[:self.max_utt_len]\n",
    "            response[-1]=EOS_ID\n",
    "            res_len=self.max_utt_len\n",
    "\n",
    "        return context, context_len, utt_lens, floors, response, res_len\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.data_len\n",
    "    \n",
    "\n",
    "def load_dict(filename):\n",
    "    return json.loads(open(filename, \"r\").readline())\n",
    "\n",
    "def load_vecs(fin):         \n",
    "    \"\"\"read vectors (2D numpy array) from a hdf5 file\"\"\"\n",
    "    h5f = tables.open_file(fin)\n",
    "    h5vecs= h5f.root.vecs\n",
    "    \n",
    "    vecs=np.zeros(shape=h5vecs.shape,dtype=h5vecs.dtype)\n",
    "    vecs[:]=h5vecs[:]\n",
    "    h5f.close()\n",
    "    return vecs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Models\n",
    "Define your model(including its dependent sub-modules) here. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "import torch.nn.init as weight_init\n",
    "from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence\n",
    "from torch import Tensor\n",
    "from typing import Optional, Any\n",
    "\n",
    "class TransformerEncoder(nn.Module):\n",
    "    \"\"\"\n",
    "\n",
    "    Args:\n",
    "        encoder_layer: an instance of the TransformerEncoderLayer() class (required).\n",
    "        num_layers: the number of sub-encoder-layers in the encoder (required).\n",
    "        norm: the layer normalization component (optional).\n",
    "\n",
    "    \"\"\"\n",
    "    __constants__ = ['norm']\n",
    "\n",
    "    def __init__(self, num_layers, d_model,n_head,norm=None):\n",
    "        super(TransformerEncoder, self).__init__()\n",
    "        encoder_layer = nn.TransformerEncoderLayer(d_model, n_head)\n",
    "        self.layers = _get_clones(encoder_layer, num_layers)\n",
    "        self.num_layers = num_layers\n",
    "        self.norm = norm\n",
    "\n",
    "    def forward(self, src: Tensor, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:\n",
    "        r\"\"\"Pass the input through the encoder layers in turn.\n",
    "\n",
    "        Args:\n",
    "            src: the sequence to the encoder (required).\n",
    "            mask: the mask for the src sequence (optional).\n",
    "            src_key_padding_mask: the mask for the src keys per batch (optional).\n",
    "\n",
    "        Shape:\n",
    "            see the docs in Transformer class.\n",
    "        \"\"\"\n",
    "        output = src\n",
    "\n",
    "        for mod in self.layers:\n",
    "            output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)\n",
    "\n",
    "        if self.norm is not None:\n",
    "            output = self.norm(output)\n",
    "\n",
    "        return output\n",
    "\n",
    "\n",
    "class TransformerDecoder(nn.Module):\n",
    "    r\"\"\"TransformerDecoder is a stack of N decoder layers\n",
    "\n",
    "    Args:\n",
    "        decoder_layer: an instance of the TransformerDecoderLayer() class (required).\n",
    "        num_layers: the number of sub-decoder-layers in the decoder (required).\n",
    "        norm: the layer normalization component (optional).\n",
    "\n",
    "    \"\"\"\n",
    "    __constants__ = ['norm']\n",
    "\n",
    "    def __init__(self, decoder_layer, num_layers, norm=None):\n",
    "        super(TransformerDecoder, self).__init__()\n",
    "        self.layers = _get_clones(decoder_layer, num_layers)\n",
    "        self.num_layers = num_layers\n",
    "        self.norm = norm\n",
    "\n",
    "    def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None,\n",
    "                memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None,\n",
    "                memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:\n",
    "        r\"\"\"Pass the inputs (and mask) through the decoder layer in turn.\n",
    "\n",
    "        Args:\n",
    "            tgt: the sequence to the decoder (required).\n",
    "            memory: the sequence from the last layer of the encoder (required).\n",
    "            tgt_mask: the mask for the tgt sequence (optional).\n",
    "            memory_mask: the mask for the memory sequence (optional).\n",
    "            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).\n",
    "            memory_key_padding_mask: the mask for the memory keys per batch (optional).\n",
    "        \"\"\"\n",
    "        output = tgt\n",
    "\n",
    "        for mod in self.layers:\n",
    "            output = mod(output, memory, tgt_mask=tgt_mask,\n",
    "                         memory_mask=memory_mask,\n",
    "                         tgt_key_padding_mask=tgt_key_padding_mask,\n",
    "                         memory_key_padding_mask=memory_key_padding_mask)\n",
    "\n",
    "        if self.norm is not None:\n",
    "            output = self.norm(output)\n",
    "\n",
    "        return output\n",
    "    \n",
    "class RNNEncoder(nn.Module):\n",
    "    def __init__(self, embedder, input_size, hidden_size, bidir, n_layers, dropout=0.5):\n",
    "        super(RNNEncoder, self).__init__()\n",
    "        \n",
    "        self.hidden_size = hidden_size\n",
    "        self.n_layers = n_layers\n",
    "        self.bidir = bidir\n",
    "        assert type(self.bidir)==bool\n",
    "        self.dropout=dropout\n",
    "        \n",
    "        self.embedding = embedder # nn.Embedding(vocab_size, emb_size)\n",
    "        self.rnn = nn.GRU(input_size, hidden_size, n_layers, batch_first=True, bidirectional=bidir)\n",
    "        self.init_h = nn.Parameter(torch.randn(self.n_layers*(1+self.bidir), 1, self.hidden_size), requires_grad=True)#learnable h0\n",
    "        self.init_weights()\n",
    "        \n",
    "    def init_weights(self):\n",
    "        \"\"\"adopted from https://gist.github.com/jeasinema/ed9236ce743c8efaf30fa2ff732749f5\"\"\"\n",
    "        for w in self.rnn.parameters(): # initialize the gate weights with orthogonal\n",
    "            if len(w.shape)>1: \n",
    "                weight_init.orthogonal_(w.data)\n",
    "            else:\n",
    "                weight_init.normal_(w.data)\n",
    "                \n",
    "    \n",
    "    def forward(self, inputs, input_lens=None, init_h=None): \n",
    "        # init_h: [n_layers*n_dir x batch_size x hid_size]\n",
    "        if self.embedding is not None:\n",
    "            inputs=self.embedding(inputs)  # input: [batch_sz x seq_len] -> [batch_sz x seq_len x emb_sz]\n",
    "        \n",
    "        batch_size, seq_len, emb_size=inputs.size()\n",
    "        inputs=F.dropout(inputs, self.dropout, self.training)# dropout\n",
    "        \n",
    "        if input_lens is not None:# sort and pack sequence \n",
    "            input_lens_sorted, indices = input_lens.sort(descending=True)\n",
    "            inputs_sorted = inputs.index_select(0, indices)        \n",
    "            inputs = pack_padded_sequence(inputs_sorted, input_lens_sorted.data.tolist(), batch_first=True)\n",
    "        \n",
    "        if init_h is None:\n",
    "            init_h = self.init_h.expand(-1,batch_size,-1).contiguous()# use learnable initial states, expanding along batches\n",
    "        #self.rnn.flatten_parameters() # time consuming!!\n",
    "        hids, h_n = self.rnn(inputs, init_h) # hids: [b x seq x (n_dir*hid_sz)]  \n",
    "                                                  # h_n: [(n_layers*n_dir) x batch_sz x hid_sz] (2=fw&bw)\n",
    "        if input_lens is not None: # reorder and pad\n",
    "            _, inv_indices = indices.sort()\n",
    "            hids, lens = pad_packed_sequence(hids, batch_first=True)     \n",
    "            hids = hids.index_select(0, inv_indices)\n",
    "            h_n = h_n.index_select(1, inv_indices)\n",
    "        \n",
    "        h_n = h_n.view(self.n_layers, (1+self.bidir), batch_size, self.hidden_size) #[n_layers x n_dirs x batch_sz x hid_sz]\n",
    "        h_n = h_n[-1] # get the last layer [n_dirs x batch_sz x hid_sz]\n",
    "#         print(h_n.transpose(0,1).view(batch_size,-1))\n",
    "        enc = h_n.view(batch_size,-1) #[batch_sz x (n_dirs*hid_sz)]\n",
    "#         print(enc)\n",
    "            \n",
    "        return enc, hids\n",
    "    \n",
    "class ContextEncoder(nn.Module):\n",
    "    def __init__(self, utt_encoder, input_size, hidden_size, n_layers=1, dropout=0.5):\n",
    "        super(ContextEncoder, self).__init__()     \n",
    "        self.utt_encoder=utt_encoder\n",
    "        # self.ctx_encoder= RNNEncoder(None, input_size, hidden_size, False, n_layers, dropout)\n",
    "        self.ctx_encoder=TransformerEncoder(n_layers,hidden_size,8)\n",
    "    def forward(self, context, context_lens, utt_lens): # context: [batch_sz x diag_len x max_utt_len] \n",
    "                                                      # context_lens: [batch_sz x dia_len]\n",
    "        batch_size, max_context_len, max_utt_len = context.size()\n",
    "        utts = context.view(-1, max_utt_len) # [(batch_size*diag_len) x max_utt_len]\n",
    "        utt_lens = utt_lens.view(-1)\n",
    "        utt_encs, _ = self.utt_encoder(utts, utt_lens) # [(batch_size*diag_len) x 2hid_size]\n",
    "        \n",
    "        utt_encs = utt_encs.view(batch_size, max_context_len, -1)\n",
    "        # enc, hids = self.ctx_encoder(utt_encs, context_lens)\n",
    "        utt_encs = utt_encs.transpose(0,1)\n",
    "        enc = self.ctx_encoder(utt_encs)\n",
    "        enc = enc.transpose(0,1)        \n",
    "        enc = enc.sum(dim=1)\n",
    "        \n",
    "        ###warning!\n",
    "#         enc = enc.tanh()\n",
    "        return enc\n",
    "  \n",
    "\n",
    "class RNNDecoder(nn.Module):\n",
    "    def __init__(self, embedder, input_size, hidden_size, vocab_size, n_layers=1, dropout=0.5):\n",
    "        super(RNNDecoder, self).__init__()\n",
    "        self.n_layers = n_layers\n",
    "        self.input_size= input_size # size of the input to the RNN (e.g., embedding dim)\n",
    "        self.hidden_size = hidden_size # RNN hidden size\n",
    "        self.vocab_size = vocab_size # RNN output size (vocab size)\n",
    "        self.dropout= dropout\n",
    "        \n",
    "        self.embedding = embedder\n",
    "        self.rnn = nn.GRU(input_size, hidden_size, batch_first=True)\n",
    "        self.project = nn.Linear(hidden_size, vocab_size)\n",
    "        \n",
    "        self.init_weights()\n",
    "        \n",
    "    def init_weights(self):\n",
    "        for w in self.rnn.parameters(): # initialize the gate weights with orthogonal\n",
    "            if w.dim()>1:\n",
    "                weight_init.orthogonal_(w)\n",
    "        self.project.weight.data.uniform_(-0.1, 0.1)#nn.init.xavier_normal_(self.out.weight)        \n",
    "        nn.init.constant_(self.project.bias, 0.)\n",
    "\n",
    "    def forward(self, init_h, inputs=None, lens=None, enc_hids=None, src_pad_mask=None, context=None):\n",
    "        '''\n",
    "        init_h: initial hidden state for decoder\n",
    "        enc_hids: enc_hids for attention use\n",
    "        context: context information to be paired with input\n",
    "        inputs: inputs to the decoder\n",
    "        lens: input lengths\n",
    "        '''\n",
    "        if self.embedding is not None:\n",
    "            inputs = self.embedding(inputs) # input: [batch_sz x seqlen x emb_sz]\n",
    "        batch_size, maxlen, _ = inputs.size()\n",
    "        inputs = F.dropout(inputs, self.dropout, self.training)  \n",
    "        h = init_h.unsqueeze(0) # last_hidden of decoder [n_dir x batch_sz x hid_sz]        \n",
    "\n",
    "        if context is not None:            \n",
    "            repeated_context = context.unsqueeze(1).repeat(1, maxlen, 1) # [batch_sz x max_len x hid_sz]\n",
    "            inputs = torch.cat([inputs, repeated_context], 2)\n",
    "                \n",
    "            #self.rnn.flatten_parameters()\n",
    "        hids, h = self.rnn(inputs, h)         \n",
    "        decoded = self.project(hids.contiguous().view(-1, self.hidden_size))# reshape before linear over vocab\n",
    "        decoded = decoded.view(batch_size, maxlen, self.vocab_size)\n",
    "        return decoded, h\n",
    "    \n",
    "    def sampling(self, init_h, enc_hids, src_pad_mask, context, maxlen, to_numpy=True):\n",
    "        \"\"\"\n",
    "        A simple greedy sampling\n",
    "        :param init_h: [batch_sz x hid_sz]\n",
    "        :param enc_hids: a tuple of (enc_hids, mask) for attention use. [batch_sz x seq_len x hid_sz]\n",
    "        \"\"\"\n",
    "        device = init_h.device\n",
    "        batch_size = init_h.size(0)\n",
    "        decoded_words = torch.zeros((batch_size, maxlen), dtype=torch.long, device=device)  \n",
    "        sample_lens = torch.zeros((batch_size), dtype=torch.long, device=device)\n",
    "        len_inc = torch.ones((batch_size), dtype=torch.long, device=device)\n",
    "               \n",
    "        x = torch.zeros((batch_size, 1), dtype=torch.long, device=device).fill_(SOS_ID)# [batch_sz x 1] (1=seq_len)\n",
    "        h = init_h.unsqueeze(0) # [1 x batch_sz x hid_sz]  \n",
    "        for di in range(maxlen):  \n",
    "            if self.embedding is not None:\n",
    "                x = self.embedding(x) # x: [batch_sz x 1 x emb_sz]\n",
    "            h_n, h = self.rnn(x, h) # h_n: [batch_sz x 1 x hid_sz] h: [1 x batch_sz x hid_sz]\n",
    "\n",
    "            logits = self.project(h_n) # out: [batch_sz x 1 x vocab_sz]  \n",
    "            logits = logits.squeeze(1) # [batch_size x vocab_size]                  \n",
    "            x = torch.multinomial(F.softmax(logits, dim=1), 1)  # [batch_size x 1 x 1]?\n",
    "            decoded_words[:,di] = x.squeeze()\n",
    "            len_inc=len_inc*(x.squeeze()!=EOS_ID).long() # stop increse length (set 0 bit) when EOS is met\n",
    "            sample_lens=sample_lens+len_inc            \n",
    "        \n",
    "        if to_numpy:\n",
    "            decoded_words = decoded_words.data.cpu().numpy()\n",
    "            sample_lens = sample_lens.data.cpu().numpy()\n",
    "        return decoded_words, sample_lens\n",
    "\n",
    "class MyModel(nn.Module):\n",
    "    '''The basic Hierarchical Recurrent Encoder-Decoder model. '''\n",
    "    def __init__(self, config, vocab_size):\n",
    "        super(MyModel, self).__init__()\n",
    "        self.vocab_size = vocab_size \n",
    "        self.maxlen=config['maxlen']\n",
    "        self.clip = config['clip']\n",
    "        self.init_w = config['init_w']\n",
    "        \n",
    "        self.embedder= nn.Embedding(vocab_size, config['emb_size'], padding_idx=PAD_ID)\n",
    "        self.utt_encoder = RNNEncoder(self.embedder, config['emb_size'], config['rnn_hid_utt'], False, \n",
    "                                   config['n_layers'], config['dropout']) \n",
    "                                                        # utter encoder: encode response to vector\n",
    "        self.context_encoder = ContextEncoder(self.utt_encoder, config['rnn_hid_utt']*2,\n",
    "                                              config['rnn_hid_ctx'], 1, config['dropout']) \n",
    "                                              # context encoder: encode context to vector    \n",
    "        self.decoder = RNNDecoder(self.embedder, config['emb_size'], config['rnn_hid_ctx'], vocab_size, 1, config['dropout']) # utter decoder: P(x|c,z)\n",
    "        self.optimizer = optim.Adam(list(self.context_encoder.parameters())\n",
    "                                      +list(self.decoder.parameters()),lr=config['lr'])\n",
    "\n",
    "    def forward(self, context, context_lens, utt_lens, response, res_lens):\n",
    "        c = self.context_encoder(context, context_lens, utt_lens)\n",
    "        output,_ = self.decoder(c, response[:,:-1], res_lens-1) # decode from z, c  # output: [batch x seq_len x n_tokens]   \n",
    "        dec_target = response[:,1:].clone()\n",
    "        dec_target[response[:,1:]==PAD_ID] = -100\n",
    "        loss = nn.CrossEntropyLoss()(output.view(-1, self.vocab_size), dec_target.view(-1))\n",
    "        return loss\n",
    "    \n",
    "    def train_batch(self, context, context_lens, utt_lens, response, res_lens):\n",
    "        self.context_encoder.train()\n",
    "        self.decoder.train()\n",
    "        \n",
    "        loss = self.forward(context, context_lens, utt_lens, response, res_lens)\n",
    "        \n",
    "        self.optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        # `clip_grad_norm` to prevent exploding gradient in RNNs\n",
    "        nn.utils.clip_grad_norm_(list(self.context_encoder.parameters())+list(self.decoder.parameters()), self.clip)\n",
    "        self.optimizer.step()\n",
    "        \n",
    "        return {'train_loss': loss.item()}      \n",
    "    \n",
    "    def valid(self, context, context_lens, utt_lens, response, res_lens):\n",
    "        self.context_encoder.eval()  \n",
    "        self.decoder.eval()        \n",
    "        loss = self.forward(context, context_lens, utt_lens, response, res_lens)\n",
    "        return {'valid_loss': loss.item()}\n",
    "    \n",
    "    def sample(self, context, context_lens, utt_lens, n_samples):    \n",
    "        self.context_encoder.eval()\n",
    "        self.decoder.eval()\n",
    "        with torch.no_grad():\n",
    "            c = self.context_encoder(context, context_lens, utt_lens)\n",
    "        sample_words, sample_lens = self.decoder.sampling(c, None, None, None, n_samples, self.maxlen)  \n",
    "        return sample_words, sample_lens  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Evaluation\n",
    "We provide the evaluation script as well as the BLEU score metric. \n",
    "\n",
    "**Do not change code in this block**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "from nltk.translate.bleu_score import sentence_bleu\n",
    "from nltk.translate.bleu_score import SmoothingFunction\n",
    "from collections import Counter\n",
    "\n",
    "class Metrics:\n",
    "    \"\"\"\n",
    "    \"\"\"\n",
    "    def __init__(self):\n",
    "        super(Metrics, self).__init__()\n",
    "\n",
    "    def sim_bleu(self, hyps, ref):\n",
    "        \"\"\"\n",
    "        :param ref - a list of tokens of the reference\n",
    "        :param hyps - a list of tokens of the hypothesis\n",
    "    \n",
    "        :return maxbleu - recall bleu\n",
    "        :return avgbleu - precision bleu\n",
    "        \"\"\"\n",
    "        scores = []\n",
    "        for hyp in hyps:\n",
    "            try:\n",
    "                scores.append(sentence_bleu([ref], hyp, smoothing_function=SmoothingFunction().method7,\n",
    "                                        weights=[1./4, 1./4, 1./4, 1./4]))\n",
    "            except:\n",
    "                scores.append(0.0)\n",
    "        return np.max(scores), np.mean(scores)\n",
    "    \n",
    "def evaluate(model, metrics, test_loader, vocab, repeat, f_eval):\n",
    "    ivocab = {v: k for k, v in vocab.items()}\n",
    "    device = next(model.parameters()).device\n",
    "    \n",
    "    recall_bleus, prec_bleus, avg_lens  = [], [], []\n",
    "        \n",
    "    dlg_id = 0\n",
    "    for context, context_lens, utt_lens, floors, response, res_lens in tqdm(test_loader): \n",
    "        \n",
    "        if dlg_id > 5000: break\n",
    "        \n",
    "#        max_ctx_len = max(context_lens)\n",
    "        max_ctx_len = context.size(1)\n",
    "        context, utt_lens, floors = context[:,:max_ctx_len,1:], utt_lens[:,:max_ctx_len]-1, floors[:,:max_ctx_len] \n",
    "                         # remove empty utts and the sos token in the context and reduce the context length\n",
    "        ctx, ctx_lens = context, context_lens\n",
    "        context, context_lens, utt_lens \\\n",
    "            = [tensor.to(device) for tensor in [context, context_lens, utt_lens]]\n",
    "\n",
    "#################################################\n",
    "        utt_lens[utt_lens<=0]=1\n",
    "#################################################\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            sample_words, sample_lens = model.sample(context, context_lens, utt_lens, repeat)\n",
    "        # nparray: [repeat x seq_len]       \n",
    "        \n",
    "        pred_sents, _ = indexes2sent(sample_words, vocab)\n",
    "        pred_tokens = [sent.split(' ') for sent in pred_sents]   \n",
    "        ref_str, _ =indexes2sent(response[0].numpy(), vocab, SOS_ID)\n",
    "        #ref_str = ref_str.encode('utf-8')\n",
    "        ref_tokens = ref_str.split(' ')\n",
    "        \n",
    "        max_bleu, avg_bleu = metrics.sim_bleu(pred_tokens, ref_tokens)\n",
    "        recall_bleus.append(max_bleu)\n",
    "        prec_bleus.append(avg_bleu)\n",
    "        \n",
    "        avg_lens.append(np.mean(sample_lens))\n",
    "\n",
    "        response, res_lens = [tensor.to(device) for tensor in [response, res_lens]]\n",
    "        \n",
    "        ## Write concrete results to a text file\n",
    "        dlg_id += 1 \n",
    "        if f_eval is not None:\n",
    "            f_eval.write(\"Batch {:d} \\n\".format(dlg_id))\n",
    "            # print the context\n",
    "            start = np.maximum(0, ctx_lens[0]-5)\n",
    "            for t_id in range(start, ctx_lens[0], 1):\n",
    "                context_str = indexes2sent(ctx[0, t_id].numpy(), vocab)\n",
    "                f_eval.write(\"Context {:d}-{:d}: {}\\n\".format(t_id, floors[0, t_id], context_str))\n",
    "            #print the ground truth response    \n",
    "            f_eval.write(\"Target >> {}\\n\".format(ref_str.replace(\" ' \", \"'\")))\n",
    "            for res_id, pred_sent in enumerate(pred_sents):\n",
    "                f_eval.write(\"Sample {:d} >> {}\\n\".format(res_id, pred_sent.replace(\" ' \", \"'\")))\n",
    "            f_eval.write(\"\\n\")\n",
    "    prec_bleu= float(np.mean(prec_bleus))\n",
    "    recall_bleu = float(np.mean(recall_bleus))\n",
    "    result = {'avg_len':float(np.mean(avg_lens)),\n",
    "              'recall_bleu': recall_bleu, 'prec_bleu': prec_bleu, \n",
    "              'f1_bleu': 2*(prec_bleu*recall_bleu) / (prec_bleu+recall_bleu+10e-12),\n",
    "             }\n",
    "    \n",
    "    if f_eval is not None:\n",
    "        for k, v in result.items():\n",
    "            f_eval.write(str(k) + ':'+ str(v)+' ')\n",
    "        f_eval.write('\\n')\n",
    "    print(\"Done testing\")\n",
    "    print(result)\n",
    "    \n",
    "    return result\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Training\n",
    "The training script here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "import argparse\n",
    "from datetime import datetime\n",
    "from tensorboardX import SummaryWriter # install tensorboardX (pip install tensorboardX) before importing this package\n",
    "\n",
    "def train(args, model=None, pad = 0):\n",
    "    # LOG #\n",
    "    fh = logging.FileHandler(f\"./output/logs.txt\")\n",
    "                                      # create file handler which logs even debug messages\n",
    "    logger.addHandler(fh)# add the handlers to the logger\n",
    "    \n",
    "    timestamp = datetime.now().strftime('%Y%m%d%H%M')\n",
    "    tb_writer = SummaryWriter(f\"./output/logs/{timestamp}-update-transformer\") if args.visual else None\n",
    "\n",
    "    # Set the random seed manually for reproducibility.\n",
    "    random.seed(args.seed)\n",
    "    np.random.seed(args.seed)\n",
    "    torch.manual_seed(args.seed)\n",
    "    torch.cuda.manual_seed(args.seed)\n",
    "    device = torch.device(f\"cuda:{args.gpu_id}\" if torch.cuda.is_available() else \"cpu\")\n",
    "    print(device)\n",
    "\n",
    "\n",
    "    config=get_config()\n",
    "\n",
    "    if args.visual:\n",
    "        json.dump(config, open(f'./output/config_{timestamp}.json', 'w'))# save configs\n",
    "\n",
    "    ###############################################################################\n",
    "    # Load data\n",
    "    ###############################################################################\n",
    "    data_path = args.data_path+'/'\n",
    "    train_set = DialogDataset(os.path.join(data_path, 'train.h5'), config['diaglen'], config['maxlen'])\n",
    "    valid_set = DialogDataset(os.path.join(data_path, 'valid.h5'), config['diaglen'], config['maxlen'])\n",
    "    test_set = DialogDataset(os.path.join(data_path, 'test.h5'), config['diaglen'], config['maxlen'])\n",
    "    vocab = load_dict(os.path.join(data_path, 'vocab.json'))\n",
    "    ivocab = {v: k for k, v in vocab.items()}\n",
    "    n_tokens = len(ivocab)\n",
    "    metrics=Metrics()    \n",
    "    print(\"Loaded data!\")\n",
    "\n",
    "    ###############################################################################\n",
    "    # Define the models\n",
    "    ###############################################################################\n",
    "    if model is None:\n",
    "        model = MyModel(config, n_tokens)\n",
    "\n",
    "    if args.reload_from>=0:\n",
    "        load_model(model, args.reload_from)\n",
    "        \n",
    "    model=model.to(device)\n",
    "\n",
    "    logger.info(\"Training...\")\n",
    "    flag = 0\n",
    "    best_perf = -1\n",
    "    itr_global=1\n",
    "    start_epoch=1 if args.reload_from==-1 else args.reload_from+1\n",
    "    for epoch in range(start_epoch, config['epochs']+1):\n",
    "        epoch_start_time = time.time()\n",
    "        itr_start_time = time.time()\n",
    "        \n",
    "        # shuffle (re-define) data between epochs   \n",
    "        train_loader=torch.utils.data.DataLoader(dataset=train_set, batch_size=config['batch_size'],\n",
    "                                                 shuffle=True, num_workers=1, drop_last=True)\n",
    "        n_iters=train_loader.__len__()\n",
    "        itr = 1\n",
    "        for batch in train_loader:# loop through all batches in training data\n",
    "            model.train()\n",
    "            context, context_lens, utt_lens, floors, response, res_lens = batch\n",
    "\n",
    " #           max_ctx_len = max(context_lens)\n",
    "            max_ctx_len = context.size(1)\n",
    "            context, utt_lens = context[:,:max_ctx_len,1:], utt_lens[:,:max_ctx_len]-1\n",
    "                                    # remove empty utterances in context\n",
    "                                    # remove the sos token in the context and reduce the context length\n",
    "#################################################\n",
    "            utt_lens[utt_lens<=0]=1\n",
    "#################################################\n",
    "            batch_gpu = [tensor.to(device) for tensor in [context, context_lens, utt_lens, response, res_lens]] \n",
    "            train_results = model.train_batch(*batch_gpu)\n",
    "        \n",
    "            #################################################\n",
    "            #    --------------draw graph----------\n",
    "            if flag == 0:\n",
    "                flag = 1\n",
    "                tb_writer.add_graph(model,batch_gpu)\n",
    "            #################################################\n",
    "                     \n",
    "            if itr % args.log_every == 0:\n",
    "                elapsed = time.time() - itr_start_time\n",
    "                log = '%s|%s@gpu%d epo:[%d/%d] iter:[%d/%d] step_time:%ds elapsed:%s'\\\n",
    "                %(args.model, args.dataset, args.gpu_id, epoch, config['epochs'],\n",
    "                         itr, n_iters, elapsed, timeSince(epoch_start_time,itr/n_iters))\n",
    "                logger.info(log)\n",
    "                logger.info(train_results)\n",
    "                if args.visual:\n",
    "                    tb_writer.add_scalar('train_loss', train_results['train_loss'], itr_global)\n",
    "\n",
    "                itr_start_time = time.time()    \n",
    "                \n",
    "            if itr % args.valid_every == 0 and False:\n",
    "                logger.info('Validation ')\n",
    "                valid_loader=torch.utils.data.DataLoader(dataset=valid_set, batch_size=config['batch_size'], shuffle=True, num_workers=1)\n",
    "                model.eval()    \n",
    "                valid_losses = []\n",
    "                for context, context_lens, utt_lens, floors, response, res_lens in valid_loader:\n",
    " #                   max_ctx_len = max(context_lens)\n",
    "                    max_ctx_len = context.size(1)\n",
    "                    context, utt_lens = context[:,:max_ctx_len,1:], utt_lens[:,:max_ctx_len]-1\n",
    "                             # remove empty utterances in context\n",
    "                             # remove the sos token in the context and reduce the context length\n",
    "#################################################\n",
    "                    utt_lens[utt_lens<=0]=1\n",
    "#################################################\n",
    "                    batch = [tensor.to(device) for tensor in [context, context_lens, utt_lens, response, res_lens]]\n",
    "                    valid_results = model.valid(*batch)    \n",
    "                    valid_losses.append(valid_results['valid_loss'])\n",
    "                if args.visual: tb_writer.add_scalar('valid_loss', np.mean(valid_losses), itr_global)\n",
    "                logger.info({'valid_loss':np.mean(valid_losses)})    \n",
    "                \n",
    "            itr += 1\n",
    "            itr_global+=1            \n",
    "            \n",
    "            if itr_global % args.eval_every == 0:  # evaluate the model in the validation set\n",
    "                model.eval()          \n",
    "                logger.info(\"Evaluating in the validation set..\")\n",
    "\n",
    "                valid_loader=torch.utils.data.DataLoader(dataset=valid_set, batch_size=1, shuffle=False, num_workers=1)\n",
    "\n",
    "                f_eval = open(f\"./output/tmp_results/iter{itr_global}.txt\", \"w\")\n",
    "                repeat = 10            \n",
    "                eval_results = evaluate(model, metrics, valid_loader, vocab, repeat, f_eval)\n",
    "                bleu = eval_results['recall_bleu']\n",
    "                if bleu> best_perf:\n",
    "                    save_model(model, 0)#itr_global) # save model after each epoch\n",
    "                if args.visual:\n",
    "                    tb_writer.add_scalar('recall_bleu', bleu, itr_global)\n",
    "                \n",
    "        # end of epoch ----------------------------\n",
    "               # model.adjust_lr()\n",
    "\n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Main Function (Training)\n",
    "You can change the default arguments by setting the `default` attribute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'data_path': './data2/', 'model': 'MyModel', 'dataset': 'weibo', 'visual': True, 'reload_from': -1, 'gpu_id': 0, 'log_every': 100, 'valid_every': 1000, 'eval_every': 2000, 'seed': 1111}\n",
      "cuda:0\n",
      "loading data...\n",
      "76052 entries\n",
      "loading data...\n",
      "7069 entries\n",
      "loading data...\n",
      "7069 entries\n",
      "Loaded data!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training...\n",
      "/opt/conda/envs/pytorch_py3/lib/python3.6/site-packages/ipykernel/__main__.py:138: TracerWarning: Converting a tensor to a Python list might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "/opt/conda/envs/pytorch_py3/lib/python3.6/site-packages/ipykernel/__main__.py:138: UserWarning: pack_padded_sequence has been called with a Python list of sequence lengths. The tracer cannot track the data flow of Python values, and it will treat them as constants, likely rendering the trace incorrect for any other combination of lengths.\n",
      "MyModel|weibo@gpu0 epo:[1/10] iter:[100/1188] step_time:4s elapsed:0:0:4<0:0:49\n",
      "{'train_loss': 5.529356002807617}\n",
      "MyModel|weibo@gpu0 epo:[1/10] iter:[200/1188] step_time:3s elapsed:0:0:8<0:0:40\n",
      "{'train_loss': 5.125542640686035}\n",
      "MyModel|weibo@gpu0 epo:[1/10] iter:[300/1188] step_time:3s elapsed:0:0:11<0:0:35\n",
      "{'train_loss': 4.947093486785889}\n",
      "MyModel|weibo@gpu0 epo:[1/10] iter:[400/1188] step_time:3s elapsed:0:0:15<0:0:30\n",
      "{'train_loss': 4.696111679077148}\n",
      "MyModel|weibo@gpu0 epo:[1/10] iter:[500/1188] step_time:3s elapsed:0:0:19<0:0:26\n",
      "{'train_loss': 4.543099880218506}\n",
      "MyModel|weibo@gpu0 epo:[1/10] iter:[600/1188] step_time:3s elapsed:0:0:23<0:0:22\n",
      "{'train_loss': 4.818960189819336}\n",
      "MyModel|weibo@gpu0 epo:[1/10] iter:[700/1188] step_time:3s elapsed:0:0:26<0:0:18\n",
      "{'train_loss': 4.486680507659912}\n",
      "MyModel|weibo@gpu0 epo:[1/10] iter:[800/1188] step_time:3s elapsed:0:0:30<0:0:14\n",
      "{'train_loss': 4.598004341125488}\n",
      "MyModel|weibo@gpu0 epo:[1/10] iter:[900/1188] step_time:3s elapsed:0:0:34<0:0:10\n",
      "{'train_loss': 4.32649040222168}\n",
      "MyModel|weibo@gpu0 epo:[1/10] iter:[1000/1188] step_time:3s elapsed:0:0:37<0:0:7\n",
      "{'train_loss': 4.380548477172852}\n",
      "MyModel|weibo@gpu0 epo:[1/10] iter:[1100/1188] step_time:3s elapsed:0:0:41<0:0:3\n",
      "{'train_loss': 4.3980183601379395}\n",
      "MyModel|weibo@gpu0 epo:[2/10] iter:[100/1188] step_time:3s elapsed:0:0:3<0:0:40\n",
      "{'train_loss': 4.167244911193848}\n",
      "MyModel|weibo@gpu0 epo:[2/10] iter:[200/1188] step_time:3s elapsed:0:0:7<0:0:36\n",
      "{'train_loss': 4.187685012817383}\n",
      "MyModel|weibo@gpu0 epo:[2/10] iter:[300/1188] step_time:3s elapsed:0:0:11<0:0:33\n",
      "{'train_loss': 4.005294322967529}\n",
      "MyModel|weibo@gpu0 epo:[2/10] iter:[400/1188] step_time:3s elapsed:0:0:14<0:0:29\n",
      "{'train_loss': 4.17413330078125}\n",
      "MyModel|weibo@gpu0 epo:[2/10] iter:[500/1188] step_time:3s elapsed:0:0:18<0:0:25\n",
      "{'train_loss': 4.314266681671143}\n",
      "MyModel|weibo@gpu0 epo:[2/10] iter:[600/1188] step_time:3s elapsed:0:0:22<0:0:21\n",
      "{'train_loss': 4.102077960968018}\n",
      "MyModel|weibo@gpu0 epo:[2/10] iter:[700/1188] step_time:3s elapsed:0:0:25<0:0:18\n",
      "{'train_loss': 4.178990840911865}\n",
      "MyModel|weibo@gpu0 epo:[2/10] iter:[800/1188] step_time:3s elapsed:0:0:29<0:0:14\n",
      "{'train_loss': 4.124073505401611}\n",
      "Evaluating in the validation set..\n",
      " 71%|███████   | 5001/7069 [01:30<00:37, 55.11it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done testing\n",
      "{'avg_len': 8.373125374925015, 'recall_bleu': 0.30624121766100043, 'prec_bleu': 0.30624121766100043, 'f1_bleu': 0.30624121765600043}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "MyModel|weibo@gpu0 epo:[2/10] iter:[900/1188] step_time:94s elapsed:0:2:4<0:0:39\n",
      "{'train_loss': 4.083339691162109}\n",
      "MyModel|weibo@gpu0 epo:[2/10] iter:[1000/1188] step_time:3s elapsed:0:2:7<0:0:24\n",
      "{'train_loss': 4.162360191345215}\n",
      "MyModel|weibo@gpu0 epo:[2/10] iter:[1100/1188] step_time:3s elapsed:0:2:11<0:0:10\n",
      "{'train_loss': 3.9043235778808594}\n",
      "MyModel|weibo@gpu0 epo:[3/10] iter:[100/1188] step_time:3s elapsed:0:0:3<0:0:40\n",
      "{'train_loss': 3.9930450916290283}\n",
      "MyModel|weibo@gpu0 epo:[3/10] iter:[200/1188] step_time:3s elapsed:0:0:7<0:0:37\n",
      "{'train_loss': 4.082245826721191}\n",
      "MyModel|weibo@gpu0 epo:[3/10] iter:[300/1188] step_time:3s elapsed:0:0:11<0:0:33\n",
      "{'train_loss': 4.193918704986572}\n",
      "MyModel|weibo@gpu0 epo:[3/10] iter:[400/1188] step_time:3s elapsed:0:0:14<0:0:29\n",
      "{'train_loss': 3.9984970092773438}\n",
      "MyModel|weibo@gpu0 epo:[3/10] iter:[500/1188] step_time:3s elapsed:0:0:18<0:0:25\n",
      "{'train_loss': 4.1031999588012695}\n",
      "MyModel|weibo@gpu0 epo:[3/10] iter:[600/1188] step_time:3s elapsed:0:0:22<0:0:21\n",
      "{'train_loss': 3.992549180984497}\n",
      "MyModel|weibo@gpu0 epo:[3/10] iter:[700/1188] step_time:3s elapsed:0:0:25<0:0:18\n",
      "{'train_loss': 4.26853084564209}\n",
      "MyModel|weibo@gpu0 epo:[3/10] iter:[800/1188] step_time:3s elapsed:0:0:29<0:0:14\n",
      "{'train_loss': 4.1404290199279785}\n",
      "MyModel|weibo@gpu0 epo:[3/10] iter:[900/1188] step_time:3s elapsed:0:0:33<0:0:10\n",
      "{'train_loss': 3.845642328262329}\n",
      "MyModel|weibo@gpu0 epo:[3/10] iter:[1000/1188] step_time:3s elapsed:0:0:37<0:0:6\n",
      "{'train_loss': 3.698251485824585}\n",
      "MyModel|weibo@gpu0 epo:[3/10] iter:[1100/1188] step_time:3s elapsed:0:0:40<0:0:3\n",
      "{'train_loss': 3.942115306854248}\n",
      "MyModel|weibo@gpu0 epo:[4/10] iter:[100/1188] step_time:3s elapsed:0:0:3<0:0:40\n",
      "{'train_loss': 3.9470624923706055}\n",
      "MyModel|weibo@gpu0 epo:[4/10] iter:[200/1188] step_time:3s elapsed:0:0:7<0:0:36\n",
      "{'train_loss': 3.873811960220337}\n",
      "MyModel|weibo@gpu0 epo:[4/10] iter:[300/1188] step_time:3s elapsed:0:0:11<0:0:32\n",
      "{'train_loss': 3.7955269813537598}\n",
      "MyModel|weibo@gpu0 epo:[4/10] iter:[400/1188] step_time:3s elapsed:0:0:14<0:0:29\n",
      "{'train_loss': 3.5860140323638916}\n",
      "Evaluating in the validation set..\n",
      " 71%|███████   | 5001/7069 [01:30<00:37, 55.01it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done testing\n",
      "{'avg_len': 8.291141771645671, 'recall_bleu': 0.3139334681137361, 'prec_bleu': 0.3139334681137361, 'f1_bleu': 0.3139334681087361}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "MyModel|weibo@gpu0 epo:[4/10] iter:[500/1188] step_time:94s elapsed:0:1:49<0:2:30\n",
      "{'train_loss': 3.991703748703003}\n",
      "MyModel|weibo@gpu0 epo:[4/10] iter:[600/1188] step_time:3s elapsed:0:1:53<0:1:51\n",
      "{'train_loss': 3.9677443504333496}\n",
      "MyModel|weibo@gpu0 epo:[4/10] iter:[700/1188] step_time:3s elapsed:0:1:57<0:1:21\n",
      "{'train_loss': 3.878000020980835}\n",
      "MyModel|weibo@gpu0 epo:[4/10] iter:[800/1188] step_time:3s elapsed:0:2:0<0:0:58\n",
      "{'train_loss': 4.135827541351318}\n",
      "MyModel|weibo@gpu0 epo:[4/10] iter:[900/1188] step_time:3s elapsed:0:2:4<0:0:39\n",
      "{'train_loss': 3.533134937286377}\n",
      "MyModel|weibo@gpu0 epo:[4/10] iter:[1000/1188] step_time:3s elapsed:0:2:8<0:0:24\n",
      "{'train_loss': 3.5958003997802734}\n",
      "MyModel|weibo@gpu0 epo:[4/10] iter:[1100/1188] step_time:3s elapsed:0:2:11<0:0:10\n",
      "{'train_loss': 3.8423359394073486}\n",
      "MyModel|weibo@gpu0 epo:[5/10] iter:[100/1188] step_time:3s elapsed:0:0:3<0:0:40\n",
      "{'train_loss': 3.7531774044036865}\n",
      "MyModel|weibo@gpu0 epo:[5/10] iter:[200/1188] step_time:3s elapsed:0:0:7<0:0:37\n",
      "{'train_loss': 3.7024171352386475}\n",
      "MyModel|weibo@gpu0 epo:[5/10] iter:[300/1188] step_time:3s elapsed:0:0:11<0:0:33\n",
      "{'train_loss': 3.515455484390259}\n",
      "MyModel|weibo@gpu0 epo:[5/10] iter:[400/1188] step_time:3s elapsed:0:0:14<0:0:29\n",
      "{'train_loss': 3.516704559326172}\n",
      "MyModel|weibo@gpu0 epo:[5/10] iter:[500/1188] step_time:3s elapsed:0:0:18<0:0:25\n",
      "{'train_loss': 3.774460792541504}\n",
      "MyModel|weibo@gpu0 epo:[5/10] iter:[600/1188] step_time:3s elapsed:0:0:22<0:0:21\n",
      "{'train_loss': 3.518766403198242}\n",
      "MyModel|weibo@gpu0 epo:[5/10] iter:[700/1188] step_time:3s elapsed:0:0:26<0:0:18\n",
      "{'train_loss': 3.5113394260406494}\n",
      "MyModel|weibo@gpu0 epo:[5/10] iter:[800/1188] step_time:3s elapsed:0:0:29<0:0:14\n",
      "{'train_loss': 3.8112027645111084}\n",
      "MyModel|weibo@gpu0 epo:[5/10] iter:[900/1188] step_time:3s elapsed:0:0:33<0:0:10\n",
      "{'train_loss': 3.6757986545562744}\n",
      "MyModel|weibo@gpu0 epo:[5/10] iter:[1000/1188] step_time:3s elapsed:0:0:37<0:0:6\n",
      "{'train_loss': 3.662677526473999}\n",
      "MyModel|weibo@gpu0 epo:[5/10] iter:[1100/1188] step_time:3s elapsed:0:0:40<0:0:3\n",
      "{'train_loss': 3.512507200241089}\n",
      "Evaluating in the validation set..\n",
      " 71%|███████   | 5001/7069 [01:30<00:37, 55.17it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done testing\n",
      "{'avg_len': 8.430913817236553, 'recall_bleu': 0.30944281176172556, 'prec_bleu': 0.30944281176172556, 'f1_bleu': 0.30944281175672556}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "MyModel|weibo@gpu0 epo:[6/10] iter:[100/1188] step_time:94s elapsed:0:1:34<0:17:8\n",
      "{'train_loss': 3.5358195304870605}\n",
      "MyModel|weibo@gpu0 epo:[6/10] iter:[200/1188] step_time:3s elapsed:0:1:38<0:8:5\n",
      "{'train_loss': 3.4332456588745117}\n",
      "MyModel|weibo@gpu0 epo:[6/10] iter:[300/1188] step_time:3s elapsed:0:1:41<0:5:1\n",
      "{'train_loss': 3.5960922241210938}\n",
      "MyModel|weibo@gpu0 epo:[6/10] iter:[400/1188] step_time:3s elapsed:0:1:45<0:3:28\n",
      "{'train_loss': 3.845384120941162}\n",
      "MyModel|weibo@gpu0 epo:[6/10] iter:[500/1188] step_time:3s elapsed:0:1:49<0:2:30\n",
      "{'train_loss': 3.6513259410858154}\n",
      "MyModel|weibo@gpu0 epo:[6/10] iter:[600/1188] step_time:3s elapsed:0:1:53<0:1:50\n",
      "{'train_loss': 3.6359498500823975}\n",
      "MyModel|weibo@gpu0 epo:[6/10] iter:[700/1188] step_time:3s elapsed:0:1:56<0:1:21\n",
      "{'train_loss': 3.6793899536132812}\n",
      "MyModel|weibo@gpu0 epo:[6/10] iter:[800/1188] step_time:3s elapsed:0:2:0<0:0:58\n",
      "{'train_loss': 3.673205852508545}\n",
      "MyModel|weibo@gpu0 epo:[6/10] iter:[900/1188] step_time:3s elapsed:0:2:4<0:0:39\n",
      "{'train_loss': 3.7047626972198486}\n",
      "MyModel|weibo@gpu0 epo:[6/10] iter:[1000/1188] step_time:3s elapsed:0:2:7<0:0:24\n",
      "{'train_loss': 3.632887125015259}\n",
      "MyModel|weibo@gpu0 epo:[6/10] iter:[1100/1188] step_time:3s elapsed:0:2:11<0:0:10\n",
      "{'train_loss': 3.6154098510742188}\n",
      "MyModel|weibo@gpu0 epo:[7/10] iter:[100/1188] step_time:3s elapsed:0:0:3<0:0:40\n",
      "{'train_loss': 3.433790922164917}\n",
      "MyModel|weibo@gpu0 epo:[7/10] iter:[200/1188] step_time:3s elapsed:0:0:7<0:0:36\n",
      "{'train_loss': 3.4456982612609863}\n",
      "MyModel|weibo@gpu0 epo:[7/10] iter:[300/1188] step_time:3s elapsed:0:0:11<0:0:33\n",
      "{'train_loss': 3.279827356338501}\n",
      "MyModel|weibo@gpu0 epo:[7/10] iter:[400/1188] step_time:3s elapsed:0:0:14<0:0:29\n",
      "{'train_loss': 3.3232622146606445}\n",
      "MyModel|weibo@gpu0 epo:[7/10] iter:[500/1188] step_time:3s elapsed:0:0:18<0:0:25\n",
      "{'train_loss': 3.5868260860443115}\n",
      "MyModel|weibo@gpu0 epo:[7/10] iter:[600/1188] step_time:3s elapsed:0:0:22<0:0:21\n",
      "{'train_loss': 3.4944944381713867}\n",
      "MyModel|weibo@gpu0 epo:[7/10] iter:[700/1188] step_time:3s elapsed:0:0:26<0:0:18\n",
      "{'train_loss': 3.3088834285736084}\n",
      "MyModel|weibo@gpu0 epo:[7/10] iter:[800/1188] step_time:3s elapsed:0:0:29<0:0:14\n",
      "{'train_loss': 3.644766330718994}\n",
      "Evaluating in the validation set..\n",
      " 71%|███████   | 5001/7069 [01:30<00:37, 55.34it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done testing\n",
      "{'avg_len': 8.473705258948211, 'recall_bleu': 0.3085645340566451, 'prec_bleu': 0.3085645340566451, 'f1_bleu': 0.3085645340516451}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "MyModel|weibo@gpu0 epo:[7/10] iter:[900/1188] step_time:94s elapsed:0:2:4<0:0:39\n",
      "{'train_loss': 3.3136956691741943}\n",
      "MyModel|weibo@gpu0 epo:[7/10] iter:[1000/1188] step_time:3s elapsed:0:2:7<0:0:24\n",
      "{'train_loss': 3.6939213275909424}\n",
      "MyModel|weibo@gpu0 epo:[7/10] iter:[1100/1188] step_time:3s elapsed:0:2:11<0:0:10\n",
      "{'train_loss': 3.533986806869507}\n",
      "MyModel|weibo@gpu0 epo:[8/10] iter:[100/1188] step_time:3s elapsed:0:0:3<0:0:40\n",
      "{'train_loss': 3.358433246612549}\n",
      "MyModel|weibo@gpu0 epo:[8/10] iter:[200/1188] step_time:3s elapsed:0:0:7<0:0:36\n",
      "{'train_loss': 3.487255096435547}\n",
      "MyModel|weibo@gpu0 epo:[8/10] iter:[300/1188] step_time:3s elapsed:0:0:11<0:0:32\n",
      "{'train_loss': 3.503608226776123}\n",
      "MyModel|weibo@gpu0 epo:[8/10] iter:[400/1188] step_time:3s elapsed:0:0:14<0:0:28\n",
      "{'train_loss': 3.2705886363983154}\n",
      "MyModel|weibo@gpu0 epo:[8/10] iter:[500/1188] step_time:3s elapsed:0:0:18<0:0:25\n",
      "{'train_loss': 3.373830795288086}\n",
      "MyModel|weibo@gpu0 epo:[8/10] iter:[600/1188] step_time:3s elapsed:0:0:22<0:0:21\n",
      "{'train_loss': 3.264432430267334}\n",
      "MyModel|weibo@gpu0 epo:[8/10] iter:[700/1188] step_time:3s elapsed:0:0:25<0:0:17\n",
      "{'train_loss': 3.4094693660736084}\n",
      "MyModel|weibo@gpu0 epo:[8/10] iter:[800/1188] step_time:3s elapsed:0:0:29<0:0:14\n",
      "{'train_loss': 3.4738399982452393}\n",
      "MyModel|weibo@gpu0 epo:[8/10] iter:[900/1188] step_time:3s elapsed:0:0:32<0:0:10\n",
      "{'train_loss': 3.3389391899108887}\n",
      "MyModel|weibo@gpu0 epo:[8/10] iter:[1000/1188] step_time:3s elapsed:0:0:36<0:0:6\n",
      "{'train_loss': 3.4993550777435303}\n",
      "MyModel|weibo@gpu0 epo:[8/10] iter:[1100/1188] step_time:3s elapsed:0:0:40<0:0:3\n",
      "{'train_loss': 3.585380792617798}\n",
      "MyModel|weibo@gpu0 epo:[9/10] iter:[100/1188] step_time:3s elapsed:0:0:3<0:0:40\n",
      "{'train_loss': 3.326112985610962}\n",
      "MyModel|weibo@gpu0 epo:[9/10] iter:[200/1188] step_time:3s elapsed:0:0:7<0:0:36\n",
      "{'train_loss': 3.277453899383545}\n",
      "MyModel|weibo@gpu0 epo:[9/10] iter:[300/1188] step_time:3s elapsed:0:0:11<0:0:32\n",
      "{'train_loss': 3.5662989616394043}\n",
      "MyModel|weibo@gpu0 epo:[9/10] iter:[400/1188] step_time:3s elapsed:0:0:14<0:0:29\n",
      "{'train_loss': 3.3841631412506104}\n",
      "Evaluating in the validation set..\n",
      " 71%|███████   | 5001/7069 [01:28<00:36, 56.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done testing\n",
      "{'avg_len': 8.345930813837233, 'recall_bleu': 0.31324515458029106, 'prec_bleu': 0.31324515458029106, 'f1_bleu': 0.31324515457529106}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "MyModel|weibo@gpu0 epo:[9/10] iter:[500/1188] step_time:92s elapsed:0:1:47<0:2:27\n",
      "{'train_loss': 3.265955686569214}\n",
      "MyModel|weibo@gpu0 epo:[9/10] iter:[600/1188] step_time:3s elapsed:0:1:51<0:1:48\n",
      "{'train_loss': 3.562735080718994}\n",
      "MyModel|weibo@gpu0 epo:[9/10] iter:[700/1188] step_time:3s elapsed:0:1:54<0:1:19\n",
      "{'train_loss': 3.4401681423187256}\n",
      "MyModel|weibo@gpu0 epo:[9/10] iter:[800/1188] step_time:3s elapsed:0:1:58<0:0:57\n",
      "{'train_loss': 3.152111768722534}\n",
      "MyModel|weibo@gpu0 epo:[9/10] iter:[900/1188] step_time:3s elapsed:0:2:2<0:0:39\n",
      "{'train_loss': 3.5151991844177246}\n",
      "MyModel|weibo@gpu0 epo:[9/10] iter:[1000/1188] step_time:3s elapsed:0:2:5<0:0:23\n",
      "{'train_loss': 3.7213807106018066}\n",
      "MyModel|weibo@gpu0 epo:[9/10] iter:[1100/1188] step_time:3s elapsed:0:2:9<0:0:10\n",
      "{'train_loss': 3.54030179977417}\n",
      "MyModel|weibo@gpu0 epo:[10/10] iter:[100/1188] step_time:3s elapsed:0:0:3<0:0:40\n",
      "{'train_loss': 3.243582248687744}\n",
      "MyModel|weibo@gpu0 epo:[10/10] iter:[200/1188] step_time:3s elapsed:0:0:7<0:0:36\n",
      "{'train_loss': 3.3471577167510986}\n",
      "MyModel|weibo@gpu0 epo:[10/10] iter:[300/1188] step_time:3s elapsed:0:0:11<0:0:32\n",
      "{'train_loss': 3.1908364295959473}\n",
      "MyModel|weibo@gpu0 epo:[10/10] iter:[400/1188] step_time:3s elapsed:0:0:14<0:0:29\n",
      "{'train_loss': 3.280076742172241}\n",
      "MyModel|weibo@gpu0 epo:[10/10] iter:[500/1188] step_time:3s elapsed:0:0:18<0:0:25\n",
      "{'train_loss': 3.237961530685425}\n",
      "MyModel|weibo@gpu0 epo:[10/10] iter:[600/1188] step_time:3s elapsed:0:0:22<0:0:21\n",
      "{'train_loss': 3.299494504928589}\n",
      "MyModel|weibo@gpu0 epo:[10/10] iter:[700/1188] step_time:3s elapsed:0:0:25<0:0:17\n",
      "{'train_loss': 3.377887725830078}\n",
      "MyModel|weibo@gpu0 epo:[10/10] iter:[800/1188] step_time:3s elapsed:0:0:29<0:0:14\n",
      "{'train_loss': 3.1742634773254395}\n",
      "MyModel|weibo@gpu0 epo:[10/10] iter:[900/1188] step_time:3s elapsed:0:0:33<0:0:10\n",
      "{'train_loss': 3.3920092582702637}\n",
      "MyModel|weibo@gpu0 epo:[10/10] iter:[1000/1188] step_time:3s elapsed:0:0:36<0:0:6\n",
      "{'train_loss': 3.067640542984009}\n",
      "MyModel|weibo@gpu0 epo:[10/10] iter:[1100/1188] step_time:3s elapsed:0:0:40<0:0:3\n",
      "{'train_loss': 3.2566628456115723}\n"
     ]
    }
   ],
   "source": [
    "\n",
    "if __name__ == '__main__':\n",
    "\n",
    "    parser = argparse.ArgumentParser(description='Dialog Pytorch')\n",
    "    # Path Arguments\n",
    "    parser.add_argument('--data_path', type=str, default='./data2/', help='location of the data corpus')\n",
    "    parser.add_argument('--model', type=str, default='MyModel', help='model name')\n",
    "    parser.add_argument('--dataset', type=str, default='weibo', help='name of dataset.')\n",
    "    parser.add_argument('-v','--visual', action='store_true', default=True, help='visualize training status in tensorboard')\n",
    "    parser.add_argument('--reload_from', type=int, default=-1, help='reload from a trained ephoch')\n",
    "    parser.add_argument('--gpu_id', type=int, default=0, help='GPU ID')\n",
    "\n",
    "    # Evaluation Arguments\n",
    "    parser.add_argument('--log_every', type=int, default=100, help='interval to log autoencoder training results')\n",
    "    parser.add_argument('--valid_every', type=int, default=1000, help='interval to validation')\n",
    "    parser.add_argument('--eval_every', type=int, default=2000, help='interval to evaluation to concrete results')\n",
    "    parser.add_argument('--seed', type=int, default=1111, help='random seed')\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    args = parser.parse_args(args=[])\n",
    "    print(vars(args))\n",
    "\n",
    "    # make output directory if it doesn't already exist\n",
    "    os.makedirs(f'./output/models', exist_ok=True)\n",
    "    os.makedirs(f'./output/tmp_results', exist_ok=True)\n",
    "        \n",
    "    torch.backends.cudnn.benchmark = True # speed up training by using cudnn\n",
    "    torch.backends.cudnn.deterministic = True # fix the random seed in cudnn\n",
    "    \n",
    "    model = train(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Main Function (Test)\n",
    "\n",
    "**Please do not change code here except the default arguments**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'data_path': './data2/', 'model': 'MyModel', 'reload_from': 0, 'n_samples': 10, 'seed': 1111}\n",
      "loading data...\n",
      "7069 entries\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 71%|███████   | 5001/7069 [02:28<01:01, 33.66it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done testing\n",
      "{'avg_len': 8.360927814437112, 'recall_bleu': 0.3130616087813351, 'prec_bleu': 0.3130616087813351, 'f1_bleu': 0.3130616087763351}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "\n",
    "def test(args):\n",
    "    conf = get_config()\n",
    "    # Set the random seed manually for reproducibility.\n",
    "    random.seed(args.seed)\n",
    "    np.random.seed(args.seed)\n",
    "    torch.manual_seed(args.seed)\n",
    "    if torch.cuda.is_available():\n",
    "        torch.cuda.manual_seed(args.seed)\n",
    "    else:\n",
    "        print(\"Note that our pre-trained models require CUDA to evaluate.\")\n",
    "    \n",
    "    # Load data\n",
    "    data_path=args.data_path+'/'\n",
    "    test_set=DialogDataset(data_path+'test.h5', conf['diaglen'], conf['maxlen'])\n",
    "    test_loader=torch.utils.data.DataLoader(dataset=test_set, batch_size=1, shuffle=False, num_workers=1)\n",
    "    vocab = load_dict(data_path+'vocab.json')\n",
    "    n_tokens = len(vocab)\n",
    "\n",
    "    metrics=Metrics()\n",
    "    \n",
    "    # Load model checkpoints    \n",
    "    model = MyModel(conf, n_tokens)\n",
    "    load_model(model, 0)\n",
    "    #model=model.to(device)\n",
    "    model.eval()\n",
    "    \n",
    "    f_eval = open(\"./output/results.txt\", \"w\")\n",
    "    repeat = args.n_samples\n",
    "    \n",
    "    evaluate(model, metrics, test_loader, vocab, repeat, f_eval)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    parser = argparse.ArgumentParser(description='PyTorch DialogGAN for Eval')\n",
    "    parser.add_argument('--data_path', type=str, default='./data2/', help='location of the data corpus')\n",
    "#     parser.add_argument('--dataset', type=str, default='weibo', help='name of dataset, SWDA or DailyDial')\n",
    "    parser.add_argument('--model', type=str, default='MyModel', help='model name')\n",
    "    parser.add_argument('--reload_from', type=int, default=0, \n",
    "                        help='directory to load models from')\n",
    "    \n",
    "    parser.add_argument('--n_samples', type=int, default=10, help='2Number of responses to sampling')\n",
    "    parser.add_argument('--seed', type=int, default=1111, help='random seed')\n",
    "    args = parser.parse_args(args=[])\n",
    "    print(vars(args))\n",
    "    test(args)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_py3",
   "language": "python",
   "name": "conda_pytorch_py3"
  },
  "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"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "source": [],
    "metadata": {
     "collapsed": false
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}