text: # Copyright 2017--2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License"). You may not
  # use this file except in compliance with the License. A copy of the License
  # is located at
  #
  #     http://aws.amazon.com/apache2.0/
  #
  # or in the "license" file accompanying this file. This file is distributed on
  # an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
  # express or implied. See the License for the specific language governing
  # permissions and limitations under the License.
  
  """
  Decoders for sequence-to-sequence models.
  """
  import logging
  from abc import abstractmethod
  from itertools import islice
  from typing import Dict, List, Optional, Tuple, Type, Union
  
  import torch as pt
  
  from . import constants as C
  from . import layers
  from . import transformer
  from .transformer import TransformerConfig
  
  logger = logging.getLogger(__name__)
  DecoderConfig = Union[TransformerConfig]  # type: ignore
  
  
  def get_decoder(config: DecoderConfig, inference_only: bool = False) -> 'Decoder':
      return Decoder.get_decoder(config, inference_only)
  
  
  class Decoder(pt.nn.Module):
      """
      Generic decoder interface.
      A decoder needs to implement code to decode a target sequence known in advance (decode_sequence),
      and code to decode a single word given its decoder state (decode_step).
      The latter is typically used for inference graphs in beam search.
      For the inference module to be able to keep track of decoder's states
      a decoder provides methods to return initial states (init_states), state variables and their shapes.
      """
  
      __registry = {}  # type: Dict[Type[DecoderConfig], Type['Decoder']]
  
      @classmethod
      def register(cls, config_type: Type[DecoderConfig]):
          """
          Registers decoder type for configuration. Suffix is appended to decoder prefix.
  
          :param config_type: Configuration type for decoder.
  
          :return: Class decorator.
          """
  
          def wrapper(target_cls):
              cls.__registry[config_type] = target_cls
              return target_cls
  
          return wrapper
  
      @classmethod
      def get_decoder(cls, config: DecoderConfig, inference_only: bool) -> 'Decoder':
          """
          Creates decoder based on config type.
  
          :param config: Decoder config.
          :param inference_only: Create a decoder that is only used for inference.
  
          :return: Decoder instance.
          """
          config_type = type(config)
          if config_type not in cls.__registry:
              raise ValueError('Unsupported decoder configuration %s' % config_type.__name__)
          decoder_cls = cls.__registry[config_type]
          return decoder_cls(config=config, inference_only=inference_only)  # type: ignore
  
      @abstractmethod
      def __init__(self):
          super().__init__()
  
      @abstractmethod
      def state_structure(self) -> str:
          raise NotImplementedError()
  
      @abstractmethod
      def init_state_from_encoder(self,
                                  encoder_outputs: pt.Tensor,
                                  encoder_valid_length: Optional[pt.Tensor] = None,
                                  target_embed: Optional[pt.Tensor] = None) -> List[pt.Tensor]:
          raise NotImplementedError()
  
      @abstractmethod
      def decode_seq(self, inputs: pt.Tensor, states: List[pt.Tensor]) -> pt.Tensor:
          """
          Decodes a sequence of embedded target words and returns sequence of last decoder
          representations for each time step.
  
          :param inputs: Encoded source: (batch_size, source_encoded_max_length, encoder_depth).
          :param states: List of initial states, as given by init_state_from_encoder().
          :return: Decoder output. Shape: (batch_size, target_embed_max_length, decoder_depth).
          """
          raise NotImplementedError()
  
      @abstractmethod
      def get_num_hidden(self):
          raise NotImplementedError()
  
  
  @Decoder.register(TransformerConfig)
  class TransformerDecoder(Decoder):
      """
      Transformer decoder as in Vaswani et al, 2017: Attention is all you need.
      In training, computation scores for each position of the known target sequence are computed in parallel,
      yielding most of the speedup.
      At inference time, the decoder block is evaluated again and again over a maximum length input sequence that is
      initially filled with zeros and grows during beam search with predicted tokens. Appropriate masking at every
      time-step ensures correct self-attention scores and is updated with every step.
  
      :param config: Transformer configuration.
      :param inference_only: Only use the model for inference enabling some optimizations,
                             such as disabling the auto-regressive mask.
      """
  
      def __init__(self, config: TransformerConfig, inference_only: bool = False) -> None:
          Decoder.__init__(self)
          pt.nn.Module.__init__(self)
          self.config = config
          self.inference_only = inference_only
          self.pos_embedding = layers.PositionalEmbeddings(weight_type=self.config.positional_embedding_type,
                                                           num_embed=self.config.model_size,
                                                           max_seq_len=self.config.max_seq_len_target,
                                                           scale_up_input=True,
                                                           scale_down_positions=False)
          self.autoregressive_mask = transformer.AutoRegressiveMask()
  
          self.layers = pt.nn.ModuleList(  # using ModuleList because we have additional inputs
              transformer.TransformerDecoderBlock(config, inference_only=self.inference_only)
              for _ in range(config.num_layers))
  
          self.final_process = transformer.TransformerProcessBlock(sequence=config.preprocess_sequence,
                                                                   dropout=config.dropout_prepost,
                                                                   num_hidden=self.config.model_size)
          if self.config.dropout_prepost > 0.0:
              self.dropout = pt.nn.Dropout(p=self.config.dropout_prepost, inplace=inference_only)
  
      def state_structure(self) -> str:
          """
          Returns the structure of states used for manipulation of the states.
          Each state is either labeled 's' for step, 'b' for source_mask, 'd' for decoder, or 'e' for encoder.
          """
          structure = ''
          if self.inference_only:
              structure += C.STEP_STATE + C.MASK_STATE + C.ENCODER_STATE * self.config.num_layers
          else:
              structure += C.STEP_STATE + C.ENCODER_STATE + C.MASK_STATE
  
          total_num_states = sum(layer.num_state_tensors for layer in self.layers)
          structure += C.DECODER_STATE * total_num_states
  
          return structure
  
      def init_state_from_encoder(self,
                                  encoder_outputs: pt.Tensor,
                                  encoder_valid_length: Optional[pt.Tensor] = None,
                                  target_embed: Optional[pt.Tensor] = None) -> List[pt.Tensor]:
          """
          Returns the initial states given encoder output. States for teacher-forced training are encoder outputs
          and a valid length mask for encoder outputs.
          At inference, this method returns the following state tuple:
          valid length bias, step state,
          [projected encoder attention keys, projected encoder attention values] * num_layers,
          [autoregressive state dummies] * num_layers.
  
          :param encoder_outputs: Encoder outputs. Shape: (batch, source_length, encoder_dim).
          :param encoder_valid_length: Valid lengths of encoder outputs. Shape: (batch,).
          :param target_embed: Target-side embedding layer output. Shape: (batch, target_length, target_embedding_dim).
          :return: Initial states.
          """
          source_max_len = encoder_outputs.size()[1]
          if target_embed is None:  # Inference: initial step = 0. Shape: (batch_size, 1)
              steps = pt.zeros_like(encoder_valid_length).unsqueeze(1)
              # (batch * heads, 1, source_max_len)
              source_mask = layers.prepare_source_length_mask(encoder_valid_length, self.config.attention_heads,
                                                              source_max_len)
              # Shape: (batch, heads, 1, src_max_len)
              source_mask = source_mask.view(-1, self.config.attention_heads, 1, source_max_len)
          else:  # Training: steps up to target length. Shape: (1, target_length)
              target_length = target_embed.size()[1]
              steps = pt.arange(0, target_length, device=target_embed.device).unsqueeze(0)
              # (batch * heads, 1, source_max_len)
              source_mask = layers.prepare_source_length_mask(encoder_valid_length, self.config.attention_heads,
                                                              source_max_len)
              source_mask = source_mask.expand(-1, target_length, -1)  # Shape: (batch * heads, trg_max_len, src_max_len)
  
              # Shape: (batch, heads, trg_max_len, src_max_len)
              source_mask = source_mask.view(-1, self.config.attention_heads, target_length, source_max_len)
  
          if self.inference_only:
              # Encoder projection caching, therefore we don't pass the encoder_outputs
              states = [steps, source_mask]
              encoder_outputs_t = encoder_outputs.transpose(1, 0)  # time-major layout
              for layer in self.layers:
                  enc_att_kv = layer.enc_attention.ff_kv(encoder_outputs_t)
                  states.append(enc_att_kv)
          else:
              # NO encoder projection caching
              states = [steps, encoder_outputs.transpose(1, 0), source_mask]
  
          _batch_size = encoder_outputs.size()[0]
          _device = encoder_outputs.device
          _dtype = encoder_outputs.dtype
          dummy_autoregr_states = [pt.zeros(layer.get_states_shape(_batch_size), device=_device, dtype=_dtype)
                                   for layer in self.layers
                                   for _ in range(layer.num_state_tensors)]
  
          states += dummy_autoregr_states
          return states
  
      def decode_seq(self, inputs: pt.Tensor, states: List[pt.Tensor]) -> pt.Tensor:
          """
          Decodes a sequence of embedded target words and returns sequence of last decoder
          representations for each time step.
  
          :param inputs: Encoded source: (batch_size, source_encoded_max_length, encoder_depth).
          :param states: List of initial states, as given by init_state_from_encoder().
          :return: Decoder output. Shape: (batch_size, target_embed_max_length, decoder_depth).
          """
          outputs, _ = self.forward(inputs, states)
          return outputs
  
      def forward(self, step_input: pt.Tensor, states: List[pt.Tensor]) -> Tuple[pt.Tensor, List[pt.Tensor]]:
          target_mask = None
          if self.inference_only:
              steps, source_mask, *other = states
              source_encoded = None  # use constant pre-computed key value projections from the states
              enc_att_kv = other[:self.config.num_layers]
              autoregr_states = other[self.config.num_layers:]
          else:
              if any(layer.needs_mask for layer in self.layers):
                  target_mask = self.autoregressive_mask(step_input)  # mask: (length, length)
              steps, source_encoded, source_mask, *autoregr_states = states
              enc_att_kv = [None for _ in range(self.config.num_layers)]
  
          if any(layer.num_state_tensors > 1 for layer in self.layers):
              # separates autoregressive states by layer
              states_iter = iter(autoregr_states)
              autoregr_states = [list(islice(states_iter, 0, layer.num_state_tensors)) for layer in self.layers]  # type: ignore
  
          batch, heads, target_max_len, source_max_len = source_mask.size()
          source_mask_view = source_mask.view(batch * heads, target_max_len, source_max_len)
  
          # target: (batch_size, length, model_size)
          target = self.pos_embedding(step_input, steps)
          # (length, batch_size, model_size)
          target = target.transpose(1, 0)
  
          if self.config.dropout_prepost > 0.0:
              target = self.dropout(target)
  
          new_autoregr_states = []  # type: List[pt.Tensor]
          for layer, layer_autoregr_state, layer_enc_att_kv in zip(self.layers, autoregr_states, enc_att_kv):
              target, new_layer_autoregr_state = layer(target=target,
                                                       target_mask=target_mask,
                                                       source=source_encoded,
                                                       source_mask=source_mask_view,
                                                       autoregr_states=layer_autoregr_state,
                                                       enc_att_kv=layer_enc_att_kv)
  
              new_autoregr_states += [*new_layer_autoregr_state]
  
          target = self.final_process(target)
          target = target.transpose(1, 0)
  
          # Inference: increment steps by 1 (discarded in training)
          steps = steps + 1
  
          if self.inference_only:
              # pass in cached encoder states
              encoder_attention_keys_values = states[2:2 + self.config.num_layers]
              new_states = [steps, states[1]] + encoder_attention_keys_values + new_autoregr_states
          else:
              new_states = [steps, states[1], states[2]] + new_autoregr_states
  
          return target, new_states
  
      def get_num_hidden(self):
          return self.config.model_size
  