text: # This file contains Att2in2, AdaAtt, AdaAttMO, TopDown model
  
  # AdaAtt is from Knowing When to Look: Adaptive Attention via A Visual Sentinel for Image Captioning
  # https://arxiv.org/abs/1612.01887
  # AdaAttMO is a modified version with maxout lstm
  
  # Att2in is from Self-critical Sequence Training for Image Captioning
  # https://arxiv.org/abs/1612.00563
  # In this file we only have Att2in2, which is a slightly different version of att2in,
  # in which the img feature embedding and word embedding is the same as what in adaatt.
  
  # TopDown is from Bottom-Up and Top-Down Attention for Image Captioning and VQA
  # https://arxiv.org/abs/1707.07998
  # However, it may not be identical to the author's architecture.
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import numpy as np
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from torch.autograd import Variable
  import misc.utils as utils
  from torch.nn.utils.rnn import PackedSequence, pack_padded_sequence, pad_packed_sequence
  
  from .CaptionModel import CaptionModel
  from torchtext.vocab import GloVe
  
  bad_endings = ['a','an','the','in','for','at','of','with','before','after','on','upon','near','to','is','are','am']
  bad_endings += ['the']
  
  def sort_pack_padded_sequence(input, lengths):
      sorted_lengths, indices = torch.sort(lengths, descending=True)
      tmp = pack_padded_sequence(input[indices], sorted_lengths, batch_first=True)
      inv_ix = indices.clone()
      inv_ix[indices] = torch.arange(0,len(indices)).type_as(inv_ix)
      return tmp, inv_ix
  
  def pad_unsort_packed_sequence(input, inv_ix):
      tmp, _ = pad_packed_sequence(input, batch_first=True)
      tmp = tmp[inv_ix]
      return tmp
  
  def pack_wrapper(module, att_feats, att_masks):
      if att_masks is not None:
          packed, inv_ix = sort_pack_padded_sequence(att_feats, att_masks.data.long().sum(1))
          return pad_unsort_packed_sequence(PackedSequence(module(packed[0]), packed[1]), inv_ix)
      else:
          return module(att_feats)
  
  class AttModel(CaptionModel):
      def __init__(self, opt):
          super(AttModel, self).__init__()
          self.vocab_size = opt.vocab_size
          self.input_encoding_size = opt.input_encoding_size
          #self.rnn_type = opt.rnn_type
          self.rnn_size = opt.rnn_size
          self.num_layers = opt.num_layers
          self.drop_prob_lm = opt.drop_prob_lm
          self.seq_length = getattr(opt, 'max_length', 20) or opt.seq_length # maximum sample length
          self.fc_feat_size = opt.fc_feat_size
          self.att_feat_size = opt.att_feat_size
          self.att_hid_size = opt.att_hid_size
  
          self.use_bn = getattr(opt, 'use_bn', 0)
  
          self.ss_prob = 0.0 # Schedule sampling probability
  
          embedding_layer = None
          
          if opt.enable_cbs_support:
              glove_vectors = self.initialize_glove(opt)
              embedding_layer = nn.Embedding.from_pretrained(glove_vectors, freeze=True)
              # embedding_layer = nn.Embedding(self.vocab_size + 1, self.input_encoding_size)
          else:
              embedding_layer = nn.Embedding(self.vocab_size + 1, self.input_encoding_size)
          self.embed = nn.Sequential(embedding_layer,
                                  nn.ReLU(),
                                  nn.Dropout(self.drop_prob_lm))
          self.fc_embed = nn.Sequential(nn.Linear(self.fc_feat_size, self.rnn_size),
                                      nn.ReLU(),
                                      nn.Dropout(self.drop_prob_lm))
          self.att_embed = nn.Sequential(*(
                                      ((nn.BatchNorm1d(self.att_feat_size),) if self.use_bn else ())+
                                      (nn.Linear(self.att_feat_size, self.rnn_size),
                                      nn.ReLU(),
                                      nn.Dropout(self.drop_prob_lm))+
                                      ((nn.BatchNorm1d(self.rnn_size),) if self.use_bn==2 else ())))
  
          self.logit_layers = getattr(opt, 'logit_layers', 1)
          if self.logit_layers == 1:
              if opt.enable_cbs_support:
                  temp_linear_layer = nn.Linear(self.input_encoding_size, self.vocab_size + 1, bias = False)
                  temp_linear_layer.weight = embedding_layer.weight
                  self.logit = nn.Sequential(
                      nn.Linear(self.rnn_size, self.input_encoding_size),
                      nn.Tanh(), 
                      temp_linear_layer)
              else:
                  self.logit = nn.Linear(self.rnn_size, self.vocab_size + 1)
          else:
              assert False
              self.logit = [[nn.Linear(self.rnn_size, self.rnn_size), nn.ReLU(), nn.Dropout(0.5)] for _ in range(opt.logit_layers - 1)]
              self.logit = nn.Sequential(*(reduce(lambda x,y:x+y, self.logit) + [nn.Linear(self.rnn_size, self.vocab_size + 1)]))
          self.ctx2att = nn.Linear(self.rnn_size, self.att_hid_size)
  
          # For remove bad endding
          self.vocab = opt.vocab
          self.bad_endings_ix = [int(k) for k,v in self.vocab.items() if v in bad_endings]
  
      def initialize_glove(self, opt):
          # Taken from https://github.com/nocaps-org/updown-baseline/blob/266081042bc2f0c3e6676ee0b5e204036cb6a74a/updown/models/updown_captioner.py#L148
          glove = GloVe(name="42B", dim=300)
          glove_vectors = torch.zeros(self.vocab_size + 1, 300)
          for index, word in opt.vocab.items():
              if word in glove.stoi:
                  glove_vectors[int(index)] = glove.vectors[glove.stoi[word]]
              else:
                  # elif word != 'UNK':
                  # Nocaps paper - pad,unk - 0; Boundary - random
                  # We give 0 to boundary
                  # Random to pad and unk
                  glove_vectors[int(index)] = 2 * torch.randn(300) - 1
          return glove_vectors
  
      def init_hidden(self, bsz):
          weight = next(self.parameters())
          return (weight.new_zeros(self.num_layers, bsz, self.rnn_size),
                  weight.new_zeros(self.num_layers, bsz, self.rnn_size))
  
      def clip_att(self, att_feats, att_masks):
          # Clip the length of att_masks and att_feats to the maximum length
          if att_masks is not None:
              max_len = att_masks.data.long().sum(1).max()
              att_feats = att_feats[:, :max_len].contiguous()
              att_masks = att_masks[:, :max_len].contiguous()
          return att_feats, att_masks
  
      def _prepare_feature(self, fc_feats, att_feats, att_masks):
          att_feats, att_masks = self.clip_att(att_feats, att_masks)
  
          # embed fc and att feats
          fc_feats = self.fc_embed(fc_feats)
          att_feats = pack_wrapper(self.att_embed, att_feats, att_masks)
          
          # Project the attention feats first to reduce memory and computation comsumptions.
          p_att_feats = self.ctx2att(att_feats)
          
          return fc_feats, att_feats, p_att_feats, att_masks
  
      def _forward(self, fc_feats, att_feats, seq, att_masks=None):
          batch_size = fc_feats.size(0)
          state = self.init_hidden(batch_size)
  
          outputs = fc_feats.new_zeros(batch_size, seq.size(1) - 1, self.vocab_size+1)
  
          # Prepare the features
          p_fc_feats, p_att_feats, pp_att_feats, p_att_masks = self._prepare_feature(fc_feats, att_feats, att_masks)
          # pp_att_feats is used for attention, we cache it in advance to reduce computation cost
  
          for i in range(seq.size(1) - 1):
              if self.training and i >= 1 and self.ss_prob > 0.0: # otherwiste no need to sample
                  sample_prob = fc_feats.new(batch_size).uniform_(0, 1)
                  sample_mask = sample_prob < self.ss_prob
                  if sample_mask.sum() == 0:
                      it = seq[:, i].clone()
                  else:
                      sample_ind = sample_mask.nonzero().view(-1)
                      it = seq[:, i].data.clone()
                      #prob_prev = torch.exp(outputs[-1].data.index_select(0, sample_ind)) # fetch prev distribution: shape Nx(M+1)
                      #it.index_copy_(0, sample_ind, torch.multinomial(prob_prev, 1).view(-1))
                      # prob_prev = torch.exp(outputs[-1].data) # fetch prev distribution: shape Nx(M+1)
                      prob_prev = torch.exp(outputs[:, i-1].detach()) # fetch prev distribution: shape Nx(M+1)
                      it.index_copy_(0, sample_ind, torch.multinomial(prob_prev, 1).view(-1).index_select(0, sample_ind))
              else:
                  it = seq[:, i].clone()          
              # break if all the sequences end
              if i >= 1 and seq[:, i].sum() == 0:
                  break
  
              output, state = self.get_logprobs_state(it, p_fc_feats, p_att_feats, pp_att_feats, p_att_masks, state)
              outputs[:, i] = output
  
          return outputs
  
      def get_logprobs_state(self, it, fc_feats, att_feats, p_att_feats, att_masks, state, output_logsoftmax=1):
  
          xt = self.embed(it)
  
          output, state = self.core(xt, fc_feats, att_feats, p_att_feats, state, att_masks)
          if output_logsoftmax:
              logprobs = F.log_softmax(self.logit(output), dim=1)
          else:
              logprobs = self.logit(output)
  
          return logprobs, state
  
      def _sample_beam(self, fc_feats, att_feats, att_masks=None, opt={}):
          beam_size = opt.get('beam_size', 10)
          batch_size = fc_feats.size(0)
  
          p_fc_feats, p_att_feats, pp_att_feats, p_att_masks = self._prepare_feature(fc_feats, att_feats, att_masks)
  
          assert beam_size <= self.vocab_size + 1, 'lets assume this for now, otherwise this corner case causes a few headaches down the road. can be dealt with in future if needed'
          seq = torch.LongTensor(self.seq_length, batch_size).zero_()
          seqLogprobs = torch.FloatTensor(self.seq_length, batch_size, self.vocab_size + 1)
          # lets process every image independently for now, for simplicity
  
          self.done_beams = [[] for _ in range(batch_size)]
          for k in range(batch_size):
              state = self.init_hidden(beam_size)
              tmp_fc_feats = p_fc_feats[k:k+1].expand(beam_size, p_fc_feats.size(1))
              tmp_att_feats = p_att_feats[k:k+1].expand(*((beam_size,)+p_att_feats.size()[1:])).contiguous()
              tmp_p_att_feats = pp_att_feats[k:k+1].expand(*((beam_size,)+pp_att_feats.size()[1:])).contiguous()
              tmp_att_masks = p_att_masks[k:k+1].expand(*((beam_size,)+p_att_masks.size()[1:])).contiguous() if att_masks is not None else None
  
              for t in range(1):
                  if t == 0: # input <bos>
                      it = fc_feats.new_zeros([beam_size], dtype=torch.long)
  
                  logprobs, state = self.get_logprobs_state(it, tmp_fc_feats, tmp_att_feats, tmp_p_att_feats, tmp_att_masks, state)
  
              self.done_beams[k] = self.beam_search(state, logprobs, tmp_fc_feats, tmp_att_feats, tmp_p_att_feats, tmp_att_masks, opt=opt)
              seq[:, k] = self.done_beams[k][0]['seq'] # the first beam has highest cumulative score
              seqLogprobs[:, k] = self.done_beams[k][0]['logps']
          # return the samples and their log likelihoods
          return seq.transpose(0, 1), seqLogprobs.transpose(0, 1)
  
      def get_seq_logprobs(self, fc_feats, att_feats, att_masks, seq):
          batch_size = fc_feats.size(0)
          state = self.init_hidden(batch_size)
  
          p_fc_feats, p_att_feats, pp_att_feats, p_att_masks = self._prepare_feature(fc_feats, att_feats, att_masks)
  
          seqLogprobs = fc_feats.new_zeros(batch_size, self.seq_length, self.vocab_size + 1)
          
          for t in range(self.seq_length + 1):
              if t == 0: # input <bos>
                  it = fc_feats.new_zeros(batch_size, dtype=torch.long)
  
              logprobs, state = self.get_logprobs_state(it, p_fc_feats, p_att_feats, pp_att_feats, p_att_masks, state)
  
              # sample the next word
              if t == self.seq_length: # skip if we achieve maximum length
                  break
              
              # it = torch.distributions.Categorical(logits=logprobs.detach()).sample()
              it = seq[:,t]
  
              # stop when all finished
              if t == 0:
                  unfinished = it > 0
              else:
                  unfinished = unfinished * (it > 0)
              it = it * unfinished.type_as(it)
              
              # it = seq[:,t]
              seqLogprobs[:,t] = logprobs
              
              # quit loop if all sequences have finished
              if unfinished.sum() == 0:
                  break
  
          return seqLogprobs
  
      def _sample(self, fc_feats, att_feats, att_masks=None, opt={}):
          sample_max = opt.get('sample_max', 1)
          beam_size = opt.get('beam_size', 1)
          temperature = opt.get('temperature', 1.0)
          sample_n = int(opt.get('sample_n', 1))
          output_logsoftmax = opt.get('output_logsoftmax', 1)
          decoding_constraint = opt.get('decoding_constraint', 0)
          block_trigrams = opt.get('block_trigrams', 0)
          remove_bad_endings = opt.get('remove_bad_endings', 0)
          if beam_size > 1:
              return self._sample_beam(fc_feats, att_feats, att_masks, opt)
  
          batch_size = fc_feats.size(0)
          state = self.init_hidden(batch_size*sample_n)
  
          p_fc_feats, p_att_feats, pp_att_feats, p_att_masks = self._prepare_feature(fc_feats, att_feats, att_masks)
  
          if sample_n > 1:
              p_fc_feats = p_fc_feats.unsqueeze(1).repeat(1,sample_n,1).reshape((batch_size*sample_n,)+p_fc_feats.shape[1:])
              p_att_feats = p_att_feats.unsqueeze(1).repeat(1,sample_n,1,1).reshape((batch_size*sample_n,)+p_att_feats.shape[1:])
              pp_att_feats = pp_att_feats.unsqueeze(1).repeat(1,sample_n,1,1).reshape((batch_size*sample_n,)+pp_att_feats.shape[1:])
              p_att_masks = p_att_masks.unsqueeze(1).repeat(1,sample_n,1,1).reshape((batch_size*sample_n,)+p_att_masks.shape[1:]) if att_masks is not None else None
  
          trigrams = [] # will be a list of batch_size dictionaries
          
          seq = fc_feats.new_zeros((batch_size*sample_n, self.seq_length), dtype=torch.long)
          seqLogprobs = fc_feats.new_zeros(batch_size*sample_n, self.seq_length, self.vocab_size + 1)
          for t in range(self.seq_length + 1):
              if t == 0: # input <bos>
                  it = fc_feats.new_zeros(batch_size*sample_n, dtype=torch.long)
  
              logprobs, state = self.get_logprobs_state(it, p_fc_feats, p_att_feats, pp_att_feats, p_att_masks, state, output_logsoftmax=output_logsoftmax)
              
              if decoding_constraint and t > 0:
                  tmp = logprobs.new_zeros(logprobs.size())
                  tmp.scatter_(1, seq[:,t-1].data.unsqueeze(1), float('-inf'))
                  logprobs = logprobs + tmp
  
              if remove_bad_endings and t > 0:
                  tmp = logprobs.new_zeros(logprobs.size())
                  prev_bad = np.isin(seq[:,t-1].data.cpu().numpy(), self.bad_endings_ix)
                  # Impossible to generate remove_bad_endings
                  tmp[torch.from_numpy(prev_bad.astype('uint8')), 0] = float('-inf')
                  logprobs = logprobs + tmp
  
              # Mess with trigrams
              if block_trigrams and t >= 3:
                  # Store trigram generated at last step
                  prev_two_batch = seq[:,t-3:t-1]
                  for i in range(batch_size): # = seq.size(0)
                      prev_two = (prev_two_batch[i][0].item(), prev_two_batch[i][1].item())
                      current  = seq[i][t-1]
                      if t == 3: # initialize
                          trigrams.append({prev_two: [current]}) # {LongTensor: list containing 1 int}
                      elif t > 3:
                          if prev_two in trigrams[i]: # add to list
                              trigrams[i][prev_two].append(current)
                          else: # create list
                              trigrams[i][prev_two] = [current]
                  # Block used trigrams at next step
                  prev_two_batch = seq[:,t-2:t]
                  mask = torch.zeros(logprobs.size(), requires_grad=False).cuda() # batch_size x vocab_size
                  for i in range(batch_size):
                      prev_two = (prev_two_batch[i][0].item(), prev_two_batch[i][1].item())
                      if prev_two in trigrams[i]:
                          for j in trigrams[i][prev_two]:
                              mask[i,j] += 1
                  # Apply mask to log probs
                  #logprobs = logprobs - (mask * 1e9)
                  alpha = 2.0 # = 4
                  logprobs = logprobs + (mask * -0.693 * alpha) # ln(1/2) * alpha (alpha -> infty works best)
  
              # sample the next word
              if t == self.seq_length: # skip if we achieve maximum length
                  break
              if sample_max == 1:
                  sampleLogprobs, it = torch.max(logprobs.data, 1)
                  it = it.view(-1).long()
              elif sample_max == 2: # gumbel softmax
                  # ref: https://gist.github.com/yzh119/fd2146d2aeb329d067568a493b20172f
                  def sample_gumbel(shape, eps=1e-20):
                      U = torch.rand(shape).cuda()
                      return -torch.log(-torch.log(U + eps) + eps)
                  def gumbel_softmax_sample(logits, temperature):
                      y = logits + sample_gumbel(logits.size())
                      return F.log_softmax(y / temperature, dim=-1)
                  _logprobs = gumbel_softmax_sample(logprobs, temperature)
                  _, it = torch.max(_logprobs.data, 1)
                  sampleLogprobs = logprobs.gather(1, it.unsqueeze(1)) # gather the logprobs at sampled positions
              else:
                  if sample_max < 0:
                      # negative means topk sampling.
                      tmp = torch.empty_like(logprobs).fill_(float('-inf'))
                      topk, indices = torch.topk(logprobs, -sample_max, dim=1)
                      tmp = tmp.scatter(1, indices, topk)
                      logprobs = tmp
                  logprobs = logprobs / temperature
                  it = torch.distributions.Categorical(logits=logprobs.detach()).sample()
                  # sampleLogprobs = logprobs.gather(1, it.unsqueeze(1)) # gather the logprobs at sampled positions
  
              # stop when all finished
              if t == 0:
                  unfinished = it > 0
              else:
                  unfinished = unfinished * (it > 0)
              it = it * unfinished.type_as(it)
              seq[:,t] = it
              seqLogprobs[:,t] = logprobs
              # quit loop if all sequences have finished
              if unfinished.sum() == 0:
                  break
  
          return seq, seqLogprobs
  
  class Attention(nn.Module):
      def __init__(self, opt):
          super(Attention, self).__init__()
          self.rnn_size = opt.rnn_size
          self.att_hid_size = opt.att_hid_size
  
          self.h2att = nn.Linear(self.rnn_size, self.att_hid_size)
          self.alpha_net = nn.Linear(self.att_hid_size, 1)
  
      def forward(self, h, att_feats, p_att_feats, att_masks=None):
          # The p_att_feats here is already projected
          att_size = att_feats.numel() // att_feats.size(0) // att_feats.size(-1)
          att = p_att_feats.view(-1, att_size, self.att_hid_size)
          
          att_h = self.h2att(h)                        # batch * att_hid_size
          att_h = att_h.unsqueeze(1).expand_as(att)            # batch * att_size * att_hid_size
          dot = att + att_h                                   # batch * att_size * att_hid_size
          dot = torch.tanh(dot)                                # batch * att_size * att_hid_size
          dot = dot.view(-1, self.att_hid_size)               # (batch * att_size) * att_hid_size
          dot = self.alpha_net(dot)                           # (batch * att_size) * 1
          dot = dot.view(-1, att_size)                        # batch * att_size
          
          weight = F.softmax(dot, dim=1)                             # batch * att_size
          if att_masks is not None:
              weight = weight * att_masks.view(-1, att_size).float()
              weight = weight / weight.sum(1, keepdim=True) # normalize to 1
          att_feats_ = att_feats.view(-1, att_size, att_feats.size(-1)) # batch * att_size * att_feat_size
          att_res = torch.bmm(weight.unsqueeze(1), att_feats_).squeeze(1) # batch * att_feat_size
  
          return att_res
  
  class Att2in2Core(nn.Module):
      def __init__(self, opt):
          super(Att2in2Core, self).__init__()
          self.input_encoding_size = opt.input_encoding_size
          self.rnn_size = opt.rnn_size
          self.drop_prob_lm = opt.drop_prob_lm
          self.fc_feat_size = opt.fc_feat_size
          self.att_feat_size = opt.att_feat_size
          self.att_hid_size = opt.att_hid_size
          
          # Build a LSTM
          self.a2c = nn.Linear(self.rnn_size, 2 * self.rnn_size)
          self.i2h = nn.Linear(self.input_encoding_size, 5 * self.rnn_size)
          self.h2h = nn.Linear(self.rnn_size, 5 * self.rnn_size)
          self.dropout = nn.Dropout(self.drop_prob_lm)
  
          self.attention = Attention(opt)
  
      def forward(self, xt, fc_feats, att_feats, p_att_feats, state, att_masks=None):
          att_res = self.attention(state[0][-1], att_feats, p_att_feats, att_masks)
  
          all_input_sums = self.i2h(xt) + self.h2h(state[0][-1])
          sigmoid_chunk = all_input_sums.narrow(1, 0, 3 * self.rnn_size)
          sigmoid_chunk = torch.sigmoid(sigmoid_chunk)
          in_gate = sigmoid_chunk.narrow(1, 0, self.rnn_size)
          forget_gate = sigmoid_chunk.narrow(1, self.rnn_size, self.rnn_size)
          out_gate = sigmoid_chunk.narrow(1, self.rnn_size * 2, self.rnn_size)
  
          in_transform = all_input_sums.narrow(1, 3 * self.rnn_size, 2 * self.rnn_size) + \
              self.a2c(att_res)
          in_transform = torch.max(\
              in_transform.narrow(1, 0, self.rnn_size),
              in_transform.narrow(1, self.rnn_size, self.rnn_size))
          next_c = forget_gate * state[1][-1] + in_gate * in_transform
          next_h = out_gate * torch.tanh(next_c)
  
          output = self.dropout(next_h)
          state = (next_h.unsqueeze(0), next_c.unsqueeze(0))
          return output, state
  
  class Att2inCore(Att2in2Core):
      def __init__(self, opt):
          super(Att2inCore, self).__init__(opt)
          del self.a2c
          self.a2c = nn.Linear(self.att_feat_size, 2 * self.rnn_size)
  
  class Att2in2Model(AttModel):
      def __init__(self, opt):
          super(Att2in2Model, self).__init__(opt)
          self.core = Att2in2Core(opt)
          delattr(self, 'fc_embed')
          self.fc_embed = lambda x : x