text: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
  import random
  from params import Params
  from utils import Vocab, Hypothesis, word_detector
  from typing import Union, List
  
  DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  eps = 1e-31
  
  
  class EncoderRNN(nn.Module):
  
    def __init__(self, embed_size, hidden_size, bidi=True, *, rnn_drop: float=0):
      super(EncoderRNN, self).__init__()
      self.hidden_size = hidden_size
      self.num_directions = 2 if bidi else 1
      self.gru = nn.GRU(embed_size, hidden_size, bidirectional=bidi, dropout=rnn_drop)
  
    def forward(self, embedded, hidden, input_lengths=None):
      """
      :param embedded: (src seq len, batch size, embed size)
      :param hidden: (num directions, batch size, encoder hidden size)
      :param input_lengths: list containing the non-padded length of each sequence in this batch;
                            if set, we use `PackedSequence` to skip the PAD inputs and leave the
                            corresponding encoder states as zeros
      :return: (src seq len, batch size, hidden size * num directions = decoder hidden size)
  
      Perform multi-step encoding.
      """
      if input_lengths is not None:
        embedded = pack_padded_sequence(embedded, input_lengths)
  
      output, hidden = self.gru(embedded, hidden)
  
      if input_lengths is not None:
        output, _ = pad_packed_sequence(output)
  
      if self.num_directions > 1:
        # hidden: (num directions, batch, hidden) => (1, batch, hidden * 2)
        batch_size = hidden.size(1)
        hidden = hidden.transpose(0, 1).contiguous().view(1, batch_size,
                                                          self.hidden_size * self.num_directions)
      return output, hidden
  
    def init_hidden(self, batch_size):
      return torch.zeros(self.num_directions, batch_size, self.hidden_size, device=DEVICE)
  
  
  class DecoderRNN(nn.Module):
  
    def __init__(self, vocab_size, embed_size, hidden_size, *, enc_attn=True, dec_attn=True,
                 enc_attn_cover=True, pointer=True, tied_embedding=None, out_embed_size=None,
                 in_drop: float=0, rnn_drop: float=0, out_drop: float=0, enc_hidden_size=None):
      super(DecoderRNN, self).__init__()
      self.vocab_size = vocab_size
      self.hidden_size = hidden_size
      self.combined_size = self.hidden_size
      self.enc_attn = enc_attn
      self.dec_attn = dec_attn
      self.enc_attn_cover = enc_attn_cover
      self.pointer = pointer
      self.out_embed_size = out_embed_size
      if tied_embedding is not None and self.out_embed_size and embed_size != self.out_embed_size:
        print("Warning: Output embedding size %d is overriden by its tied embedding size %d."
              % (self.out_embed_size, embed_size))
        self.out_embed_size = embed_size
  
      self.in_drop = nn.Dropout(in_drop) if in_drop > 0 else None
      self.gru = nn.GRU(embed_size, self.hidden_size, dropout=rnn_drop)
  
      if enc_attn:
        if not enc_hidden_size: enc_hidden_size = self.hidden_size
        self.enc_bilinear = nn.Bilinear(self.hidden_size, enc_hidden_size, 1)
        self.combined_size += enc_hidden_size
        if enc_attn_cover:
          self.cover_weight = nn.Parameter(torch.rand(1))
  
      if dec_attn:
        self.dec_bilinear = nn.Bilinear(self.hidden_size, self.hidden_size, 1)
        self.combined_size += self.hidden_size
  
      self.out_drop = nn.Dropout(out_drop) if out_drop > 0 else None
      if pointer:
        self.ptr = nn.Linear(self.combined_size, 1)
  
      if tied_embedding is not None and embed_size != self.combined_size:
        # use pre_out layer if combined size is different from embedding size
        self.out_embed_size = embed_size
  
      if self.out_embed_size:  # use pre_out layer
        self.pre_out = nn.Linear(self.combined_size, self.out_embed_size)
        size_before_output = self.out_embed_size
      else:  # don't use pre_out layer
        size_before_output = self.combined_size
  
      self.out = nn.Linear(size_before_output, vocab_size)
      if tied_embedding is not None:
        self.out.weight = tied_embedding.weight
  
    def forward(self, embedded, hidden, encoder_states=None, decoder_states=None, coverage_vector=None, *,
                encoder_word_idx=None, ext_vocab_size: int=None, log_prob: bool=True):
      """
      :param embedded: (batch size, embed size)
      :param hidden: (1, batch size, decoder hidden size)
      :param encoder_states: (src seq len, batch size, hidden size), for attention mechanism
      :param decoder_states: (past dec steps, batch size, hidden size), for attention mechanism
      :param encoder_word_idx: (src seq len, batch size), for pointer network
      :param ext_vocab_size: the dynamic vocab size, determined by the max num of OOV words contained
                             in any src seq in this batch, for pointer network
      :param log_prob: return log probability instead of probability
      :return: tuple of four things:
               1. word prob or log word prob, (batch size, dynamic vocab size);
               2. RNN hidden state after this step, (1, batch size, decoder hidden size);
               3. attention weights over encoder states, (batch size, src seq len);
               4. prob of copying by pointing as opposed to generating, (batch size, 1)
  
      Perform single-step decoding.
      """
      batch_size = embedded.size(0)
      combined = torch.zeros(batch_size, self.combined_size, device=DEVICE)
  
      if self.in_drop: embedded = self.in_drop(embedded)
  
      output, hidden = self.gru(embedded.unsqueeze(0), hidden)  # unsqueeze and squeeze are necessary
      combined[:, :self.hidden_size] = output.squeeze(0)        # as RNN expects a 3D tensor (step=1)
      offset = self.hidden_size
      enc_attn, prob_ptr = None, None  # for visualization
  
      if self.enc_attn or self.pointer:
        # energy and attention: (num encoder states, batch size, 1)
        num_enc_steps = encoder_states.size(0)
        enc_total_size = encoder_states.size(2)
        enc_energy = self.enc_bilinear(hidden.expand(num_enc_steps, batch_size, -1).contiguous(),
                                       encoder_states)
        if self.enc_attn_cover and coverage_vector is not None:
          enc_energy += self.cover_weight * torch.log(coverage_vector.transpose(0, 1).unsqueeze(2) + eps)
        # transpose => (batch size, num encoder states, 1)
        enc_attn = F.softmax(enc_energy, dim=0).transpose(0, 1)
        if self.enc_attn:
          # context: (batch size, encoder hidden size, 1)
          enc_context = torch.bmm(encoder_states.permute(1, 2, 0), enc_attn)
          combined[:, offset:offset+enc_total_size] = enc_context.squeeze(2)
          offset += enc_total_size
        enc_attn = enc_attn.squeeze(2)
  
      if self.dec_attn:
        if decoder_states is not None and len(decoder_states) > 0:
          dec_energy = self.dec_bilinear(hidden.expand_as(decoder_states).contiguous(),
                                         decoder_states)
          dec_attn = F.softmax(dec_energy, dim=0).transpose(0, 1)
          dec_context = torch.bmm(decoder_states.permute(1, 2, 0), dec_attn)
          combined[:, offset:offset + self.hidden_size] = dec_context.squeeze(2)
        offset += self.hidden_size
  
      if self.out_drop: combined = self.out_drop(combined)
  
      # generator
      if self.out_embed_size:
        out_embed = self.pre_out(combined)
      else:
        out_embed = combined
      logits = self.out(out_embed)  # (batch size, vocab size)
  
      # pointer
      if self.pointer:
        output = torch.zeros(batch_size, ext_vocab_size, device=DEVICE)
        # distribute probabilities between generator and pointer
        prob_ptr = F.sigmoid(self.ptr(combined))  # (batch size, 1)
        prob_gen = 1 - prob_ptr
        # add generator probabilities to output
        gen_output = F.softmax(logits, dim=1)  # can't use log_softmax due to adding probabilities
        output[:, :self.vocab_size] = prob_gen * gen_output
        # add pointer probabilities to output
        ptr_output = enc_attn
        output.scatter_add_(1, encoder_word_idx.transpose(0, 1), prob_ptr * ptr_output)
        if log_prob: output = torch.log(output + eps)
      else:
        if log_prob: output = F.log_softmax(logits, dim=1)
        else: output = F.softmax(logits, dim=1)
  
      return output, hidden, enc_attn, prob_ptr
  
  
  class Seq2SeqOutput(object):
  
    def __init__(self, encoder_outputs: torch.Tensor, encoder_hidden: torch.Tensor,
                 decoded_tokens: torch.Tensor, loss: Union[torch.Tensor, float]=0,
                 loss_value: float=0, enc_attn_weights: torch.Tensor=None,
                 ptr_probs: torch.Tensor=None):
      self.encoder_outputs = encoder_outputs
      self.encoder_hidden = encoder_hidden
      self.decoded_tokens = decoded_tokens  # (out seq len, batch size)
      self.loss = loss  # scalar
      self.loss_value = loss_value  # float value, excluding coverage loss
      self.enc_attn_weights = enc_attn_weights  # (out seq len, batch size, src seq len)
      self.ptr_probs = ptr_probs  # (out seq len, batch size)
  
  
  class Seq2Seq(nn.Module):
  
    def __init__(self, vocab: Vocab, params: Params, max_dec_steps=None):
      """
      :param vocab: mainly for info about special tokens and vocab size
      :param params: model hyper-parameters
      :param max_dec_steps: max num of decoding steps (only effective at test time, as during
                            training the num of steps is determined by the `target_tensor`); it is
                            safe to change `self.max_dec_steps` as the network architecture is
                            independent of src/tgt seq lengths
  
      Create the seq2seq model; its encoder and decoder will be created automatically.
      """
      super(Seq2Seq, self).__init__()
      self.vocab = vocab
      self.vocab_size = len(vocab)
      if vocab.embeddings is not None:
        self.embed_size = vocab.embeddings.shape[1]
        if params.embed_size is not None and self.embed_size != params.embed_size:
          print("Warning: Model embedding size %d is overriden by pre-trained embedding size %d."
                % (params.embed_size, self.embed_size))
        embedding_weights = torch.from_numpy(vocab.embeddings)
      else:
        self.embed_size = params.embed_size
        embedding_weights = None
      self.max_dec_steps = params.max_tgt_len + 1 if max_dec_steps is None else max_dec_steps
      self.enc_attn = params.enc_attn
      self.enc_attn_cover = params.enc_attn_cover
      self.dec_attn = params.dec_attn
      self.pointer = params.pointer
      self.cover_loss = params.cover_loss
      self.cover_func = params.cover_func
      enc_total_size = params.hidden_size * 2 if params.enc_bidi else params.hidden_size
      if params.dec_hidden_size:
        dec_hidden_size = params.dec_hidden_size
        self.enc_dec_adapter = nn.Linear(enc_total_size, dec_hidden_size)
      else:
        dec_hidden_size = enc_total_size
        self.enc_dec_adapter = None
  
      self.embedding = nn.Embedding(self.vocab_size, self.embed_size, padding_idx=vocab.PAD,
                                    _weight=embedding_weights)
      self.encoder = EncoderRNN(self.embed_size, params.hidden_size, params.enc_bidi,
                                rnn_drop=params.enc_rnn_dropout)
      self.decoder = DecoderRNN(self.vocab_size, self.embed_size, dec_hidden_size,
                                enc_attn=params.enc_attn, dec_attn=params.dec_attn,
                                pointer=params.pointer, out_embed_size=params.out_embed_size,
                                tied_embedding=self.embedding if params.tie_embed else None,
                                in_drop=params.dec_in_dropout, rnn_drop=params.dec_rnn_dropout,
                                out_drop=params.dec_out_dropout, enc_hidden_size=enc_total_size)
  
    def filter_oov(self, tensor, ext_vocab_size):
      """Replace any OOV index in `tensor` with UNK"""
      if ext_vocab_size and ext_vocab_size > self.vocab_size:
        result = tensor.clone()
        result[tensor >= self.vocab_size] = self.vocab.UNK
        return result
      return tensor
  
    def get_coverage_vector(self, enc_attn_weights):
      """Combine the past attention weights into one vector"""
      if self.cover_func == 'max':
        coverage_vector, _ = torch.max(torch.cat(enc_attn_weights), dim=0)
      elif self.cover_func == 'sum':
        coverage_vector = torch.sum(torch.cat(enc_attn_weights), dim=0)
      else:
        raise ValueError('Unrecognized cover_func: ' + self.cover_func)
      return coverage_vector
  
    def forward(self, input_tensor, target_tensor=None, input_lengths=None, criterion=None, *,
                forcing_ratio=0, partial_forcing=True, ext_vocab_size=None, sample=False,
                saved_out: Seq2SeqOutput=None, visualize: bool=None, include_cover_loss: bool=False)\
            -> Seq2SeqOutput:
      """
      :param input_tensor: tensor of word indices, (src seq len, batch size)
      :param target_tensor: tensor of word indices, (tgt seq len, batch size)
      :param input_lengths: see explanation in `EncoderRNN`
      :param criterion: the loss function; if set, loss will be returned
      :param forcing_ratio: see explanation in `Params` (requires `target_tensor`, training only)
      :param partial_forcing: see explanation in `Params` (training only)
      :param ext_vocab_size: see explanation in `DecoderRNN`
      :param sample: if True, the returned `decoded_tokens` will be based on random sampling instead
                     of greedily selecting the token of the highest probability at each step
      :param saved_out: the output of this function in a previous run; if set, the encoding step will
                        be skipped and we reuse the encoder states saved in this object
      :param visualize: whether to return data for attention and pointer visualization; if None,
                        return if no `criterion` is provided
      :param include_cover_loss: whether to include coverage loss in the returned `loss_value`
  
      Run the seq2seq model for training or testing.
      """
      input_length = input_tensor.size(0)
      batch_size = input_tensor.size(1)
      log_prob = not (sample or self.decoder.pointer)  # don't apply log too soon in these cases
      if visualize is None:
        visualize = criterion is None
      if visualize and not (self.enc_attn or self.pointer):
        visualize = False  # nothing to visualize
  
      if target_tensor is None:
        target_length = self.max_dec_steps
      else:
        target_length = target_tensor.size(0)
  
      if forcing_ratio == 1:
        # if fully teacher-forced, it may be possible to eliminate the for-loop over decoder steps
        # for generality, this optimization is not investigated
        use_teacher_forcing = True
      elif forcing_ratio > 0:
        if partial_forcing:
          use_teacher_forcing = None  # decide later individually in each step
        else:
          use_teacher_forcing = random.random() < forcing_ratio
      else:
        use_teacher_forcing = False
  
      if saved_out:  # reuse encoder states of a previous run
        encoder_outputs = saved_out.encoder_outputs
        encoder_hidden = saved_out.encoder_hidden
        assert input_length == encoder_outputs.size(0)
        assert batch_size == encoder_outputs.size(1)
      else:  # run the encoder
        encoder_hidden = self.encoder.init_hidden(batch_size)
        # encoder_embedded: (input len, batch size, embed size)
        encoder_embedded = self.embedding(self.filter_oov(input_tensor, ext_vocab_size))
        encoder_outputs, encoder_hidden = \
          self.encoder(encoder_embedded, encoder_hidden, input_lengths)
  
      # initialize return values
      r = Seq2SeqOutput(encoder_outputs, encoder_hidden,
                        torch.zeros(target_length, batch_size, dtype=torch.long))
      if visualize:
        r.enc_attn_weights = torch.zeros(target_length, batch_size, input_length)
        if self.pointer:
          r.ptr_probs = torch.zeros(target_length, batch_size)
  
      decoder_input = torch.tensor([self.vocab.SOS] * batch_size, device=DEVICE)
      if self.enc_dec_adapter is None:
        decoder_hidden = encoder_hidden
      else:
        decoder_hidden = self.enc_dec_adapter(encoder_hidden)
      decoder_states = []
      enc_attn_weights = []
  
      for di in range(target_length):
        decoder_embedded = self.embedding(self.filter_oov(decoder_input, ext_vocab_size))
        if enc_attn_weights:
          coverage_vector = self.get_coverage_vector(enc_attn_weights)
        else:
          coverage_vector = None
        decoder_output, decoder_hidden, dec_enc_attn, dec_prob_ptr = \
          self.decoder(decoder_embedded, decoder_hidden, encoder_outputs,
                       torch.cat(decoder_states) if decoder_states else None, coverage_vector,
                       encoder_word_idx=input_tensor, ext_vocab_size=ext_vocab_size,
                       log_prob=log_prob)
        if self.dec_attn:
          decoder_states.append(decoder_hidden)
        # save the decoded tokens
        if not sample:
          _, top_idx = decoder_output.data.topk(1)  # top_idx shape: (batch size, k=1)
        else:
          prob_distribution = torch.exp(decoder_output) if log_prob else decoder_output
          top_idx = torch.multinomial(prob_distribution, 1)
        top_idx = top_idx.squeeze(1).detach()  # detach from history as input
        r.decoded_tokens[di] = top_idx
        # compute loss
        if criterion:
          if target_tensor is None:
            gold_standard = top_idx  # for sampling
          else:
            gold_standard = target_tensor[di]
          if not log_prob:
            decoder_output = torch.log(decoder_output + eps)  # necessary for NLLLoss
          nll_loss = criterion(decoder_output, gold_standard)
          r.loss += nll_loss
          r.loss_value += nll_loss.item()
        # update attention history and compute coverage loss
        if self.enc_attn_cover or (criterion and self.cover_loss > 0):
          if coverage_vector is not None and criterion and self.cover_loss > 0:
            coverage_loss = torch.sum(torch.min(coverage_vector, dec_enc_attn)) / batch_size \
                            * self.cover_loss
            r.loss += coverage_loss
            if include_cover_loss: r.loss_value += coverage_loss.item()
          enc_attn_weights.append(dec_enc_attn.unsqueeze(0))
        # save data for visualization
        if visualize:
          r.enc_attn_weights[di] = dec_enc_attn.data
          if self.pointer:
            r.ptr_probs[di] = dec_prob_ptr.squeeze(1).data
        # decide the next input
        if use_teacher_forcing or (use_teacher_forcing is None and random.random() < forcing_ratio):
          decoder_input = target_tensor[di]  # teacher forcing
        else:
          decoder_input = top_idx
      
      return r
  
    def beam_search(self, input_tensor, input_lengths=None, ext_vocab_size=None, beam_size=4, *,
                    min_out_len=1, max_out_len=None, len_in_words=True) -> List[Hypothesis]:
      """
      :param input_tensor: tensor of word indices, (src seq len, batch size); for now, batch size has
                           to be 1
      :param input_lengths: see explanation in `EncoderRNN`
      :param ext_vocab_size: see explanation in `DecoderRNN`
      :param beam_size: the beam size
      :param min_out_len: required minimum output length
      :param max_out_len: required maximum output length (if None, use the model's own value)
      :param len_in_words: if True, count output length in words instead of tokens (i.e. do not count
                           punctuations)
      :return: list of the best decoded sequences, in descending order of probability
  
      Use beam search to generate summaries.
      """
      batch_size = input_tensor.size(1)
      assert batch_size == 1
      if max_out_len is None:
        max_out_len = self.max_dec_steps - 1  # max_out_len doesn't count EOS
  
      # encode
      encoder_hidden = self.encoder.init_hidden(batch_size)
      # encoder_embedded: (input len, batch size, embed size)
      encoder_embedded = self.embedding(self.filter_oov(input_tensor, ext_vocab_size))
      encoder_outputs, encoder_hidden = \
        self.encoder(encoder_embedded, encoder_hidden, input_lengths)
      if self.enc_dec_adapter is None:
        decoder_hidden = encoder_hidden
      else:
        decoder_hidden = self.enc_dec_adapter(encoder_hidden)
      # turn batch size from 1 to beam size (by repeating)
      # if we want dynamic batch size, the following must be created for all possible batch sizes
      encoder_outputs = encoder_outputs.expand(-1, beam_size, -1).contiguous()
      input_tensor = input_tensor.expand(-1, beam_size).contiguous()
  
      # decode
      hypos = [Hypothesis([self.vocab.SOS], [], decoder_hidden, [], [], 1)]
      results, backup_results = [], []
      step = 0
      while hypos and step < 2 * max_out_len:  # prevent infinitely generating punctuations
        # make batch size equal to beam size (n_hypos <= beam size)
        n_hypos = len(hypos)
        if n_hypos < beam_size:
          hypos.extend(hypos[-1] for _ in range(beam_size - n_hypos))
        # assemble existing hypotheses into a batch
        decoder_input = torch.tensor([h.tokens[-1] for h in hypos], device=DEVICE)
        decoder_hidden = torch.cat([h.dec_hidden for h in hypos], 1)
        if self.dec_attn and step > 0:  # dim 0 is decoding step, dim 1 is beam batch
          decoder_states = torch.cat([torch.cat(h.dec_states, 0) for h in hypos], 1)
        else:
          decoder_states = None
        if self.enc_attn_cover:
          enc_attn_weights = [torch.cat([h.enc_attn_weights[i] for h in hypos], 1)
                              for i in range(step)]
        else:
          enc_attn_weights = []
        if enc_attn_weights:
          coverage_vector = self.get_coverage_vector(enc_attn_weights)  # shape: (beam size, src len)
        else:
          coverage_vector = None
        # run the decoder over the assembled batch
        decoder_embedded = self.embedding(self.filter_oov(decoder_input, ext_vocab_size))
        decoder_output, decoder_hidden, dec_enc_attn, dec_prob_ptr = \
          self.decoder(decoder_embedded, decoder_hidden, encoder_outputs,
                       decoder_states, coverage_vector,
                       encoder_word_idx=input_tensor, ext_vocab_size=ext_vocab_size)
        top_v, top_i = decoder_output.data.topk(beam_size)  # shape of both: (beam size, beam size)
        # create new hypotheses
        new_hypos = []
        for in_idx in range(n_hypos):
          for out_idx in range(beam_size):
            new_tok = top_i[in_idx][out_idx].item()
            new_prob = top_v[in_idx][out_idx].item()
            if len_in_words:
              non_word = not self.vocab.is_word(new_tok)
            else:
              non_word = new_tok == self.vocab.EOS  # only SOS & EOS don't count
            new_hypo = hypos[in_idx].create_next(new_tok, new_prob,
                                                 decoder_hidden[0][in_idx].unsqueeze(0).unsqueeze(0),
                                                 self.dec_attn,
                                                 dec_enc_attn[in_idx].unsqueeze(0).unsqueeze(0)
                                                 if dec_enc_attn is not None else None, non_word)
            new_hypos.append(new_hypo)
        # process the new hypotheses
        new_hypos = sorted(new_hypos, key=lambda h: -h.avg_log_prob)
        hypos = []
        new_complete_results, new_incomplete_results = [], []
        for nh in new_hypos:
          length = len(nh)
          if nh.tokens[-1] == self.vocab.EOS:  # a complete hypothesis
            if len(new_complete_results) < beam_size and min_out_len <= length <= max_out_len:
              new_complete_results.append(nh)
          elif len(hypos) < beam_size and length < max_out_len:  # an incomplete hypothesis
            hypos.append(nh)
          elif length == max_out_len and len(new_incomplete_results) < beam_size:
            new_incomplete_results.append(nh)
        if new_complete_results:
          results.extend(new_complete_results)
        elif new_incomplete_results:
          backup_results.extend(new_incomplete_results)
        step += 1
      if not results:  # if no sequence ends with EOS within desired length, fallback to sequences
        results = backup_results  # that are "truncated" at the end to max_out_len
      return sorted(results, key=lambda h: -h.avg_log_prob)[:beam_size]
  