def Transformer(vocab_size,
d_feature=512,
d_feedforward=2048,
n_layers=6,
n_heads=8,
dropout=0.1,
max_len=2048,
mode='train'):
"""Transformer.
This model expects on input a pair (source, target).
Args:
vocab_size: int: vocab size (shared source and target).
d_feature: int: depth of embedding
d_feedforward: 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' or 'eval'
Returns:
the Transformer model.
"""
positional_embedder = [
tl.Embedding(d_feature, vocab_size),
tl.Dropout(rate=dropout, mode=mode),
tl.PositionalEncoding(max_len=max_len),
]
encoder = [
tl.Branch(positional_embedder, tl.PaddingMask()),
[
EncoderBlock(d_feature, d_feedforward, n_heads, dropout, mode)
for _ in range(n_layers)
],
tl.LayerNorm(),
]
return tl.Model(
tl.Parallel([], tl.ShiftRight()),
tl.Parallel(encoder, positional_embedder),
tl.Select(inputs=(('encoder', 'mask'), 'decoder'),
output=('decoder', ('mask', 'decoder'), 'encoder')),
# (encoder_mask, decoder_input) -> encoder-decoder mask
tl.Parallel([], tl.EncoderDecoderMask(), []),
[
EncoderDecoder(d_feature, d_feedforward, n_heads, dropout, mode)
for _ in range(n_layers)
],
tl.Select(0), # Drop mask and encoder.
tl.LayerNorm(),
tl.Dense(vocab_size),
tl.LogSoftmax(),
)
def Transformer(vocab_size,
feature_depth=512,
feedforward_depth=2048,
num_layers=6,
num_heads=8,
dropout=0.1,
max_len=2048,
mode='train'):
"""Transformer.
This model expects on input a pair (source, target).
Args:
vocab_size: int: vocab size (shared source and target).
feature_depth: int: depth of embedding
feedforward_depth: int: depth of feed-forward layer
num_layers: int: number of encoder/decoder layers
num_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' or 'eval'
Returns:
the Transformer model.
"""
embedding = layers.Serial(layers.Embedding(feature_depth, vocab_size),
layers.Dropout(rate=dropout, mode=mode),
layers.PositionalEncoding(max_len=max_len))
encoder = layers.Serial(
layers.Branch(), # Branch input to create embedding and mask.
layers.Parallel(embedding, layers.PaddingMask()),
layers.Serial(*[
EncoderLayer(feature_depth, feedforward_depth, num_heads, dropout,
mode) for _ in range(num_layers)
]),
layers.Parallel(layers.LayerNorm(), layers.Identity()))
stack = [
EncoderDecoderLayer(feature_depth, feedforward_depth, num_heads,
dropout, mode) for _ in range(num_layers)
]
return layers.Serial(
layers.Parallel(layers.Identity(), layers.ShiftRight()),
layers.Parallel(encoder, embedding),
layers.UnnestBranches(), # (encoder, encoder_mask, decoder_input)
layers.Reorder(output=(0, (1, 2), 2)),
layers.
Parallel( # (encoder_mask, decoder_input) -> encoder-decoder mask
layers.Identity(), layers.EncoderDecoderMask(), layers.Identity()),
layers.Serial(*stack),
layers.ThirdBranch(),
layers.LayerNorm(),
layers.Dense(vocab_size),
layers.LogSoftmax())
def Transformer(vocab_size,
feature_depth=512,
feedforward_depth=2048,
num_layers=6,
num_heads=8,
dropout=0.1,
max_len=2048,
mode='train'):
"""Transformer.
This model expects on input a pair (source, target).
Args:
vocab_size: int: vocab size (shared source and target).
feature_depth: int: depth of embedding
feedforward_depth: int: depth of feed-forward layer
num_layers: int: number of encoder/decoder layers
num_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' or 'eval'
Returns:
the Transformer model.
"""
embedding = tl.Serial(tl.Embedding(feature_depth, vocab_size),
tl.Dropout(rate=dropout, mode=mode),
tl.PositionalEncoding(max_len=max_len))
encoder = tl.Serial(
tl.Branch(embedding, tl.PaddingMask()),
tl.Serial(*[
EncoderLayer(feature_depth, feedforward_depth, num_heads, dropout,
mode) for _ in range(num_layers)
]), tl.Parallel(tl.LayerNorm(), tl.NoOp()))
stack = [
EncoderDecoderLayer(feature_depth, feedforward_depth, num_heads,
dropout, mode) for _ in range(num_layers)
]
return tl.Serial(
tl.Parallel(tl.NoOp(), tl.ShiftRight()),
tl.Parallel(encoder, embedding),
tl.Select(inputs=(('encoder', 'mask'), 'decoder'),
output=('encoder', ('mask', 'decoder'), 'decoder')),
tl.Parallel( # (encoder_mask, decoder_input) -> encoder-decoder mask
tl.NoOp(), tl.EncoderDecoderMask(), tl.NoOp()),
tl.Serial(*stack),
tl.Select(2), # Drop encoder and mask.
tl.LayerNorm(),
tl.Dense(vocab_size),
tl.LogSoftmax())
def Transformer(input_vocab_size,
output_vocab_size=None,
d_model=512,
d_ff=2048,
n_layers=6,
n_heads=8,
dropout=0.1,
max_len=2048,
mode='train'):
"""Returns a Transformer model.
This model expects an input pair: target, source.
Args:
input_vocab_size: int: vocab size of the source.
output_vocab_size: int (optional): vocab size of the target. If None, the
source and target are assumed to have the same vocab.
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' or 'eval'
Returns:
A Transformer model as a layer that maps from a target, source pair to
activations over a vocab set.
"""
in_embed = [ # tokens
tl.Embedding(d_model, input_vocab_size), # vecs
tl.Dropout(rate=dropout, mode=mode), # vecs
tl.PositionalEncoding(max_len=max_len), # vecs
]
if output_vocab_size is None:
output_vocab_size = input_vocab_size
out_embed = in_embed
else:
out_embed = [ # tokens
tl.Embedding(d_model, output_vocab_size), # vecs
tl.Dropout(rate=dropout, mode=mode), # vecs
tl.PositionalEncoding(max_len=max_len), # vecs
]
encoder_stack = ( # masks vectors --> masks vectors
[
EncoderBlock(d_model, d_ff, n_heads, dropout, i, mode)
for i in range(n_layers)
])
encoder_decoder_stack = ( # vecs_d masks vecs_e --> vecs_d masks vecs_e
[
EncoderDecoder(d_model, d_ff, n_heads, dropout, i, mode)
for i in range(n_layers)
])
# Input: encoder_side_tokens, decoder_side_tokens
return tl.Model( # tokens_e tokens_d
tl.Swap(), # toks_d toks_e
# Encode.
tl.Parallel( # toks_d toks_e
[],
[
tl.Dup(), # ______ toks_e toks_e
tl.Parallel(in_embed, tl.PaddingMask()), # ______ vecs_e masks
encoder_stack, # ______ vecs_e masks
tl.LayerNorm(), # ______ vecs_e .....
tl.Swap()
]), # ______ masks vecs_e
# Decode. # toks_d masks vecs_e
tl.ShiftRight(), # toks_d ..... ......
out_embed, # vecs_d ..... ......
tl.Dup(), # vecs_d vecs_d ..... ......
tl.Parallel([], tl.EncoderDecoderMask()), # ______ masks ......
encoder_decoder_stack, # vecs_d masks vecs_e
tl.Parallel([], tl.Drop(), tl.Drop()), # vecs_d
tl.LayerNorm(), # vecs_d
tl.Dense(output_vocab_size), # vecs_d
tl.LogSoftmax(), # vecs_d
)
