code: #-*- coding:utf-8 -*-
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License 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.
  # ==============================================================================
  
  # pylint: disable=import-error, too-few-public-methods, unused-argument
  
  """model_helper.py: methods for models"""
  
  __author__ = "Kyungmin Lee"
  __email__ = "sephiroce@snu.ac.kr"
  
  import math
  import tensorflow as tf
  from tfsr.helper.common_helper import Constants
  
  #######################
  # Positional Encoding #
  #######################
  def get_pos_enc(length, hidden_size, min_timescale=1.0, max_timescale=1.0e4):
    """Return positional encoding.
    I borrowed from the official transformer model.
    URL: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py
  
    Calculates the position encoding as a mix of sine and cosine functions with
    geometrically increasing wavelengths.
    Defined and formulated in Attention is All You Need, section 3.5.
    Args:
      length: Sequence length.
      hidden_size: Size of the
      min_timescale: Minimum scale that will be applied at each position
      max_timescale: Maximum scale that will be applied at each position
    Returns:
      Tensor with shape [length, hidden_size]
    """
    # We compute the positional encoding in float32 even if the model uses
    # float16, as many of the ops used, like log and exp, are numerically unstable
    # in float16.
    position = tf.cast(tf.range(length), tf.float32)
    num_timescales = hidden_size // 2
    log_timescale_increment = (
        math.log(float(max_timescale) / float(min_timescale)) /
        (tf.cast(num_timescales, tf.float32) - 1))
    inv_timescales = min_timescale * tf.exp(
        tf.cast(tf.range(num_timescales), tf.float32) * -log_timescale_increment)
    scaled_time = tf.expand_dims(position, 1) * tf.expand_dims(inv_timescales, 0)
    signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1)
    return signal
  
  
  ###########
  # Masking #
  ###########
  def create_padding_mask(seq):
    """
    Mask all the pad tokens in the batch of sequence. It ensures that the model
    does not treat padding as the input. The mask indicates where pad value 0
    is present: it outputs a 1 at those locations, and a 0 otherwise.
  
    seq: a sequence padded with zeros
    """
    seq = tf.cast(tf.math.equal(seq, 0), tf.float32)
  
    # add extra dimensions to add the padding
    # to the attention logits.
    return seq[:, tf.newaxis, tf.newaxis, :]  # (batch_size, 1, 1, seq_len)
  
  
  def get_padding_bias(inp_len, strides=4):
    """Calculate bias tensor from padding values in tensor.
  
    Bias tensor that is added to the pre-softmax multi-headed attention logits,
    which has shape [batch_size, num_heads, length, length].
    The tensor is zero at non-padding locations, and -1e9 (negative infinity)
    at padding locations.
  
    Args:
      inp_len: int tensor with shape [batch_size], represents input speech lengths
      strides: time domain strides * the number of cnn layers
  
    Returns:
      Attention bias tensor of shape [batch_size, 1, 1, inp_len].
  
    """
    with tf.name_scope("attention_bias"):
      inp_len = tf.math.ceil(inp_len / strides)
      attention_bias = tf.abs(tf.sequence_mask(inp_len, dtype=tf.dtypes.float32) - 1.0)
    return attention_bias[:, tf.newaxis, tf.newaxis, :]
  
  
  def create_look_ahead_mask(size):
    """
    The look-ahead mask is used to mask the future tokens in a sequence. In other
    words, the mask indicates which entries should not be used.
  
    This means that to predict the third word, only the first and second word will
    be used. Similarly to predict the fourth word, only the first, second and the
    third word will be used and so on.
  
    size: the length of label sequences
    """
    mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)
    return mask  # (seq_len, seq_len)
  
  
  def create_combined_mask(tar):
    # Used in the 1st attention block in the decoder.
    # It is used to pad and mask future tokens in the input received by
    # the decoder.
    look_ahead_mask = create_look_ahead_mask(tf.shape(tar)[1])
    dec_target_padding_mask = create_padding_mask(tar)
    return tf.maximum(dec_target_padding_mask, look_ahead_mask)
  
  
  def feat_mask(args):
    """
    I do not know why it is not working!
    tf.cond(tf.rank(args[0]) == 4,
    lambda: tf.einsum('ijkl, ij->ijkl', args[0], mask),
    lambda: tf.einsum('ijk, ij->ijk', args[0], mask))
  
    [input_sequence, input_lengths]
    :param args:
    :return:
    """
    lengths = tf.math.ceil(tf.cast(args[1], tf.dtypes.int32) / args[2])
    mask = tf.sequence_mask(lengths, dtype=args[0].dtype)
  
    result = tf.einsum('ijkl, ij->ijkl', args[0], mask)
    return result
  
  
  def feat_mask2(args):
    """
    [input_sequence, input_lengths]
    :param args:
    :return:
    """
    lengths = tf.math.ceil(tf.cast(args[1], tf.dtypes.int32) / args[2])
    mask = tf.sequence_mask(lengths, dtype=args[0].dtype)
  
    result = tf.einsum('ijk, ij->ijk', args[0], mask)
    return result
  
  
  def get_init(init):
    if init == Constants.INIT_FANAVG:
      return tf.keras.initializers.VarianceScaling(scale=1.0, mode='fan_avg',
                                                   distribution='uniform',
                                                   seed=None)
    elif init == Constants.INIT_UNIFORM:
      return tf.keras.initializers.RandomUniform(minval=-0.05, maxval=0.05,
                                                 seed=None)
    return Constants.INIT_GLOROT
  
  
  def get_decoder_self_attention_bias(length, dtype=tf.float32):
    """Calculate bias for decoder that maintains model's autoregressive property.
    Creates a tensor that masks out locations that correspond to illegal
    connections, so prediction at position i cannot draw information from future
    positions.
    Args:
      length: int length of sequences in batch.
      dtype: The dtype of the return value.
    Returns:
      float tensor of shape [1, 1, length, length]
    """
    neg_inf = -1e9
    with tf.name_scope("decoder_self_attention_bias"):
      valid_locs = tf.linalg.band_part(tf.ones([length, length], dtype=dtype),
                                       -1, 0)
      valid_locs = tf.reshape(valid_locs, [1, 1, length, length])
      decoder_bias = neg_inf * (1.0 - valid_locs)
    return decoder_bias
  
  #####################
  # Attention penalty #
  #####################
  def create_attention_penalty(config, logger):
    #pylint: disable=too-many-boolean-expressions
    # Creating attention penalty
    if (config.model_ap_encoder or config.model_ap_decoder or config.model_ap_encdec) and \
        config.model_ap_width_zero is not None and config.model_ap_width_zero > 0 and \
        config.model_ap_width_stripe is not None and config.model_ap_width_stripe > 0 and \
        config.model_ap_scale is not None and config.model_ap_scale > 0.0:
      att_pen = AttentionPenalty(max_len=2500, # 100s= 100Kms = 2500 * 4 * 10 ms
                                 num_head=config.model_att_head_num,
                                 zero_width=config.model_ap_width_zero,
                                 stripe_width=config.model_ap_width_stripe,
                                 scale=config.model_ap_scale)
      logger.info("An attention penalty board was built with a zero width %d, "
                  "a stripe width %d and a scale factor %f",
                  config.model_ap_width_zero, config.model_ap_width_stripe,
                  config.model_ap_scale)
      logger.info("Attention penalties mask will be applied for")
      if config.model_ap_encoder:
        logger.info("> Encoder self-attention")
      if config.model_ap_decoder:
        logger.info("> Decoder self-attention")
      if config.model_ap_encdec:
        logger.info("> Encoder-Decoder attention")
    else:
      att_pen = None
      logger.info("Attention penalties will not be applied.")
  
    return att_pen
  
  
  class AttentionPenalty:
    #pylint: disable=too-many-arguments
    def __init__(self, max_len, num_head, zero_width, stripe_width, scale):
      att_penalty = tf.ones([max_len, max_len])
      self.eap = tf.zeros(([num_head, max_len, max_len]))
      for i in range(zero_width - 1, max_len, stripe_width):
        self.eap += tf.abs(1 - tf.linalg.band_part(att_penalty, i, i,
                                                   name="attention_penalty_board"))
      self.num_head = num_head
      self.max_len = max_len
      self.eap *= scale
  
    @property
    def big_penalty_map(self):
      return self.eap
  
    def create_eap(self, inp_len):
      #pylint: disable=too-many-locals
      inp_len = tf.cast(inp_len, tf.int32)
      enc_max_len = tf.math.reduce_max(inp_len, keepdims=False)
      return self.eap[:, :enc_max_len, :enc_max_len]
  
    def create_pens(self, ap_enc, ap_dec, inp_len, tar_len=None):
      #pylint: disable=too-many-locals
      inp_len = tf.cast(inp_len, tf.int32)
      enc_max_len = tf.math.reduce_max(inp_len, keepdims=False)
      enc_att_pen = self.eap[:, :enc_max_len, :enc_max_len]
  
      enc_dec_att_pen, dec_att_pen, dec_max_len = None, None, None
  
      if tar_len is not None:
        dec_max_len = tf.cast(tf.math.reduce_max(tar_len, keepdims=False), tf.int32)
  
        if ap_enc:
          enc_dec_att_pen = self.eap[ :, :dec_max_len, :enc_max_len]
  
        if ap_dec:
          dec_att_pen = self.eap[ :, :dec_max_len, :dec_max_len]
  
      return enc_att_pen, enc_dec_att_pen, dec_att_pen
  
    def get_enc_att_pen(self, inp_len):
      #pylint: disable=too-many-locals
      inp_len = tf.cast(inp_len, tf.int32)
      enc_max_len = tf.math.reduce_max(inp_len, keepdims=False)
      return self.eap[:, :enc_max_len, :enc_max_len]
  
apis: tensorflow.zeros tensorflow.cast tensorflow.math.equal tensorflow.keras.initializers.VarianceScaling tensorflow.math.ceil tensorflow.name_scope tensorflow.shape tensorflow.math.reduce_max tensorflow.sequence_mask tensorflow.sin tensorflow.cos tensorflow.range tensorflow.keras.initializers.RandomUniform tensorflow.maximum tensorflow.linalg.band_part tensorflow.reshape tensorflow.expand_dims tensorflow.ones tensorflow.einsum
api_extract: [(122, 'tensorflow.maximum', 'tf.maximum', (['dec_target_padding_mask', 'look_ahead_mask'], {}), True, 'import tensorflow as tf\n'), (137, 'tensorflow.sequence_mask', 'tf.sequence_mask', (['lengths'], {'dtype': 'args[0].dtype'}), True, 'import tensorflow as tf\n'), (139, 'tensorflow.einsum', 'tf.einsum', (['"""ijkl, ij->ijkl"""', 'args[0]', 'mask'], {}), True, 'import tensorflow as tf\n'), (150, 'tensorflow.sequence_mask', 'tf.sequence_mask', (['lengths'], {'dtype': 'args[0].dtype'}), True, 'import tensorflow as tf\n'), (152, 'tensorflow.einsum', 'tf.einsum', (['"""ijk, ij->ijk"""', 'args[0]', 'mask'], {}), True, 'import tensorflow as tf\n'), (49, 'tensorflow.range', 'tf.range', (['length'], {}), True, 'import tensorflow as tf\n'), (56, 'tensorflow.expand_dims', 'tf.expand_dims', (['position', '(1)'], {}), True, 'import tensorflow as tf\n'), (56, 'tensorflow.expand_dims', 'tf.expand_dims', (['inv_timescales', '(0)'], {}), True, 'import tensorflow as tf\n'), (72, 'tensorflow.math.equal', 'tf.math.equal', (['seq', '(0)'], {}), True, 'import tensorflow as tf\n'), (95, 'tensorflow.name_scope', 'tf.name_scope', (['"""attention_bias"""'], {}), True, 'import tensorflow as tf\n'), (96, 'tensorflow.math.ceil', 'tf.math.ceil', (['(inp_len / strides)'], {}), True, 'import tensorflow as tf\n'), (158, 'tensorflow.keras.initializers.VarianceScaling', 'tf.keras.initializers.VarianceScaling', ([], {'scale': '(1.0)', 'mode': '"""fan_avg"""', 'distribution': '"""uniform"""', 'seed': 'None'}), True, 'import tensorflow as tf\n'), (179, 'tensorflow.name_scope', 'tf.name_scope', (['"""decoder_self_attention_bias"""'], {}), True, 'import tensorflow as tf\n'), (182, 'tensorflow.reshape', 'tf.reshape', (['valid_locs', '[1, 1, length, length]'], {}), True, 'import tensorflow as tf\n'), (222, 'tensorflow.ones', 'tf.ones', (['[max_len, max_len]'], {}), True, 'import tensorflow as tf\n'), (223, 'tensorflow.zeros', 'tf.zeros', (['[num_head, max_len, max_len]'], {}), True, 'import tensorflow as tf\n'), (237, 'tensorflow.cast', 'tf.cast', (['inp_len', 'tf.int32'], {}), True, 'import tensorflow as tf\n'), (238, 'tensorflow.math.reduce_max', 'tf.math.reduce_max', (['inp_len'], {'keepdims': '(False)'}), True, 'import tensorflow as tf\n'), (243, 'tensorflow.cast', 'tf.cast', (['inp_len', 'tf.int32'], {}), True, 'import tensorflow as tf\n'), (244, 'tensorflow.math.reduce_max', 'tf.math.reduce_max', (['inp_len'], {'keepdims': '(False)'}), True, 'import tensorflow as tf\n'), (262, 'tensorflow.cast', 'tf.cast', (['inp_len', 'tf.int32'], {}), True, 'import tensorflow as tf\n'), (263, 'tensorflow.math.reduce_max', 'tf.math.reduce_max', (['inp_len'], {'keepdims': '(False)'}), True, 'import tensorflow as tf\n'), (53, 'tensorflow.cast', 'tf.cast', (['num_timescales', 'tf.float32'], {}), True, 'import tensorflow as tf\n'), (57, 'tensorflow.sin', 'tf.sin', (['scaled_time'], {}), True, 'import tensorflow as tf\n'), (57, 'tensorflow.cos', 'tf.cos', (['scaled_time'], {}), True, 'import tensorflow as tf\n'), (112, 'tensorflow.ones', 'tf.ones', (['(size, size)'], {}), True, 'import tensorflow as tf\n'), (120, 'tensorflow.shape', 'tf.shape', (['tar'], {}), True, 'import tensorflow as tf\n'), (136, 'tensorflow.cast', 'tf.cast', (['args[1]', 'tf.dtypes.int32'], {}), True, 'import tensorflow as tf\n'), (149, 'tensorflow.cast', 'tf.cast', (['args[1]', 'tf.dtypes.int32'], {}), True, 'import tensorflow as tf\n'), (162, 'tensorflow.keras.initializers.RandomUniform', 'tf.keras.initializers.RandomUniform', ([], {'minval': '(-0.05)', 'maxval': '(0.05)', 'seed': 'None'}), True, 'import tensorflow as tf\n'), (180, 'tensorflow.ones', 'tf.ones', (['[length, length]'], {'dtype': 'dtype'}), True, 'import tensorflow as tf\n'), (97, 'tensorflow.sequence_mask', 'tf.sequence_mask', (['inp_len'], {'dtype': 'tf.dtypes.float32'}), True, 'import tensorflow as tf\n'), (250, 'tensorflow.math.reduce_max', 'tf.math.reduce_max', (['tar_len'], {'keepdims': '(False)'}), True, 'import tensorflow as tf\n'), (55, 'tensorflow.range', 'tf.range', (['num_timescales'], {}), True, 'import tensorflow as tf\n'), (225, 'tensorflow.linalg.band_part', 'tf.linalg.band_part', (['att_penalty', 'i', 'i'], {'name': '"""attention_penalty_board"""'}), True, 'import tensorflow as tf\n')]