SE125Proj2_baseline.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A Baseline Implementation for SE125 Project 2\n",
    "\n",
    "We provide a baseline 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": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Requirement already satisfied: tqdm in c:\\users\\54600\\anaconda3\\lib\\site-packages (4.50.2)\n",
      "Requirement already satisfied: cython in c:\\users\\54600\\anaconda3\\lib\\site-packages (0.29.21)\n",
      "Requirement already satisfied: tables in c:\\users\\54600\\anaconda3\\lib\\site-packages (3.6.1)\n",
      "Requirement already satisfied: numpy>=1.9.3 in c:\\users\\54600\\anaconda3\\lib\\site-packages (from tables) (1.19.2)\n",
      "Requirement already satisfied: numexpr>=2.6.2 in c:\\users\\54600\\anaconda3\\lib\\site-packages (from tables) (2.7.1)\n",
      "Requirement already satisfied: tensorboardX in c:\\users\\54600\\anaconda3\\lib\\site-packages (2.1)\n",
      "Requirement already satisfied: protobuf>=3.8.0 in c:\\users\\54600\\anaconda3\\lib\\site-packages (from tensorboardX) (3.14.0)\n",
      "Requirement already satisfied: numpy in c:\\users\\54600\\anaconda3\\lib\\site-packages (from tensorboardX) (1.19.2)\n",
      "Requirement already satisfied: six in c:\\users\\54600\\anaconda3\\lib\\site-packages (from tensorboardX) (1.15.0)\n",
      "Requirement already satisfied: nltk in c:\\users\\54600\\anaconda3\\lib\\site-packages (3.5)\n",
      "Requirement already satisfied: joblib in c:\\users\\54600\\anaconda3\\lib\\site-packages (from nltk) (0.17.0)\n",
      "Requirement already satisfied: regex in c:\\users\\54600\\anaconda3\\lib\\site-packages (from nltk) (2020.10.15)\n",
      "Requirement already satisfied: tqdm in c:\\users\\54600\\anaconda3\\lib\\site-packages (from nltk) (4.50.2)\n",
      "Requirement already satisfied: click in c:\\users\\54600\\anaconda3\\lib\\site-packages (from nltk) (7.1.2)\n"
     ]
    }
   ],
   "source": [
    "!pip install tqdm\n",
    "!pip install cython\n",
    "!pip install tables\n",
    "!pip install tensorboardX\n",
    "!pip install nltk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "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": {},
   "outputs": [],
   "source": [
    "# Padding, StartOfSentence, EndOfSentence, Unknown(OutOfVocabulary)\n",
    "PAD_ID, SOS_ID, EOS_ID, UNK_ID = [0, 1, 2, 3]\n",
    "\n",
    "# time format\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",
    "\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": {},
   "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": {},
   "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",
    "        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",
    "        \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": {},
   "outputs": [],
   "source": [
    "import torch.nn.init as weight_init\n",
    "from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence\n",
    "from tensorboardX import SummaryWriter\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",
    "        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",
    "        enc = h_n.view(batch_size,-1) #[batch_sz x (n_dirs*hid_sz)]\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",
    "\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",
    "        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'], True, \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 \n",
    "\n",
    "model = MyModel()\n",
    "dummy_input = torch.rand(20, 1, 28, 28)  # 假设输入20张1*28*28的图片\n",
    "with SummaryWriter(comment='MyModel') as w:\n",
    "    w.add_graph(model, (dummy_input,))"
   ]
  },
  {
   "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": 8,
   "metadata": {},
   "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": 9,
   "metadata": {},
   "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}\") 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+args.dataset+'/'\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",
    "    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",
    "            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": 10,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "{'data_path': './data/', 'model': 'MyModel', 'dataset': 'dailydialog', '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",
      "6740 entries\n",
      "Loaded data!\n",
      "Training...\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='./data/', 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('--dataset', type=str, default='dailydialog', help='name of dataset.')\n",
    "    # parser.add_argument('-v','--visual', action='store_true', default=False, help='visualize training status in tensorboard')\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=1, help='GPU ID')\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": {},
   "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": null,
   "metadata": {},
   "outputs": [],
   "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+args.dataset+'/'\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='./data/', 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('--dataset', type=str, default='dailydialog', 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='Number 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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3",
   "language": "python"
  },
  "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.8.5-final"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}