text: # Copyright (c) 2019 NVIDIA Corporation
  """
  This package contains Transformer for translation Neural Module
  """
  __all__ = ['TransformerEncoderNM',
             'TransformerDecoderNM',
             'GreedyLanguageGeneratorNM',
             'BeamSearchTranslatorNM']
  
  import math
  
  from nemo.backends.pytorch.nm import TrainableNM, LossNM
  from nemo.core.neural_types import *
  
  from ..transformer import (TransformerEmbedding,
                             TransformerEncoder,
                             TransformerDecoder,
                             GreedySequenceGenerator,
                             BeamSearchSequenceGenerator)
  from ..transformer.utils import transformer_weights_init
  
  
  class TransformerEncoderNM(TrainableNM):
      """
      Neural module which consists of embedding layer followed by Transformer
      encoder.
  
      Args:
          vocab_size: size of the vocabulary (number of tokens)
          d_model: hidden size (d_model) of the Transformer
          max_seq_length: maximum allowed length of input sequences, feeding
              longer sequences will cause an error
          embedding_dropout: dropout ratio applied to embeddings
          learn_positional_encodings: bool, whether to learn positional encoding
              or use fixed sinusoidal encodings
          num_layers: number of layers in Transformer encoder
          mask_future: bool, whether to apply triangular future masking to the
              sequence of hidden states (which allows to use it for LM)
          num_attn_heads: number of attention heads
          d_inner: number of neurons in the intermediate part of feed-forward
              network (FFN)
          ffn_dropout: dropout ratio applied to FFN
          attn_score_dropout: dropout ratio applied to attention scores
          attn_layer_dropout: dropout ratio applied to the output of attn layer
          hidden_act: activation function applied in intermediate FFN module
      """
  
      @property
      def input_ports(self):
          """Returns definitions of module input ports.
  
          input_ids:
              0: AxisType(BatchTag)
  
              1: AxisType(TimeTag)
  
          input_mask_src:
              0: AxisType(BatchTag)
  
              1: AxisType(TimeTag)
          """
          return {
              "input_ids": NeuralType({
                  0: AxisType(BatchTag),
                  1: AxisType(TimeTag)
              }),
              "input_mask_src": NeuralType({
                  0: AxisType(BatchTag),
                  1: AxisType(TimeTag)
              }),
          }
  
      @property
      def output_ports(self):
          """Returns definitions of module output ports.
  
          hidden_states:
              0: AxisType(BatchTag)
  
              1: AxisType(TimeTag)
  
              2: AxisType(ChannelTag)
          """
          return {
              "hidden_states": NeuralType({
                  0: AxisType(BatchTag),
                  1: AxisType(TimeTag),
                  2: AxisType(ChannelTag)
              })
          }
  
      def __init__(self,
                   vocab_size,
                   d_model,
                   d_inner,
                   max_seq_length,
                   num_layers,
                   num_attn_heads,
                   ffn_dropout=0.0,
                   embedding_dropout=0.0,
                   attn_score_dropout=0.0,
                   attn_layer_dropout=0.0,
                   learn_positional_encodings=False,
                   hidden_act='relu',
                   mask_future=False,
                   **kwargs):
          TrainableNM.__init__(self, **kwargs)
  
          self.embedding_layer = TransformerEmbedding(
              vocab_size=vocab_size,
              hidden_size=d_model,
              max_sequence_length=max_seq_length,
              embedding_dropout=embedding_dropout,
              learn_positional_encodings=learn_positional_encodings)
          self.encoder = TransformerEncoder(
              num_layers=num_layers,
              hidden_size=d_model,
              mask_future=mask_future,
              num_attention_heads=num_attn_heads,
              inner_size=d_inner,
              ffn_dropout=ffn_dropout,
              hidden_act=hidden_act,
              attn_score_dropout=attn_score_dropout,
              attn_layer_dropout=attn_layer_dropout)
  
          std_init_range = 1 / math.sqrt(d_model)
          self.apply(lambda module: transformer_weights_init(module,
                                                             std_init_range))
          self.to(self._device)
  
      def forward(self, input_ids, input_mask_src):
          hidden_states = self.embedding_layer(input_ids)
          hidden_states = self.encoder(hidden_states, input_mask_src)
          return hidden_states
  
  
  class TransformerDecoderNM(TrainableNM):
      """
      Neural module which consists of embedding layer followed by Transformer
      decoder.
  
      Args:
          vocab_size: size of the vocabulary (number of tokens)
          d_model: hidden size (d_model) of the Transformer
          max_seq_length: maximum allowed length of input sequences, feeding
              longer sequences will cause an error
          embedding_dropout: dropout ratio applied to embeddings
          learn_positional_encodings: bool, whether to learn positional encoding
              or use fixed sinusoidal encodings
          num_layers: number of layers in Transformer decoder
          num_attn_heads: number of attention heads
          d_inner: number of neurons in the intermediate part of feed-forward
              network (FFN)
          ffn_dropout: dropout ratio applied to FFN
          attn_score_dropout: dropout ratio applied to attention scores
          attn_layer_dropout: dropout ratio applied to the output of attn layer
          hidden_act: activation function applied in intermediate FFN module
      """
  
      @property
      def input_ports(self):
          """Returns definitions of module input ports.
  
          input_ids_tgt:
              0: AxisType(BatchTag)
  
              1: AxisType(TimeTag)
  
          hidden_states_src:
              0: AxisType(BatchTag)
  
              1: AxisType(TimeTag)
  
              2: AxisType(ChannelTag)
  
          input_mask_src:
              0: AxisType(BatchTag)
  
              1: AxisType(TimeTag)
  
          input_mask_tgt:
              0: AxisType(BatchTag)
  
              1: AxisType(TimeTag)
          """
          return {
              "input_ids_tgt": NeuralType({
                  0: AxisType(BatchTag),
                  1: AxisType(TimeTag)
              }),
              "hidden_states_src": NeuralType({
                  0: AxisType(BatchTag),
                  1: AxisType(TimeTag),
                  2: AxisType(ChannelTag)
              }),
              "input_mask_src": NeuralType({
                  0: AxisType(BatchTag),
                  1: AxisType(TimeTag)
              }),
              "input_mask_tgt": NeuralType({
                  0: AxisType(BatchTag),
                  1: AxisType(TimeTag)
              })
          }
  
      @property
      def output_ports(self):
          """Returns definitions of module output ports.
  
          hidden_states:
              0: AxisType(BatchTag)
  
              1: AxisType(TimeTag)
  
              2: AxisType(ChannelTag)
          """
          return {
              "hidden_states": NeuralType({
                  0: AxisType(BatchTag),
                  1: AxisType(TimeTag),
                  2: AxisType(ChannelTag)
              })
          }
  
      def __init__(self,
                   vocab_size,
                   d_model,
                   d_inner,
                   num_layers,
                   max_seq_length,
                   num_attn_heads,
                   ffn_dropout=0.0,
                   embedding_dropout=0.0,
                   attn_score_dropout=0.0,
                   attn_layer_dropout=0.0,
                   learn_positional_encodings=False,
                   hidden_act='relu',
                   **kwargs):
          TrainableNM.__init__(self, **kwargs)
  
          self.embedding_layer = TransformerEmbedding(
              vocab_size=vocab_size,
              hidden_size=d_model,
              max_sequence_length=max_seq_length,
              embedding_dropout=embedding_dropout,
              learn_positional_encodings=learn_positional_encodings
          )
          self.decoder = TransformerDecoder(
              num_layers=num_layers,
              hidden_size=d_model,
              num_attention_heads=num_attn_heads,
              inner_size=d_inner,
              ffn_dropout=ffn_dropout,
              hidden_act=hidden_act,
              attn_score_dropout=attn_score_dropout,
              attn_layer_dropout=attn_layer_dropout
          )
  
          std_init_range = 1 / math.sqrt(d_model)
          self.apply(lambda module: transformer_weights_init(module,
                                                             std_init_range))
          self.to(self._device)
  
      def forward(self,
                  input_ids_tgt,
                  hidden_states_src,
                  input_mask_src,
                  input_mask_tgt):
          hidden_states_tgt = self.embedding_layer(input_ids_tgt)
          hidden_states = self.decoder(hidden_states_tgt,
                                       input_mask_tgt,
                                       hidden_states_src,
                                       input_mask_src)
          return hidden_states
  
  
  class GreedyLanguageGeneratorNM(TrainableNM):
      """
      Neural module for greedy text generation with language model
  
      Args:
          decoder: module which maps input_ids into hidden_states
          log_softmax: module which maps hidden_states into log_probs
          max_seq_length: maximum allowed length of generated sequences
          pad_token: index of padding token in the vocabulary
          bos_token: index of beginning of sequence token in the vocabulary
          eos_token: index of end of sequence token in the vocabulary
          batch_size: size of the batch of generated sequences if no starting
              tokens are provided
      """
  
      @property
      def input_ports(self):
          """Returns definitions of module input ports.
  
          input_ids:
              0: AxisType(BatchTag)
  
              1: AxisType(TimeTag)
          """
          return {
              "input_ids": NeuralType({
                  0: AxisType(BatchTag),
                  1: AxisType(TimeTag)
              })
          }
  
      @property
      def output_ports(self):
          """Returns definitions of module output ports.
  
          output_ids:
              0: AxisType(BatchTag)
  
              1: AxisType(TimeTag)
          """
          return {
              "output_ids": NeuralType({
                  0: AxisType(BatchTag),
                  1: AxisType(TimeTag)
              })
          }
  
      def __init__(self,
                   decoder,
                   log_softmax,
                   max_seq_length,
                   pad_token,
                   bos_token,
                   eos_token,
                   batch_size=1,
                   **kwargs):
          TrainableNM.__init__(self, **kwargs)
  
          self.generator = GreedySequenceGenerator(
              decoder.embedding_layer,
              decoder.decoder,
              log_softmax,
              max_sequence_length=max_seq_length,
              pad=pad_token,
              bos=bos_token,
              eos=eos_token,
              batch_size=batch_size
          )
  
      @property
      def num_weights(self):
          return 0
  
      def forward(self, input_ids):
          output_ids = self.generator(decoder_input_ids=input_ids)
          return output_ids
  
  
  class BeamSearchTranslatorNM(TrainableNM):
      """
      Neural module for beam search translation generation
  
      Args:
          decoder: module which maps input_ids into hidden_states
          log_softmax: module which maps hidden_states into log_probs
          max_seq_length: maximum allowed length of generated sequences
          pad_token: index of padding token in the vocabulary
          bos_token: index of beginning of sequence token in the vocabulary
          eos_token: index of end of sequence token in the vocabulary
          batch_size: size of the batch of generated sequences if no starting
              tokens are provided
          beam_size: size of the beam
          max_delta_length: maximum allowed difference between generated output
              and input sequence in case of conditional decoding
          length_penalty: parameter which penalizes shorter sequences
      """
  
      @property
      def input_ports(self):
          """Returns definitions of module input ports.
  
          hidden_states_src:
              0: AxisType(BatchTag)
  
              1: AxisType(TimeTag)
  
              2: AxisType(ChannelTag)
  
          input_mask_src:
              0: AxisType(BatchTag)
  
              1: AxisType(TimeTag)
          """
          return {
              "hidden_states_src":
              NeuralType({
                  0: AxisType(BatchTag),
                  1: AxisType(TimeTag),
                  2: AxisType(ChannelTag)
              }),
              "input_mask_src":
              NeuralType({
                  0: AxisType(BatchTag),
                  1: AxisType(TimeTag)
              })
          }
  
      @property
      def output_ports(self):
          """Returns definitions of module output ports.
  
          output_ids:
              0: AxisType(BatchTag)
  
              1: AxisType(TimeTag)
          """
          return {
              "output_ids": NeuralType({
                  0: AxisType(BatchTag),
                  1: AxisType(TimeTag)
              })
          }
  
      @property
      def num_weights(self):
          return 0
  
      def __init__(self,
                   decoder,
                   log_softmax,
                   max_seq_length,
                   pad_token,
                   bos_token,
                   eos_token,
                   batch_size=1,
                   beam_size=4,
                   max_delta_length=50,
                   length_penalty=0,
                   **kwargs):
          TrainableNM.__init__(self, **kwargs)
  
          self.generator = BeamSearchSequenceGenerator(
              decoder.embedding_layer,
              decoder.decoder,
              log_softmax,
              max_sequence_length=max_seq_length,
              max_delta_length=max_delta_length,
              pad=pad_token,
              bos=bos_token,
              eos=eos_token,
              batch_size=batch_size,
              beam_size=beam_size,
              len_pen=length_penalty
          )
  
      def forward(self, hidden_states_src, input_mask_src):
          output_ids = self.generator(encoder_hidden_states=hidden_states_src,
                                      encoder_input_mask=input_mask_src)
          return output_ids
  