# coding=utf-8
# Copyright 2021 The Trax Authors.
#
# 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.

# Lint as: python3
"""Layer-Skipping Transformer Models.

Like in https://arxiv.org/pdf/1909.11556.pdf
"""

from trax import layers as tl
from trax.layers.assert_shape import assert_shape
from trax.models import transformer


def LargerThan(val):
  """Checks if the input is larger than a certain value."""
  return tl.Fn('LargerThan', lambda x: x > val)


@assert_shape('...s->...sv')
def SkippingTransformerLM(vocab_size,
                          d_model=512,
                          d_ff=2048,
                          n_layers=6,
                          n_heads=8,
                          dropout=0.1,
                          max_len=2048,
                          mode='train',
                          ff_activation=tl.Relu,
                          skip_fraction=0.4):
  """Returns a Skipping Transformer language model.

  The input to the model is a tensor of tokens. (This model uses only the
  decoder part of the overall Transformer.)

  Args:
    vocab_size: int: vocab size
    d_model: int:  depth of embedding
    d_ff: int: depth of feed-forward layer
    n_layers: int: number of encoder/decoder layers
    n_heads: int: number of attention heads
    dropout: float: dropout rate (how much to drop out)
    max_len: int: maximum symbol length for positional encoding
    mode: str: 'train', 'eval' or 'predict', predict mode is for fast inference
    ff_activation: the non-linearity in feed-forward layer
    skip_fraction: fraction of times to skip some layers

  Returns:
    A Transformer language model as a layer that maps from a tensor of tokens
    to activations over a vocab set.
  """
  embedder = [
      tl.Embedding(vocab_size, d_model),
      tl.Dropout(rate=dropout, mode=mode),
      tl.PositionalEncoding(max_len=max_len, mode=mode),
  ]

  @assert_shape('...sd,->...sd,')
  def ConditionedBlock(current_layer_num):
    return tl.Serial(
        # stack: embedding, n_layers_to_keep
        tl.Select([1, 0, 1]),  # n_layers_to_keep, embedding, n_layers_to_keep
        tl.Cond(
            # if n_layers_to_keep > current_layer_num
            LargerThan(float(current_layer_num)),
            # then: run block
            tl.Serial(transformer._DecoderBlock(  # pylint: disable=g-complex-comprehension,protected-access
                d_model, d_ff, n_heads, dropout, [], mode, ff_activation)),
            # else: run noop
            tl.Serial()
            )
        # stack: embedding, n_layers_to_keep
        )

  if mode == 'train':
    if skip_fraction == 0.0:
      minimum_layers = float(n_layers)
      maximum_layers = float(n_layers)
    else:
      minimum_layers = 0.0
      maximum_layers = float(n_layers) / skip_fraction
  else:
    minimum_layers = maximum_layers = float(n_layers)

  return tl.Serial(
      tl.ShiftRight(mode=mode),
      embedder,
      # stack: embedding
      tl.RandomUniform(minimum_layers, maximum_layers, sync=True),
      # stack: n_layers_to_keep, embedding
      tl.Swap(),
      # stack: embedding, n_layers_to_keep
      [ConditionedBlock(i) for i in range(n_layers)],
      # stack: embedding, n_layers_to_keep
      tl.AssertShape('...sd,'),
      tl.Select([0], n_in=2),  # stack: embedding
      tl.AssertShape('...sd'),
      tl.LayerNorm(),
      tl.Dense(vocab_size),
  )


@assert_shape('...s->...sv')
def EveryOtherLayerDropTransformerLM(vocab_size,
                                     d_model=512,
                                     d_ff=2048,
                                     n_layers=6,
                                     n_heads=8,
                                     dropout=0.1,
                                     max_len=2048,
                                     mode='train',
                                     ff_activation=tl.Relu,
                                     skip_mode='even',
                                     skip_fraction=0.5,
                                     eval_skip_fraction=0.0):
  """Returns an "EveryOther" LayerDrop Transformer language model.

  During each training step it either runs all layers, or skips a subset of
  layers. This subset is the same every time, and it is specified by
  "skip_mode".
  The input to the model is a tensor of tokens. (This model uses only the
  decoder part of the overall Transformer.)

  Args:
    vocab_size: int: vocab size
    d_model: int:  depth of embedding
    d_ff: int: depth of feed-forward layer
    n_layers: int: number of encoder/decoder layers
    n_heads: int: number of attention heads
    dropout: float: dropout rate (how much to drop out)
    max_len: int: maximum symbol length for positional encoding
    mode: str: 'train', 'eval' or 'predict', predict mode is for fast inference
    ff_activation: the non-linearity in feed-forward layer
    skip_mode: which layers to skip when skipping: even/odd/1half/2half.
    skip_fraction: fraction of times to skip layers
    eval_skip_fraction: fraction of times to skip layers during eval

  Returns:
    A Transformer language model as a layer that maps from a tensor of tokens
    to activations over a vocab set.
  """
  embedder = [
      tl.Embedding(vocab_size, d_model),
      tl.Dropout(rate=dropout, mode=mode),
      tl.PositionalEncoding(max_len=max_len, mode=mode),
  ]

  if mode == 'train':
    pass
  else:
    skip_fraction = eval_skip_fraction

  skip_mode_funs = {  # which layers should be skipped?
      'even': (lambda num: num%2 == 0),  # 0th layer is even
      'odd': (lambda num: num%2 == 1),
      '1half': (lambda num: num < (n_layers/2)),
      '2half': (lambda num: num >= (n_layers/2)),
  }

  skip_mode_fun = skip_mode_funs[skip_mode]

  @assert_shape('...sd,->...sd,')
  def ConditionedBlock(current_layer_num):
    return tl.Serial(
        # stack: embedding, n_layers_to_keep
        tl.Select([1, 0, 1]),  # n_layers_to_keep, embedding, n_layers_to_keep
        tl.Cond(
            # if random() > skip_fraction OR layer not in skip_mode ...
            LargerThan(skip_fraction if skip_mode_fun(current_layer_num)
                       else 0.0),
            # then: run block
            tl.Serial(transformer._DecoderBlock(  # pylint: disable=g-complex-comprehension,protected-access
                d_model, d_ff, n_heads, dropout, [], mode, ff_activation))
            # else: noop (implicit)
            )
        # stack: embedding, n_layers_to_keep
        )

  return tl.Serial(
      tl.ShiftRight(mode=mode),
      embedder,
      # stack: embedding
      tl.RandomUniform(0., 1., sync=True),
      # stack: n_layers_to_keep, embedding
      tl.Swap(),
      # stack: embedding, n_layers_to_keep
      [ConditionedBlock(i) for i in range(n_layers)],
      # stack: embedding, n_layers_to_keep
      tl.Select([0], n_in=2),  # stack: embedding
      tl.LayerNorm(),
      tl.Dense(vocab_size),
  )


@assert_shape('...s->...sv')
def LayerDropTransformerLM(vocab_size,
                           d_model=512,
                           d_ff=2048,
                           n_layers=6,
                           n_heads=8,
                           dropout=0.1,
                           max_len=2048,
                           mode='train',
                           ff_activation=tl.Relu,
                           skip_fraction=0.4,
                           eval_skip_fraction='every_other'):
  """Returns a LayerDrop Transformer language model.

  Based on Fan, Grave, Joulin 2019, https://arxiv.org/abs/1909.11556 .

  The input to the model is a tensor of tokens. (This model uses only the
  decoder part of the overall Transformer.)

  Args:
    vocab_size: int: vocab size
    d_model: int:  depth of embedding
    d_ff: int: depth of feed-forward layer
    n_layers: int: number of encoder/decoder layers
    n_heads: int: number of attention heads
    dropout: float: dropout rate (how much to drop out)
    max_len: int: maximum symbol length for positional encoding
    mode: str: 'train', 'eval' or 'predict', predict mode is for fast inference
    ff_activation: the non-linearity in feed-forward layer
    skip_fraction: probability of skipping a layer; it can be a single
        probability or a list of probabilities different for each layer
    eval_skip_fraction: probability of skipping a layer during eval; it can be a
        single probability, or a list of probabilities different for each layer,
        or a string "every other" implementing a strategy from original paper

  Returns:
    A Transformer language model as a layer that maps from a tensor of tokens
    to activations over a vocab set.
  """
  embedder = [
      tl.Embedding(vocab_size, d_model),
      tl.Dropout(rate=dropout, mode=mode),
      tl.PositionalEncoding(max_len=max_len, mode=mode),
  ]

  if not isinstance(skip_fraction, (list, tuple)):
    # If we don't get a list of skip_fractions we use the same skip_fraction
    # for each layer.
    skip_fraction = [skip_fraction for i in range(n_layers)]
  if len(skip_fraction) != n_layers:
    raise ValueError('n_layers ({}) must be equal to len(skip_fraction) ({})'
                     .format(n_layers, len(skip_fraction)))

  if eval_skip_fraction == 'every_other':
    # 100% skipping for even-numbered layers; 0% for odd-numbered layers.
    eval_skip_fraction = [(1.0 if i % int(1./skip_fraction[i]) == 0 else 0.0)
                          if skip_fraction[i] != 0 else 0.0
                          for i in range(n_layers)]
  if eval_skip_fraction == 'same':
    # Same skip_fraction as in training.
    eval_skip_fraction = skip_fraction
  if not isinstance(eval_skip_fraction, (list, tuple)):
    # If we don't get a list of eval_skip_fractions we use the same
    # eval_skip_fraction for each layer.
    eval_skip_fraction = [eval_skip_fraction for i in range(n_layers)]
  if len(eval_skip_fraction) != n_layers:
    raise ValueError(
        'n_layers ({}) must be equal to len(eval_skip_fraction) ({})'
        .format(n_layers, len(eval_skip_fraction)))

  @assert_shape('...sd->...sd')
  def ConditionedBlock(current_layer_num):
    return tl.Serial(
        # stack: embedding
        tl.RandomUniform(0., 1, sync=True),
        # stack: random_uniform, embedding
        tl.Cond(
            # if random_uniform > skip_fraction
            LargerThan(skip_fraction[current_layer_num] if mode == 'train'
                       else eval_skip_fraction[current_layer_num]),
            # then: run block
            tl.Serial(transformer._DecoderBlock(  # pylint: disable=g-complex-comprehension,protected-access
                d_model, d_ff, n_heads, dropout, [], mode, ff_activation)),
            # else: run noop
            tl.Serial()
            )
        # stack: embedding
        )

  return tl.Serial(
      tl.ShiftRight(mode=mode),
      embedder,
      [ConditionedBlock(i) for i in range(n_layers)],
      tl.LayerNorm(),
      tl.Dense(vocab_size),
  )
