text: import sys
  import torch
  
  import argparse
  from data_utils import load_dialog_task, vectorize_data, load_candidates, vectorize_candidates, tokenize
  from six.moves import range, reduce
  from itertools import chain
  import numpy as np
  import os
  from sklearn import metrics
  from torch.autograd import Variable as V
  
  from model.mem_cnn_sim import MemCnnSim
  
  
  def init(data_dir, task_id, OOV=False):
      # load candidates
      candidates, candid2indx = load_candidates(
          data_dir, task_id)
      n_cand = len(candidates)
      print("Candidate Size", n_cand)
      indx2candid = dict(
          (candid2indx[key], key) for key in candid2indx)
  
      # load task data
      train_data, test_data, val_data = load_dialog_task(
          data_dir, task_id, candid2indx, OOV)
      data = train_data + test_data + val_data
  
      # build parameters
      word_idx, sentence_size, \
      candidate_sentence_size, memory_size, \
      vocab_size = build_vocab(data, candidates)
  
      # Variable(torch.from_numpy(candidates_vec)).view(len(candidates), sentence_size)
      candidates_vec = vectorize_candidates(
          candidates, word_idx, candidate_sentence_size)
  
      return candid2indx, \
             indx2candid, \
             candidates_vec, \
             word_idx, \
             sentence_size, \
             candidate_sentence_size, \
             memory_size, \
             vocab_size, \
             train_data, test_data, val_data
  
  
  def build_vocab(data, candidates, memory_size=50):
      vocab = reduce(lambda x, y: x | y, (set(
          list(chain.from_iterable(s)) + q) for s, q, a in data))
      vocab |= reduce(lambda x, y: x | y, (set(candidate)
                                           for candidate in candidates))
      vocab = sorted(vocab)
      word_idx = dict((c, i + 1) for i, c in enumerate(vocab))
  
      max_story_size = max(map(len, (s for s, _, _ in data)))
      mean_story_size = int(np.mean([len(s) for s, _, _ in data]))
      sentence_size = max(map(len, chain.from_iterable(s for s, _, _ in data)))
      candidate_sentence_size = max(map(len, candidates))
      query_size = max(map(len, (q for _, q, _ in data)))
      memory_size = min(memory_size, max_story_size)
      vocab_size = len(word_idx) + 1  # +1 for nil word
      sentence_size = max(query_size, sentence_size)  # for the position
      # params
      print("vocab size:", vocab_size)
      print("Longest sentence length", sentence_size)
      print("Longest candidate sentence length", candidate_sentence_size)
      print("Longest story length", max_story_size)
      print("Average story length", mean_story_size)
  
      return word_idx, \
             sentence_size, \
             candidate_sentence_size, \
             memory_size, \
             vocab_size
  
  
  def eval(utter_batch, memory_batch, answer__batch, dialog_idx, mem_cnn_sim, cuda=False):
      mem_cnn_sim.eval()
  
      total_loss = []
      preds = []
  
      dialog_len = len(utter_batch)
      for start in range(dialog_len):
          end = start
  
          loss_per_diaglo = []
  
          for j in range(start, end + 1):
  
              memory = V(torch.from_numpy(memory_batch[j])).unsqueeze(0)
              utter = V(torch.from_numpy(utter_batch[j])).unsqueeze(0)
  
              if cuda:
                  memory = transfer_to_gpu(memory)
                  utter = transfer_to_gpu(utter)
  
              context, cand_ = mem_cnn_sim(utter, memory, cands_tensor)
              pred = mem_cnn_sim.predict(context, cand_)
              preds.append(pred.data[0])
  
              # loss_per_diaglo.append(loss.data[0])
  
          total_loss += loss_per_diaglo
  
      accuracy = metrics.accuracy_score(answer__batch[:len(preds)], preds)
      print()
      print('Validation accuracy: {}'.format(accuracy))
      print('Validation loss: {}'.format(sum(total_loss)))
  
      return accuracy
  
  
  def interactive(model, indx2candid, cands_tensor, word_idx, sentence_size, memory_size, cuda=False):
      context = []
      u = None
      r = None
      nid = 1
      while True:
          line = input('--> ').strip().lower()
          if line == 'exit':
              break
          if line == 'restart':
              context = []
              nid = 1
              print("clear memory")
              continue
          u = tokenize(line)
          data = [(context, u, -1)]
          s, q, a, entity_dict = vectorize_data(data, word_idx, sentence_size, memory_size)
  
          memory = V(torch.from_numpy(np.stack(s)))
          utter = V(torch.from_numpy(np.stack(q)))
  
          if cuda:
              memory = transfer_to_gpu(memory)
              utter = transfer_to_gpu(utter)
  
          context_, cand_ = model(utter, memory, cands_tensor)
          preds = model.predict(context_, cand_)
          r = indx2candid[preds.data[0]]
          print(r)
          r = tokenize(r)
          u.append('$u')
          u.append('#' + str(nid))
          r.append('$r')
          r.append('#' + str(nid))
          context.append(u)
          context.append(r)
          nid += 1
  
  
  def transfer_to_gpu(tensor, dtype=torch.LongTensor):
      tensor_cuda = dtype(tensor.size()).cuda()
      tensor_cuda = V(tensor_cuda)
      tensor_cuda.data.copy_(tensor.data)
      return tensor_cuda
  
  
  def save_checkpoint(state, filename='checkpoint.pth.tar'):
      torch.save(state, filename)
  
  
  def load_checkpoit(model, optimizer, path_to_model):
      if os.path.isfile(path_to_model):
          print("=> loading checkpoint '{}'".format(path_to_model))
          checkpoint = torch.load(path_to_model)
          model.load_state_dict(checkpoint['state_dict'])
          optimizer.load_state_dict(checkpoint['optimizer'])
          print("=> loaded checkpoint '{}' (epoch {})"
                .format(path_to_model, checkpoint['epoch']))
      else:
          print("=> no checkpoint found at '{}'".format(path_to_model))
  
  
  def load_model(model, model_dir):
      load_checkpoit(model, model.optimizer, model_dir+'best_model')
  
  
  def test_model(mem_cnn_sim):
  
      for i in range(20):
          utter = V(torch.LongTensor([1,1,1])).unsqueeze(0)
          memory = V(torch.LongTensor([[1,2,3], [4,5,6]])).unsqueeze(0)
          cand = V(torch.LongTensor([[7,8,9], [10,11,12], [13,14,15], [16,17,18]]))
          flag = V(torch.FloatTensor([0,0,0,1]))
  
          context, cand_ = mem_cnn_sim(utter, memory, cand)
          loss = mem_cnn_sim.loss_op(context, cand_, flag)
          pred = mem_cnn_sim.predict(context, cand_)
          mem_cnn_sim.optimize(loss)
  
          print('loss: {}, pred: {}'.format(loss.data[0], pred.data[0]))
  
  
  if __name__ == '__main__':
  
      def get_bool(arg):
          if arg.lower() in ('true', 't', 'yes'):
              return True
          else:
              return False
  
  
      parser = argparse.ArgumentParser()
      parser.add_argument("-interactive", type=get_bool, default=False)
  
      args = parser.parse_args()
  
      data_dir = "data/dialog-bAbI-tasks/"
      task_id = 1
      epochs = 20
      model_dir = "task" + str(task_id) + "_model/"
      if not os.path.exists(model_dir):
          os.makedirs(model_dir)
      test_ = False
  
      cuda = torch.cuda.is_available()
  
      print('task id: {}'.format(task_id))
      if cuda: print('Cuda is available.')
  
      candid2indx, \
      indx2candid, \
      candidates_vec, \
      word_idx, \
      sentence_size, \
      candidate_sentence_size, \
      memory_size, \
      vocab_size, \
      train_data, test_data, val_data = init(data_dir, task_id)
  
      trainS, trainQ, trainA, dialog_idx = vectorize_data(
          train_data, word_idx, sentence_size, memory_size)
      valS, valQ, valA, dialog_idx_val = vectorize_data(
          val_data, word_idx, sentence_size, memory_size)
      n_train = len(trainS)
      n_val = len(valS)
  
      print("Training Size", n_train)
      print("Validation Size", n_val)
  
      param = {
              'hops': 3,
              "vocab_size": vocab_size,
              "embedding_size": 100,
              'num_filters': 100,
              "cand_vocab_size": vocab_size,
              'max_grad_norm': 40.0
               }
  
      mem_cnn_sim = MemCnnSim(param)
  
      if test_:
          test_model(mem_cnn_sim)
          input()
  
      best_validation_accuracy = 0
      time = []
  
      cands_tensor = V(torch.from_numpy(candidates_vec))
      num_cand = cands_tensor.size(0)
      num_dialog = len(dialog_idx)
  
      if cuda:
          mem_cnn_sim.cuda()
          cands_tensor = transfer_to_gpu(cands_tensor)
  
      if args.interactive:
          load_model(mem_cnn_sim, 'task{}_model/'.format(task_id))
          interactive(mem_cnn_sim, indx2candid, cands_tensor, word_idx, sentence_size, memory_size, cuda)
  
      for i in range(1, epochs+1):
          num_ = [x for x in range(len(trainS))]
          np.random.shuffle(num_)
          mem_cnn_sim.train()
          # for j, (start, end) in enumerate(dialog_idx):
          #
          #     if j%99 == 0:
          #         print('[{}/{}]\r'.format(j+1, num_dialog))
          #
          #     loss_per_diaglo = []
  
          for j, k in enumerate(num_):
  
              ans = trainA[k]
  
              memory = V(torch.from_numpy(trainS[k])).unsqueeze(0)
              utter = V(torch.from_numpy(trainQ[k])).unsqueeze(0)
  
              flag = -1 * torch.ones(num_cand)
              flag[ans] = 1
  
              flag = V(flag)
  
              if cuda:
  
                  memory = transfer_to_gpu(memory)
                  utter = transfer_to_gpu(utter)
                  flag = transfer_to_gpu(flag, dtype=torch.FloatTensor)
  
              context, cand_ = mem_cnn_sim(utter, memory, cands_tensor)
              loss = mem_cnn_sim.loss_op(context, cand_, flag)
              mem_cnn_sim.optimize(loss)
  
              if j % 100 == 0:
                  sys.stdout.write('\r{}/{}'.format(j, len(trainS)))
              # loss_per_diaglo.append(loss.data[0])
              # print('loss: {}'.format(sum(loss_per_diaglo)/len(loss_per_diaglo)))
  
          print('\ntrain accuracy')
          accuracy = eval(trainQ, trainS, trainA, dialog_idx, mem_cnn_sim, cuda)
          print('eval accuracy')
          accuracy = eval(valQ, valS, valA, dialog_idx_val, mem_cnn_sim, cuda)
  
          if accuracy > best_validation_accuracy:
              best_validation_accuracy = accuracy
              save_checkpoint({
                  'epoch': i + 1,
                  'state_dict': mem_cnn_sim.state_dict(),
                  'optimizer': mem_cnn_sim.optimizer.state_dict(),
              }, filename=model_dir + 'best_model')
  