def Transformer(source_vocab_size,
                target_vocab_size,
                mode='train',
                num_layers=6,
                feature_depth=512,
                feedforward_depth=2048,
                num_heads=8,
                dropout=0.1,
                shared_embedding=True,
                max_len=200,
                return_evals=False):
  """Transformer model.

  Args:
    source_vocab_size: int: source vocab size
    target_vocab_size: int: target vocab size
    mode: str: 'train' or 'eval'
    num_layers: int: number of encoder/decoder layers
    feature_depth: int:  depth of embedding
    feedforward_depth: int: depth of feed-forward layer
    num_heads: int: number of attention heads
    dropout: float: dropout rate (how much to drop out)
    shared_embedding: bool: specify whether source/target embeddings are tied.
    max_len: int: maximum symbol length for positional encoding
    return_evals: bool: whether to generate decode-time evaluation functions

  Returns:
    A namedtuple containing model 'init' and 'apply' functions for training and
  the 'evals' functions that itself returns a namedtuple containing evaluation
  functions for the trained encoder, decoder, and generator substax.
  """
  # Input embedding and positional encoding
  inject_position = layers.Serial(
      layers.Dropout(dropout, mode=mode),
      layers.PositionalEncoding(feature_depth, max_len=max_len)
  )
  if shared_embedding:
    assert source_vocab_size == target_vocab_size
    # Weight-shared Embedding
    embedding = layers.Share(layers.Embedding(feature_depth, source_vocab_size))
    source_embedding_layer = layers.Serial(embedding, inject_position)
    target_embedding_layer = source_embedding_layer
  else:
    source_embedding = layers.Embedding(feature_depth, source_vocab_size)
    target_embedding = layers.Embedding(feature_depth, target_vocab_size)
    source_embedding_layer = layers.Serial(source_embedding, inject_position)
    target_embedding_layer = layers.Serial(target_embedding, inject_position)

  # Multi-headed Attention and Feed-forward layers
  multi_attention = layers.MultiHeadedAttention(
      feature_depth, num_heads=num_heads, dropout=dropout, mode=mode)

  # Encoder
  @layers.Lambda
  def Encoder(source, source_mask):
    """Transformer encoder stack.

    Args:
      source: layer variable: raw source sequences
      source_mask: layer variable: self-attention mask

    Returns:
      Layer variable that outputs encoded source.
    """
    encoder_layer = layers.Serial(
        # input attends to self
        layers.Residual(layers.LayerNorm(),
                        layers.Branch(size=4),
                        layers.Parallel(layers.Identity(),  # query
                                        layers.Identity(),  # key
                                        layers.Identity(),  # value
                                        source_mask),  # attention mask
                        multi_attention,
                        layers.Dropout(dropout, mode=mode)),
        # feed-forward
        ResidualFeedForward(
            feature_depth, feedforward_depth, dropout, mode=mode),
    )
    return layers.Serial(
        source,
        source_embedding_layer,
        layers.repeat(encoder_layer, num_layers),
        layers.LayerNorm(),
    )

  # Decoder
  @layers.Lambda
  def Decoder(memory, target, target_mask, memory_mask):
    """Transformer decoder stack.

    Args:
      memory: layer variable: encoded source sequences
      target: layer variable: raw target sequences
      target_mask: layer variable: self-attention mask
      memory_mask: layer variable: memory attention mask

    Returns:
      Layer variable that outputs encoded source.
    """
    decoder_layer = layers.Serial(
        # target attends to self
        layers.Residual(layers.LayerNorm(),
                        layers.Branch(size=4),
                        layers.Parallel(layers.Identity(),  # query
                                        layers.Identity(),  # key
                                        layers.Identity(),  # value
                                        target_mask),  # attention mask
                        multi_attention,
                        layers.Dropout(dropout, mode=mode)),
        # target attends to encoded source
        layers.Residual(layers.LayerNorm(),
                        layers.Branch(size=4),
                        layers.Parallel(layers.Identity(),  # query
                                        memory,  # key
                                        memory,  # value
                                        memory_mask),  # attention mask
                        multi_attention,
                        layers.Dropout(dropout, mode=mode)),
        # feed-forward
        ResidualFeedForward(
            feature_depth, feedforward_depth, dropout, mode=mode)
    )
    return layers.Serial(
        target,
        target_embedding_layer,
        layers.repeat(decoder_layer, num_layers),
        layers.LayerNorm(),
    )

  # The Transformer
  @layers.Lambda
  def transformer(source, target, source_mask, target_mask, memory_mask):  # pylint: disable=invalid-name
    encoded_source = Encoder(source, source_mask)
    return Decoder(encoded_source, target, target_mask, memory_mask)

  # Finally, bind the generator transform to use later for inference.
  @layers.Lambda
  def Generator(encoded_target):
    return layers.Serial(
        encoded_target,
        layers.Dense(target_vocab_size),
        layers.LogSoftmax
    )

  # Model-Building and Evaluation Functions
  # Get entire model's the layer pair
  top_init, top_apply = Generator(transformer)

  # By default act as a normal constructor and emit an (init, apply) pair.
  if not return_evals:
    return (top_init, top_apply)
  else:
    raise ValueError('inference in this model is still a work in progress')
def Transformer ( source_vocab_size , target_vocab_size , mode = 'train' , num_layers = 6 , feature_depth = 512 , feedforward_depth = 2048 , num_heads = 8 , dropout = 0.1 , shared_embedding = True , max_len = 200 , return_evals = False ) : inject_position = layers . Serial ( layers . Dropout ( dropout , mode = mode ) , layers . PositionalEncoding ( feature_depth , max_len = max_len ) ) if shared_embedding : assert source_vocab_size == target_vocab_size embedding = layers . Share ( layers . Embedding ( feature_depth , source_vocab_size ) ) source_embedding_layer = layers . Serial ( embedding , inject_position ) target_embedding_layer = source_embedding_layer else : source_embedding = layers . Embedding ( feature_depth , source_vocab_size ) target_embedding = layers . Embedding ( feature_depth , target_vocab_size ) source_embedding_layer = layers . Serial ( source_embedding , inject_position ) target_embedding_layer = layers . Serial ( target_embedding , inject_position ) multi_attention = layers . MultiHeadedAttention ( feature_depth , num_heads = num_heads , dropout = dropout , mode = mode ) @ layers . Lambda def Encoder ( source , source_mask ) : encoder_layer = layers . Serial ( layers . Residual ( layers . LayerNorm ( ) , layers . Branch ( size = 4 ) , layers . Parallel ( layers . Identity ( ) , layers . Identity ( ) , layers . Identity ( ) , source_mask ) , multi_attention , layers . Dropout ( dropout , mode = mode ) ) , ResidualFeedForward ( feature_depth , feedforward_depth , dropout , mode = mode ) , ) return layers . Serial ( source , source_embedding_layer , layers . repeat ( encoder_layer , num_layers ) , layers . LayerNorm ( ) , ) @ layers . Lambda def Decoder ( memory , target , target_mask , memory_mask ) : decoder_layer = layers . Serial ( layers . Residual ( layers . LayerNorm ( ) , layers . Branch ( size = 4 ) , layers . Parallel ( layers . Identity ( ) , layers . Identity ( ) , layers . Identity ( ) , target_mask ) , multi_attention , layers . Dropout ( dropout , mode = mode ) ) , layers . Residual ( layers . LayerNorm ( ) , layers . Branch ( size = 4 ) , layers . Parallel ( layers . Identity ( ) , memory , memory , memory_mask ) , multi_attention , layers . Dropout ( dropout , mode = mode ) ) , ResidualFeedForward ( feature_depth , feedforward_depth , dropout , mode = mode ) ) return layers . Serial ( target , target_embedding_layer , layers . repeat ( decoder_layer , num_layers ) , layers . LayerNorm ( ) , ) @ layers . Lambda def transformer ( source , target , source_mask , target_mask , memory_mask ) : encoded_source = Encoder ( source , source_mask ) return Decoder ( encoded_source , target , target_mask , memory_mask ) @ layers . Lambda def Generator ( encoded_target ) : return layers . Serial ( encoded_target , layers . Dense ( target_vocab_size ) , layers . LogSoftmax ) top_init , top_apply = Generator ( transformer ) if not return_evals : return ( top_init , top_apply ) else : raise ValueError ( 'inference in this model is still a work in progress' )
Transformer model.

  Args:
    source_vocab_size: int: source vocab size
    target_vocab_size: int: target vocab size
    mode: str: 'train' or 'eval'
    num_layers: int: number of encoder/decoder layers
    feature_depth: int:  depth of embedding
    feedforward_depth: int: depth of feed-forward layer
    num_heads: int: number of attention heads
    dropout: float: dropout rate (how much to drop out)
    shared_embedding: bool: specify whether source/target embeddings are tied.
    max_len: int: maximum symbol length for positional encoding
    return_evals: bool: whether to generate decode-time evaluation functions

  Returns:
    A namedtuple containing model 'init' and 'apply' functions for training and
  the 'evals' functions that itself returns a namedtuple containing evaluation
  functions for the trained encoder, decoder, and generator substax.