def Transformer(input_vocab_size,
output_vocab_size=None,
d_model=D_MODEL,
d_ff=D_FF,
n_encoder_layers=N_LAYERS,
n_decoder_layers=N_LAYERS,
n_heads=N_HEADS,
max_len=MAX_SEQUENCE_LENGTH,
dropout=DROPOUT_RATE,
dropout_shared_axes=DROPOUT_SHARED_AXES,
mode=MODE,
ff_activation=FF_ACTIVATION_TYPE):
"""Returns a full Transformer model.
This model is an encoder-decoder that performs tokenized string-to-string
("source"-to-"target") transduction:
- inputs (2):
- source: Array representing a batch of text strings via token
IDs plus padding markers; shape is (batch_size, sequence_length),
where sequence_length <= ``max_len``. Array elements are integers in
``range(input_vocab_size)``, and 0 values mark padding positions.
- target: Array representing a batch of text strings via token
IDs plus padding markers; shape is (batch_size, sequence_length),
where sequence_length <= ``max_len``. Array elements are integers in
``range(output_vocab_size)``, and 0 values mark padding positions.
- output: 3-D array of raw activations with last/innermost dimension of
``output_vocab_size``, suitable for decoding into a batch of token
strings; shape is (batch_size, sequence_length, ``vocab_size``).
An example use would be to translate (tokenized) sentences from English to
German.
Args:
input_vocab_size: Input vocabulary size -- each element of the input tensor
should be an integer in ``range(vocab_size)``. These integers typically
represent token IDs from a vocabulary-based tokenizer.
output_vocab_size: If specified, gives the vocabulary size for the targets;
if ``None``, then input and target integers (token IDs) are assumed to
come from the same vocabulary.
d_model: Last/innermost dimension of activation arrays at most points in
the model, including the initial embedding output.
d_ff: Last/innermost dimension of special (typically wider)
:py:class:`Dense` layer in the feedforward part of each encoder block.
n_encoder_layers: Number of encoder blocks.
n_decoder_layers: Number of decoder blocks.
n_heads: Number of attention heads.
max_len: Maximum symbol length for positional encoding.
dropout: Stochastic rate (probability) for dropping an activation value
when applying dropout within encoder/decoder blocks. The same rate is
also used for attention dropout in encoder/decoder blocks.
dropout_shared_axes: Tensor axes on which to share a dropout mask.
Sharing along batch and sequence axes (``dropout_shared_axes=(0,1)``)
is a useful way to save memory and apply consistent masks to activation
vectors at different sequence positions.
mode: If ``'predict'``, use fast inference. If ``'train'``, each
encoder/decoder block will include dropout; else, it will pass all
values through unaltered.
ff_activation: Type of activation function at the end of each
encoder/decoder block; must be an activation-type subclass of
:py:class:`Layer`.
Returns:
A Transformer model as a layer that maps from a source-target tokenized
text pair to activations over a vocab set.
"""
# Avoid 'predict' mode in encoder, since encoder doesn't run stepwise.
encoder_mode = 'eval' if mode == 'predict' else mode
# Share embedding weights if no separate output vocab size.
in_embedder = tl.Embedding(input_vocab_size, d_model)
if output_vocab_size is None:
out_embedder = in_embedder
output_vocab_size = input_vocab_size
else:
out_embedder = tl.Embedding(output_vocab_size, d_model)
def _Dropout():
return tl.Dropout(rate=dropout,
shared_axes=dropout_shared_axes,
mode=mode)
def _EncBlock():
return _EncoderBlock(d_model, d_ff, n_heads, dropout,
dropout_shared_axes, mode, ff_activation)
def _Encoder():
encoder = tl.Serial(
in_embedder,
_Dropout(),
tl.PositionalEncoding(max_len=max_len, mode=encoder_mode),
[_EncBlock() for _ in range(n_encoder_layers)],
tl.LayerNorm(),
)
return tl.Cache(encoder) if mode == 'predict' else encoder
def _EncDecBlock():
return _EncoderDecoderBlock(d_model, d_ff, n_heads, dropout,
dropout_shared_axes, mode, ff_activation)
# Input to model is encoder-side tokens and decoder-side tokens: tok_d, tok_e
# Model output is decoder-side vectors and decoder-side tokens: vec_d tok_d
return tl.Serial(
tl.Select([0, 1, 1]), # Copies decoder tokens for use in loss.
# Encode.
tl.Branch([], tl.PaddingMask()), # tok_e masks tok_d tok_d
_Encoder(),
# Decode.
tl.Select([2, 1, 0]), # Re-orders inputs: tok_d masks vec_e .....
tl.ShiftRight(mode=mode),
out_embedder,
_Dropout(),
tl.PositionalEncoding(max_len=max_len, mode=mode),
tl.Branch([], tl.EncoderDecoderMask()), # vec_d masks ..... .....
[_EncDecBlock() for _ in range(n_decoder_layers)],
tl.LayerNorm(),
tl.Select([0], n_in=3), # Drops masks and encoding vectors.
# Map vectors to match output vocab size.
tl.Dense(output_vocab_size),
)
def ReformerLM(vocab_size,
d_model=512,
d_ff=2048,
d_attention_key=64,
d_attention_value=64,
n_layers=6,
n_heads=8,
dropout=0.1,
max_len=2048,
n_chunks=0,
n_attention_chunks=1,
attention_type=tl.DotProductCausalAttention,
share_qk=False,
axial_pos_shape=(),
d_axial_pos_embs=None,
ff_activation=tl.FastGelu,
ff_use_sru=0,
ff_chunk_size=0,
mode='train'):
"""Reversible transformer language model (only uses a decoder, no encoder).
Args:
vocab_size: int: vocab size
d_model: int: depth of *each half* of the two-part features
d_ff: int: depth of feed-forward layer
d_attention_key: int: depth of key vector for each attention head
d_attention_value: int: depth of value vector for each attention head
n_layers: int: number of 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
n_chunks: int: number of chunks (must match input pipeline)
n_attention_chunks: int: number of chunks for attention
attention_type: class: attention class to use, such as DotProductAttention.
share_qk: bool, whether to share queries and keys.
axial_pos_shape: tuple of ints: input shape to use for the axial position
encoding. If unset, axial position encoding is disabled.
d_axial_pos_embs: tuple of ints: depth of position embedding for each axis.
Tuple length must match axial_pos_shape, and values must sum to d_model.
ff_activation: the non-linearity in feed-forward layer
ff_use_sru: int; if > 0, we use this many SRU layers instead of feed-forward
ff_chunk_size: int; if > 0, chunk feed-forward into this-sized chunks
mode: str: 'train', 'eval', or 'predict'
Returns:
the layer.
"""
if n_chunks == 0:
n_chunks = 1
concatenate_input_chunks = []
else:
concatenate_input_chunks = tl.Concatenate(n_items=n_chunks)
d_emb = d_model
if not axial_pos_shape:
positional_encoding = tl.PositionalEncoding(
max_len=max_len, dropout=dropout, mode=mode)
elif axial_pos_shape == 'fixed-base': # TODO(lukaszkaiser): remove this HACK
positional_encoding = tl.FixedBasePositionalEncoding(mode=mode)
d_emb //= 2
elif axial_pos_shape == 'infinite': # TODO(lukaszkaiser): remove this HACK
positional_encoding = tl.InfinitePositionalEncoding(affine=False)
elif axial_pos_shape == 'infinite-affine':
# TODO(lukaszkaiser): remove this HACK
positional_encoding = tl.InfinitePositionalEncoding()
else:
assert d_axial_pos_embs is not None
positional_encoding = tl.AxialPositionalEncoding(
shape=axial_pos_shape, d_embs=d_axial_pos_embs,
dropout_broadcast_dims=tuple(range(1, len(axial_pos_shape) + 1)),
dropout=dropout, mode=mode)
positional_embedder = [
tl.Embedding(d_emb, vocab_size),
BroadcastedDropout(rate=dropout, mode=mode), # pylint: disable=no-value-for-parameter
positional_encoding,
]
decoder_blocks = []
if isinstance(attention_type, (tuple, list)):
assert n_layers % len(attention_type) == 0
else:
attention_type = [attention_type]
for layer_idx in range(n_layers):
layer_attention_type = attention_type[layer_idx % len(attention_type)]
decoder_block = DecoderBlock(
d_model, d_ff, d_attention_key, d_attention_value, n_heads,
n_attention_chunks,
attention_type=layer_attention_type,
dropout=dropout,
share_qk=(share_qk or issubclass(layer_attention_type,
tl.LSHCausalAttention)),
ff_activation=ff_activation,
ff_use_sru=ff_use_sru,
ff_chunk_size=ff_chunk_size,
mode=mode)
decoder_blocks.append(decoder_block)
return tl.Serial(
concatenate_input_chunks,
tl.ShiftRight(mode=mode),
positional_embedder,
tl.Dup(),
tl.ReversibleSerial(decoder_blocks + [
SplitForOutput(n_sections=n_chunks, axis=-2), # pylint: disable=no-value-for-parameter
]),
Map([
# TODO(kitaev): Test whether dropout should go before or after the
# LayerNorm, and whether dropout broadcasting is needed here.
tl.LayerNorm(),
BroadcastedDropout(rate=dropout, mode=mode), # pylint: disable=no-value-for-parameter
tl.Dense(vocab_size),
tl.LogSoftmax(),
], n_sections=n_chunks),
)
def ReformerShortenLM(vocab_size,
shorten_factor=1,
d_embedding=256,
d_model=512,
d_ff=2048,
d_attention_key=64,
d_attention_value=64,
n_layers=6,
n_heads=8,
dropout=0.1,
max_len=2048,
attention_type=tl.SelfAttention,
axial_pos_shape=(),
d_axial_pos_embs=None,
ff_activation=tl.FastGelu,
ff_use_sru=0,
ff_chunk_size=0,
ff_sparsity=0,
attention_chunk_size=0,
mode='train'):
"""Reversible transformer language model with shortening.
When shorten_factor is F and processing an input of shape [batch, length],
we embed the (shifted-right) input and then group each F elements (on length)
into a single vector -- so that in the end we process a tensor of shape ::
[batch, length // F, d_model]
almost until the end -- at the end it's un-shortend and a SRU is applied.
This reduces the length processed inside the main model body, effectively
making the model faster but possibly slightly less accurate.
Args:
vocab_size: int: vocab size
shorten_factor: by how much to shorten, see above
d_embedding: the depth of the embedding layer and final logits
d_model: int: depth of *each half* of the two-part features
d_ff: int: depth of feed-forward layer
d_attention_key: int: depth of key vector for each attention head
d_attention_value: int: depth of value vector for each attention head
n_layers: int: number of 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
attention_type: class: attention class to use, such as SelfAttention.
axial_pos_shape: tuple of ints: input shape to use for the axial position
encoding. If unset, axial position encoding is disabled.
d_axial_pos_embs: tuple of ints: depth of position embedding for each axis.
Tuple length must match axial_pos_shape, values must sum to d_embedding.
ff_activation: the non-linearity in feed-forward layer
ff_use_sru: int; if > 0, we use this many SRU layers instead of feed-forward
ff_chunk_size: int; if > 0, chunk feed-forward into this-sized chunks
ff_sparsity: int, if > 0 use sparse feed-forward block with this sparsity
attention_chunk_size: int, if > 0 run attention chunked at this size
mode: str: 'train' or 'eval'
Returns:
the layer.
"""
assert mode != 'predict' # TODO(lukaszkaiser,kitaev): fast inference
positional_encoding = ct.PositionalEncoder(mode, dropout, max_len,
axial_pos_shape,
d_axial_pos_embs)
positional_embedder = [
tl.Embedding(vocab_size, d_embedding),
tl.Dropout(rate=dropout, shared_axes=[-2], mode=mode), # pylint: disable=no-value-for-parameter
positional_encoding,
]
decoder_blocks = []
if isinstance(attention_type, (tuple, list)):
assert n_layers % len(attention_type) == 0
else:
attention_type = [attention_type]
for layer_idx in range(n_layers):
layer_attention_type = attention_type[layer_idx % len(attention_type)]
decoder_block = DecoderBlock(d_model,
d_ff,
d_attention_key,
d_attention_value,
n_heads,
attention_type=layer_attention_type,
dropout=dropout,
ff_activation=ff_activation,
ff_dropout=dropout,
ff_use_sru=ff_use_sru,
ff_chunk_size=ff_chunk_size,
ff_sparsity=ff_sparsity,
attention_chunk_size=attention_chunk_size,
mode=mode)
decoder_blocks.append(decoder_block)
# pylint: disable=g-long-lambda
return tl.Serial(
tl.ShiftRight(),
positional_embedder,
tl.Dup(), # Stack has (x, x), the first will be shortened
# Before shortening, we need to pad by shorten factor so as not to leak
# information into the future. To understand why, imagine shorten factor
# of 2 and sequence of length 4, so ABCD. If we shift just by 1, then we
# would have 0ABC, which gets grouped to [0A][BC] on input, which is
# predicting ABCD as targets. The problem is that [0A] has access to A
# and [BC] has access to C -- it will learn to copy it, peek into
# the future. Shifting twice to [00][AB] solves the problem as the first
# "big" symbol becomes all-0 and the rest is shifted enough.
tl.ShiftRight(n_positions=shorten_factor - 1),
tl.Fn(
'Shorten',
lambda x: jnp.reshape( # Shorten -- move to depth.
x, (x.shape[0], x.shape[1] // shorten_factor, -1)),
n_out=1),
tl.Dense(d_model),
tl.Dup(), # Stack has (short_x, short_x, x)
tl.ReversibleSerial(decoder_blocks),
tl.Select([0], n_in=2),
tl.LayerNorm(),
tl.Dropout(rate=dropout, shared_axes=[-2], mode=mode), # pylint: disable=no-value-for-parameter
tl.Dense(shorten_factor * d_embedding),
tl.Fn(
'ProlongBack',
lambda x: jnp.reshape( # Prolong back.
x, (x.shape[0], x.shape[1] * shorten_factor, -1)),
n_out=1),
tl.Concatenate(), # Concatenate with just the embeddings.
tl.CausalConv(d_embedding),
tl.Relu(),
tl.SRU(d_embedding), # One RNN layer for conditional dependence.
tl.Dense(vocab_size),
tl.LogSoftmax())
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),
)
def BERT(
d_model=768,
vocab_size=30522,
max_len=512,
type_vocab_size=2,
n_heads=12,
d_ff=3072,
n_layers=12,
head=None,
init_checkpoint=None,
mode='eval',
):
"""BERT (default hparams are for bert-base-uncased)."""
# TODO(piotrekp1): loading config from model_name
layer_norm_eps = 1e-12
d_head = d_model // n_heads
word_embeddings = tl.Embedding(vocab_size, d_model)
type_embeddings = tl.Embedding(type_vocab_size, d_model)
position_embeddings = tl.PositionalEncoding(max_len, mode=mode)
embeddings = [
tl.Select([0, 1, 0], n_in=3), # Drops 'idx' input.
tl.Parallel(word_embeddings, type_embeddings, [
tl.PaddingMask(),
tl.Fn('Squeeze', lambda x: np.squeeze(x, (1, 2)), n_out=1)
]),
tl.Add(),
position_embeddings,
tl.LayerNorm(epsilon=layer_norm_eps),
]
encoder = []
for _ in range(n_layers):
attn = tl.SelfAttention(n_heads=n_heads,
d_qk=d_head,
d_v=d_head,
bias=True,
masked=True,
mode=mode)
feed_forward = [tl.Dense(d_ff), tl.Gelu(), tl.Dense(d_model)]
encoder += [
tl.Select([0, 1, 1]), # Save a copy of the mask
tl.Residual(attn, AddBias()), # pylint: disable=no-value-for-parameter
tl.LayerNorm(epsilon=layer_norm_eps),
tl.Residual(*feed_forward),
tl.LayerNorm(epsilon=layer_norm_eps),
]
encoder += [tl.Select([0], n_in=2)] # Drop the mask
pooler = [
tl.Fn('', lambda x: (x[:, 0, :], x), n_out=2),
tl.Dense(d_model),
tl.Tanh(),
]
init_checkpoint = init_checkpoint if mode == 'train' else None
bert = PretrainedBERT(embeddings + encoder + pooler,
init_checkpoint=init_checkpoint)
if head is not None:
bert = tl.Serial(bert, head())
return bert
def ConfigurableTransformerLM(vocab_size,
d_model=512,
d_ff=2048,
n_layers=6,
n_heads=8,
max_len=2048,
dropout=0.1,
dropout_shared_axes=None,
mode='train',
ff_activation=tl.Relu,
ff_dropout=0.1,
ff_chunk_size=0,
ff_use_sru=0,
ff_sparsity=0,
ff_sparsity_type='1inN',
loss_sparsity_type='mult',
loss_sparsity=0,
loss_d_lowrank=0,
loss_sparsity_prob=None,
attention_chunk_size=0,
attention_type=tl.CausalAttention,
pos_type=None,
pos_axial_shape=None,
pos_d_axial_embs=None,
pos_start_from_zero_prob=1.0,
pos_max_offset_to_add=0):
"""Returns a Transformer language model.
This model performs autoregressive language modeling:
- input: rank 2 tensor representing a batch of text strings via token IDs
plus padding markers; shape is (batch_size, sequence_length). The tensor
elements are integers in `range(vocab_size)`, and `0` values mark padding
positions.
- output: rank 3 tensor representing a batch of log-probability
distributions for each sequence position over possible token IDs;
shape is (batch_size, sequence_length, `vocab_size`).
This model uses only the decoder part of the overall Transformer.
Args:
vocab_size: Input vocabulary size -- each element of the input tensor should
be an integer in `range(vocab_size)`. These integers typically represent
token IDs from a vocabulary-based tokenizer.
d_model: Final dimension of tensors at most points in the model, including
the initial embedding output.
d_ff: Size of special dense layer in the feed-forward part of each encoder
block.
n_layers: Number of encoder blocks. Each block includes attention, dropout,
residual, feed-forward (`Dense`), and activation layers.
n_heads: Number of attention heads.
max_len: Maximum symbol length for positional encoding.
dropout: Stochastic rate (probability) for dropping an activation value when
applying dropout within an encoder block.
dropout_shared_axes: Tensor axes on which to share a dropout mask. Sharing
along batch and sequence axes (`dropout_shared_axes=(0,1)`) is a useful
way to save memory and apply consistent masks to activation vectors at
different sequence positions.
mode: If `'predict'`, use fast inference. If `'train'`, each encoder block
will include dropout; else, it will pass all values through unaltered.
ff_activation: Type of activation function at the end of each encoder block;
must be an activation-type subclass of `Layer`.
ff_dropout: Stochastic rate (probability) for dropping an activation value
when applying dropout after the FF dense layer.
ff_chunk_size: int; if > 0, chunk feed-forward into this-sized chunks
ff_use_sru: int or pair of ints; if > 0, we use this many SRU layers
in addition to the feed-forward block (second int specifies sru size)
ff_sparsity: int, if > 0 use sparse feed-forward block with this sparsity
ff_sparsity_type: string, if ff_sparsity >0,
use SparseFF if ff_sparsity_type=`'1inN'` and
use BlockSparseFF if ff_sparsity_type=`'Block'`
loss_sparsity_type: string, type of sparsity to used in loss layer. See
SparseDenseWithOptions for options. None if no sparsity should be used.
loss_sparsity: int, the sparsity for loss layer (if used)
loss_d_lowrank: int, the dimensions for intermediate layer (if used)
loss_sparsity_prob: float, the probability for sparse version of loss to be
used. If None, only sparse version is used.
attention_chunk_size: int, if > 0 run attention chunked at this size
attention_type: The attention layer to use for the decoder part.
pos_type: string, the type of positional embeddings to use.
pos_axial_shape: tuple of ints: input shape to use for the axial position
encoding. If unset, axial position encoding is disabled.
pos_d_axial_embs: tuple of ints: depth of position embedding for each axis.
Tuple length must match pos_axial_shape, and values must sum to d_model.
pos_start_from_zero_prob: how often to start from 0 during training,
(if 1.0, we always start from position 0, if less, we randomize).
pos_max_offset_to_add: maximum offset to add to positions during training
when randomizing; this offset plus input length must still be less than
max_len for all training examples.
Returns:
A Transformer language model as a layer that maps from a tensor of tokens
to activations over a vocab set.
"""
positional_encoder = [
tl.Embedding(vocab_size, d_model),
tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode),
PositionalEncoder(mode, dropout, max_len, pos_type, pos_axial_shape,
pos_d_axial_embs, pos_start_from_zero_prob,
pos_max_offset_to_add)
]
# pylint: disable=g-complex-comprehension
decoder_blocks = [
DecoderBlock(d_model, d_ff, n_heads, dropout, dropout_shared_axes,
mode, ff_activation, ff_dropout, ff_chunk_size,
ff_use_sru, ff_sparsity, ff_sparsity_type,
attention_chunk_size, attention_type)
for i in range(n_layers)
]
# pylint: enable=g-complex-comprehension
# Assemble and return the model.
return tl.Serial( # tokens (or chunked tuple of tokens)
tl.ShiftRight(mode=mode), # toks
positional_encoder, # vecs
decoder_blocks, # vecs
tl.LayerNorm(), # vecs
tl.SparseDenseWithOptions( # vecs
vocab_size,
d_input=d_model,
sparsity_type=loss_sparsity_type,
sparsity=loss_sparsity,
d_lowrank=loss_d_lowrank,
prob_sparse=loss_sparsity_prob,
mode=mode),
)
def Embedder(vocab_size): # tokens --> vectors
return [
tl.Embedding(vocab_size, d_model),
tl.Dropout(rate=dropout, shared_axes=[-2], mode=mode),
]
def FunnelTransformer(vocab_size,
d_model=512,
d_ff=2048,
encoder_segment_lengths=(2, 2, 2),
n_decoder_blocks=2,
n_heads=8,
max_len=2048,
dropout=0.1,
dropout_shared_axes=None,
mode='train',
ff_activation=tl.Relu,
pool_layer=tl.AvgPool,
pool_size=(2, ),
separate_cls=True):
"""Returns a Full Funnel Transformer, that can be used for example for BERT.
This model outputs token-level categorical distributions over all vocab:
- input: rank 2 tensor representing a batch of text strings via token IDs
plus padding markers; shape is (batch_size, sequence_length). The tensor
elements are integers in `range(vocab_size)`, and `0` values mark padding
positions.
- output: rank 3 tensor representing a batch of log-probability
distributions over `vocab_size` categories for each token; shape is
(batch_size, sequence_length, vocab_size).
Args:
vocab_size: Input vocabulary size -- each element of the input tensor
should be an integer in `range(vocab_size)`. These integers typically
represent token IDs from a vocabulary-based tokenizer.
d_model: Final dimension of tensors at most points in the model, including
the initial embedding output.
d_ff: Size of special dense layer in the feed-forward part of each encoder
block.
encoder_segment_lengths: Tuple, where each element denotes the number of
transformer encoder blocks preceding a funnel transformer block.
There is no funnel block after the last sequence of encoder blocks,
therefore the total number of blocks in the model is equal to
`sum(encoder_segment_lengths) + len(encoder_segment_lengths) - 1`.
n_decoder_blocks: Number of transformer blocks in the upsampling decoder.
n_heads: Number of attention heads.
max_len: Maximum symbol length for positional encoding.
dropout: Stochastic rate (probability) for dropping an activation value
when applying dropout within an encoder block.
dropout_shared_axes: Tensor axes on which to share a dropout mask.
Sharing along batch and sequence axes (`dropout_shared_axes=(0,1)`) is
a useful way to save memory and apply consistent masks to activation
vectors at different sequence positions.
mode: If `'train'`, each encoder block will include dropout; else, it will
pass all values through unaltered.
ff_activation: Type of activation function at the end of each encoder
block; must be an activation-type subclass of `Layer`.
pool_layer: Type of pooling layer used for downsampling in each of the
funnel blocks; should be `tl.AvgPool` or `tl.MaxPool`.
pool_size: Shape of window that gets reduced to a single vector value.
If the layer inputs are :math:`n`-dimensional arrays, then `pool_size`
must be a tuple of length :math:`n-2`.
separate_cls: If `True`, pooling in funnel blocks is not applied to
embeddings of the first token (`cls` from BERT paper) and only final
embedding of this token is used for categorization - the rest are
discarded. If `False`, each token from the beginning is pooled and
all embeddings are averaged and mapped to output categories like in
original `TransformerEncoder` model.
"""
assert encoder_segment_lengths
positional_encoder = [
tl.Embedding(vocab_size, d_model),
tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode),
tl.PositionalEncoding(max_len=max_len)
]
n_encoder_segments = len(encoder_segment_lengths)
encoder_blocks_before_first_pooling = [
_EncoderBlock(d_model, d_ff, n_heads, dropout, dropout_shared_axes,
mode, ff_activation)
for _ in range(encoder_segment_lengths[0])
]
encoder_blocks_from_first_pooling = []
for i in range(1, n_encoder_segments):
# Building i'th segment
# Add funnel block between segments
encoder_blocks_from_first_pooling.append(
_FunnelBlock(d_model,
d_ff,
n_heads,
dropout,
dropout_shared_axes,
mode,
ff_activation,
pool_layer,
pool_size=pool_size,
strides=pool_size,
separate_cls=separate_cls))
for _ in range(encoder_segment_lengths[i]):
# Create segment_size encoder blocks
encoder_blocks_from_first_pooling.append(
_EncoderBlock(d_model, d_ff, n_heads, dropout,
dropout_shared_axes, mode, ff_activation))
decoder_blocks = [
_EncoderBlock(d_model, d_ff, n_heads, dropout, dropout_shared_axes,
mode, ff_activation) for _ in range(n_decoder_blocks)
]
total_pool_size = pool_size[0]**(len(encoder_segment_lengths) - 1)
# Assemble and return the model.
return tl.Serial( # toks
tl.Branch(positional_encoder, tl.PaddingMask()), # vecs masks
encoder_blocks_before_first_pooling, # vecs masks
tl.Select([0, 1, 0, 1]),
# vecs masks residual = vecs old_masks
encoder_blocks_from_first_pooling, # vecs masks residual masks
tl.Select([0, 2, 3]), # vecs residual masks
tl.Parallel(
# residual from first segment is taken before
# normalization, so apply it now
None,
tl.LayerNorm(),
None), # vecs norm(residual) masks
_Upsampler(total_pool_size, separate_cls), # vecs masks
decoder_blocks,
tl.Select([0], n_in=2), # vecs
tl.LayerNorm(),
tl.Dense(vocab_size),
)
def FunnelTransformerLM(vocab_size,
d_model=512,
d_ff=2048,
vanilla_layers=(0, 1),
shorten_factors=(3, ),
n_funnel_blocks=(6, ),
n_heads=8,
dropout=0.1,
dropout_shared_axes=None,
mode='train',
ff_activation=tl.FastGelu):
"""Returns a Transformer language model.
This model performs autoregressive language modeling:
- input: rank 2 tensor representing a batch of text strings via token IDs
plus padding markers; shape is (batch_size, sequence_length). The tensor
elements are integers in `range(vocab_size)`, and `0` values mark padding
positions.
- output: rank 3 tensor representing a batch of log-probability
distributions for each sequence position over possible token IDs;
shape is (batch_size, sequence_length, `vocab_size`).
This model uses only the decoder part of the overall Transformer.
Args:
vocab_size: Input vocabulary size -- each element of the input tensor
should be an integer in `range(vocab_size)`. These integers typically
represent token IDs from a vocabulary-based tokenizer.
d_model: Final dimension of tensors at most points in the model, including
the initial embedding output.
d_ff: Size of special dense layer in the feed-forward part of each encoder
block.
vanilla_layers: (pre_layers, post_layers) tuple - number of full token-level
Transformer decoder layers before and after shortening.
shorten_factors: by how much to shorten at each step - tuple of arbitrary
length denoting by how much shorten at each pooling stage.
n_funnel_blocks: number of Transformer decoder blocks after each stage of
pooling - tuple of the same length as `shorten_factors`.
n_heads: Number of attention heads.
dropout: Stochastic rate (probability) for dropping an activation value
when applying dropout within an encoder block.
dropout_shared_axes: Tensor axes on which to share a dropout mask.
Sharing along batch and sequence axes (`dropout_shared_axes=(0,1)`) is
a useful way to save memory and apply consistent masks to activation
vectors at different sequence positions.
mode: str: 'train' or 'eval'.
ff_activation: Type of activation function at the end of each encoder
block; must be an activation-type subclass of `Layer`.
Returns:
A Transformer language model as a layer that maps from a tensor of tokens
to activations over a vocab set.
"""
assert mode != 'predict' # For now, 'predict' mode is unsupported.
assert len(n_funnel_blocks) == len(shorten_factors)
token_encoder = [
tl.Embedding(vocab_size, d_model),
tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode)
]
context_bias_layer, location_bias_layer = _get_rel_att_inputs(
d_model, n_heads)
n_pre_decoder_blocks, n_post_decoder_blocks = vanilla_layers
def create_decoder_blocks(n_layers, total_pooling): # pylint: disable=invalid-name
decoder_blocks = [
# pylint: disable=g-complex-comprehension
_RelativeDecoderBlock(d_model, d_ff, n_heads, dropout,
dropout_shared_axes, mode, ff_activation,
context_bias_layer, location_bias_layer,
total_pooling) for _ in range(n_layers)
]
return decoder_blocks + [tl.LayerNorm()]
total_pooling_acc = 1
pre_decoder_blocks = create_decoder_blocks(n_pre_decoder_blocks,
total_pooling=1)
funnel_blocks = []
for shorten_factor, block_len in zip(shorten_factors, n_funnel_blocks):
funnel_blocks = funnel_blocks + [
_FunnelRelativeDecoderBlock(
d_model,
d_ff,
n_heads,
dropout,
dropout_shared_axes,
mode,
ff_activation,
context_bias_layer=context_bias_layer,
location_bias_layer=location_bias_layer,
total_pooling=total_pooling_acc,
shorten_factor=shorten_factor,
resampler_fn=_DownsamplerLM)
]
total_pooling_acc *= shorten_factor
funnel_blocks = funnel_blocks + create_decoder_blocks(
block_len, total_pooling_acc)
upsampling_layer = _FunnelRelativeDecoderBlock(
d_model,
d_ff,
n_heads,
dropout,
dropout_shared_axes,
mode,
ff_activation,
context_bias_layer=context_bias_layer,
location_bias_layer=location_bias_layer,
total_pooling=total_pooling_acc,
shorten_factor=total_pooling_acc,
resampler_fn=_UpsamplerLM)
conv_layer = tl.Serial(tl.CausalConv(d_model, total_pooling_acc),
ff_activation())
post_decoder_blocks = create_decoder_blocks(n_post_decoder_blocks,
total_pooling=1)
# Assemble and return the model.
return tl.Serial( # tokens (or chunked tuple of tokens)
tl.ShiftRight(mode=mode), # toks
token_encoder, # vecs
pre_decoder_blocks, # vecs
tl.Dup(),
tl.ShiftRight(n_positions=total_pooling_acc - 1),
funnel_blocks,
tl.Dropout(rate=dropout, shared_axes=[-2], mode=mode),
upsampling_layer,
tl.LayerNorm(),
tl.Concatenate(),
conv_layer,
post_decoder_blocks,
tl.Dense(vocab_size), # vecs
)
def PositionalEncoder(vocab_size): # tokens --> vectors
return [
tl.Embedding(vocab_size, d_model),
tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode),
tl.PositionalEncoding(max_len=max_len),
]
def FunnelTransformerEncoder(vocab_size,
n_classes=10,
d_model=512,
d_ff=2048,
encoder_segment_lengths=(2, 2, 2),
n_heads=8,
max_len=2048,
dropout=0.1,
dropout_shared_axes=None,
mode='train',
ff_activation=tl.Relu,
pool_layer=tl.AvgPool,
pool_size=(2, ),
strides=(2, ),
separate_cls=True):
"""Returns a Funnel Encoder.
This model performs text categorization:
- input: rank 2 tensor representing a batch of text strings via token IDs
plus padding markers; shape is (batch_size, sequence_length). The tensor
elements are integers in `range(vocab_size)`, and `0` values mark padding
positions.
- output: rank 2 tensor representing a batch of log-probability
distributions over N categories; shape is (batch_size, `n_classes`).
Args:
vocab_size: Input vocabulary size -- each element of the input tensor
should be an integer in `range(vocab_size)`. These integers typically
represent token IDs from a vocabulary-based tokenizer.
n_classes: Final dimension of the output tensors, representing N-way
classification.
d_model: Final dimension of tensors at most points in the model, including
the initial embedding output.
d_ff: Size of special dense layer in the feed-forward part of each encoder
block.
encoder_segment_lengths: Tuple, where each element denotes the number of
transformer encoder blocks preceding a funnel transformer block.
There is no funnel block after the last sequence of encoder blocks,
therefore the total number of blocks in the model is equal to
`sum(encoder_segment_lengths) + len(encoder_segment_lengths) - 1`.
n_heads: Number of attention heads.
max_len: Maximum symbol length for positional encoding.
dropout: Stochastic rate (probability) for dropping an activation value
when applying dropout within an encoder block.
dropout_shared_axes: Tensor axes on which to share a dropout mask.
Sharing along batch and sequence axes (`dropout_shared_axes=(0,1)`) is
a useful way to save memory and apply consistent masks to activation
vectors at different sequence positions.
mode: If `'train'`, each encoder block will include dropout; else, it will
pass all values through unaltered.
ff_activation: Type of activation function at the end of each encoder
block; must be an activation-type subclass of `Layer`.
pool_layer: Type of pooling layer used for downsampling in each of the
funnel blocks; should be `tl.AvgPool` or `tl.MaxPool`.
pool_size: Shape of window that gets reduced to a single vector value.
If the layer inputs are :math:`n`-dimensional arrays, then `pool_size`
must be a tuple of length :math:`n-2`.
strides: Offsets from the location of one window to the locations of
neighboring windows along each axis. If specified, must be a tuple of
the same length as `pool_size`. If None, then offsets of 1 along each
window axis, :math:`(1, ..., 1)`, will be used.
separate_cls: If `True`, pooling in funnel blocks is not applied to
embeddings of the first token (`cls` from BERT paper) and only final
embedding of this token is used for categorization - the rest are
discarded. If `False`, each token from the beginning is pooled and
all embeddings are averaged and mapped to output categories like in
original `TransformerEncoder` model.
Returns:
A Transformer model that maps strings (conveyed via token IDs) to
probability-like activations over a range of output classes.
"""
assert encoder_segment_lengths
positional_encoder = [
tl.Embedding(vocab_size, d_model),
tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode),
tl.PositionalEncoding(max_len=max_len)
]
encoder_blocks = []
n_encoder_segments = len(encoder_segment_lengths)
for i in range(n_encoder_segments):
# Building i'th segment
for _ in range(encoder_segment_lengths[i]):
# Create segment_size encoder blocks
encoder_blocks.append(
_EncoderBlock(d_model, d_ff, n_heads, dropout,
dropout_shared_axes, mode, ff_activation))
# If not last segment, add funnel block
if i != n_encoder_segments - 1:
encoder_blocks.append(
_FunnelBlock(d_model, d_ff, n_heads, dropout,
dropout_shared_axes, mode, ff_activation,
pool_layer, pool_size, strides, separate_cls))
cls_pooling = SelectFirst() if separate_cls else tl.Mean(axis=1)
# Assemble and return the model.
return tl.Serial( # toks
# Encode.
tl.Branch(positional_encoder, tl.PaddingMask()), # vecs masks
encoder_blocks, # vecs masks
tl.Select([0], n_in=2), # vecs
tl.LayerNorm(), # vecs
# Map to output categories.
cls_pooling, # cls
tl.Dense(n_classes), # cls
)
def TransformerEncoder(vocab_size,
n_classes=10,
d_model=512,
d_ff=2048,
n_layers=6,
n_heads=8,
dropout=0.1,
dropout_shared_axes=None,
max_len=2048,
mode='train',
ff_activation=tl.Relu):
"""Returns a Transformer-style encoder.
For each item in a batch, this model performs a sequence-to-sequence mapping:
- input: sequence of integers, usually token id's from a fixed-size
vocabulary -- integers in `range(M)`, where `M` is the vocabulary
size.
- output: same-length sequence of N-dimensional vectors, where each vector
can be interpreted as a log-probability distribution over N discrete
categories.
Args:
vocab_size: "Vocabulary size" -- input integer id's must be in
`range(vocab_size)`. Id's typically come from preprocessing text data
with a vocabulary-based tokenizer.
n_classes: Size/depth of the output vectors, intended for an N-way
classification task.
d_model: The basic embedding size (vector depth) of the model. This is the
vector size used by the initial embedding layer and at many intermediate
points in the model.
d_ff: Vector depth (typically greater than `d_model`) used in the
feed-forward (`Dense`) layer of each encoder block.
n_layers: Number of encoder blocks. Each encoder block includes attention,
dropout, residual, feed-forward (`Dense`), and activation layers.
n_heads: Number of attention heads.
dropout: Stochastic rate (probability) for dropping an activation value
when applying dropout within an encoder block.
dropout_shared_axes: Tensor axes on which to share a dropout mask.
max_len: Maximum symbol length for positional encoding.
mode: If `'train'`, each encoder block will include dropout; else, it will
pass all values through unaltered.
ff_activation: The activation function (layer) at the end of each encoder
block.
Returns:
A Transformer model as a layer that maps from token id's to activations
over a set of output classes.
"""
positional_encoder = [
tl.Embedding(vocab_size, d_model),
tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode),
tl.PositionalEncoding(max_len=max_len)]
encoder_blocks = [
_EncoderBlock(d_model, d_ff, n_heads, dropout, dropout_shared_axes,
mode, ff_activation)
for i in range(n_layers)]
# Assemble and return the model.
return tl.Serial( # toks
# Encode.
tl.Branch(
positional_encoder, tl.PaddingMask()), # vecs masks
encoder_blocks, # vecs masks
tl.Select([0], n_in=2), # vecs
tl.LayerNorm(), # vecs
# Map to output categories.
tl.Mean(axis=1), # vecs
tl.Dense(n_classes), # vecs
tl.LogSoftmax(), # vecs
)
def Transformer(input_vocab_size,
output_vocab_size=None,
d_model=512,
d_ff=2048,
n_encoder_layers=6,
n_decoder_layers=6,
n_heads=8,
dropout=0.1,
max_len=2048,
mode='train',
ff_activation=tl.Relu):
"""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_encoder_layers: int: number of encoder layers
n_decoder_layers: int: number of 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'
ff_activation: the non-linearity in feed-forward layer
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,
ff_activation) for i in range(n_encoder_layers)
])
encoder_decoder_stack = ( # vecs_d masks vecs_e --> vecs_d masks vecs_e
[
EncoderDecoder(d_model, d_ff, n_heads, dropout, i, mode,
ff_activation) for i in range(n_decoder_layers)
])
# Input: encoder_side_tokens, decoder_side_tokens
return tl.Serial( # tokens_e tokens_d
tl.Parallel([], tl.Dup()), # toks_e toks_d toks_d (for loss)
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
)
def TransformerEncoder(vocab_size,
n_classes=10,
d_model=D_MODEL,
d_ff=D_FF,
n_layers=N_LAYERS,
n_heads=N_HEADS,
max_len=MAX_SEQUENCE_LENGTH,
dropout=DROPOUT_RATE,
dropout_shared_axes=DROPOUT_SHARED_AXES,
mode=MODE,
ff_activation=FF_ACTIVATION_TYPE):
"""Returns a Transformer encoder suitable for N-way classification.
This model maps tokenized text to N-way (``n_classes``) activations:
- input: Array representing a batch of text strings via token IDs plus
padding markers; shape is (batch_size, sequence_length), where
sequence_length <= ``max_len``. Array elements are integers in
``range(vocab_size)``, and 0 values mark padding positions.
- output: Array representing a batch of raw (non-normalized) activations
over ``n_classes`` categories; shape is (batch_size, ``n_classes``).
Args:
vocab_size: Input vocabulary size -- each element of the input array
should be an integer in ``range(vocab_size)``. These integers typically
represent token IDs from a vocabulary-based tokenizer.
n_classes: Last/innermost dimension of output arrays, suitable for N-way
classification.
d_model: Last/innermost dimension of activation arrays at most points in
the model, including the initial embedding output.
d_ff: Last/innermost dimension of special (typically wider)
:py:class:`Dense` layer in the feedforward part of each encoder block.
n_layers: Number of encoder blocks. Each block includes attention, dropout,
residual, layer-norm, feedforward (:py:class:`Dense`), and activation
layers.
n_heads: Number of attention heads.
max_len: Maximum symbol length for positional encoding.
dropout: Stochastic rate (probability) for dropping an activation value
when applying dropout within encoder blocks. The same rate is also
used for attention dropout in encoder blocks.
dropout_shared_axes: Tensor axes on which to share a dropout mask.
Sharing along batch and sequence axes (``dropout_shared_axes=(0,1)``)
is a useful way to save memory and apply consistent masks to activation
vectors at different sequence positions.
mode: If ``'train'``, each encoder block will include dropout; else, it
will pass all values through unaltered.
ff_activation: Type of activation function at the end of each encoder
block; must be an activation-type subclass of :py:class:`Layer`.
Returns:
A Transformer model that maps strings (conveyed by token IDs) to
raw (non-normalized) activations over a range of output classes.
"""
def _Dropout():
return tl.Dropout(rate=dropout,
shared_axes=dropout_shared_axes,
mode=mode)
def _EncBlock():
return _EncoderBlock(d_model, d_ff, n_heads, dropout,
dropout_shared_axes, mode, ff_activation)
return tl.Serial(
tl.Branch([],
tl.PaddingMask()), # Creates masks from copy of the tokens.
tl.Embedding(vocab_size, d_model),
_Dropout(),
tl.PositionalEncoding(max_len=max_len),
[_EncBlock() for _ in range(n_layers)],
tl.Select([0], n_in=2), # Drops the masks.
tl.LayerNorm(),
tl.Mean(axis=1),
tl.Dense(n_classes),
)
def TransformerEncoder(vocab_size=vocab_size,
n_classes=10,
d_model=512,
d_ff=2048,
n_layers=6,
n_heads=8,
dropout=0.1,
dropout_shared_axes=None,
max_len=2048,
mode='train',
ff_activation=tl.Relu,
EncoderBlock=EncoderBlock):
"""
Returns a Transformer encoder model.
The input to the model is a tensor of tokens.
Args:
vocab_size (int): vocab size. Defaults to vocab_size.
n_classes (int): how many classes on output. Defaults to 10.
d_model (int): depth of embedding. Defaults to 512.
d_ff (int): depth of feed-forward layer. Defaults to 2048.
n_layers (int): number of encoder/decoder layers. Defaults to 6.
n_heads (int): number of attention heads. Defaults to 8.
dropout (float): dropout rate (how much to drop out). Defaults to 0.1.
dropout_shared_axes (int): axes on which to share dropout mask. Defaults to None.
max_len (int): maximum symbol length for positional encoding. Defaults to 2048.
mode (str): 'train' or 'eval'. Defaults to 'train'.
ff_activation (function): the non-linearity in feed-forward layer. Defaults to tl.Relu.
EncoderBlock (function): Returns the encoder block. Defaults to EncoderBlock.
Returns:
trax.layers.combinators.Serial: A Transformer model as a layer that maps
from a tensor of tokens to activations over a set of output classes.
"""
positional_encoder = [
tl.Embedding(vocab_size, d_model),
tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode),
tl.PositionalEncoding(max_len=max_len)
]
### START CODE HERE (REPLACE INSTANCES OF 'None' WITH YOUR CODE) ###
# Use the function `EncoderBlock` (implemented above) and pass in the parameters over `n_layers`
encoder_blocks = [EncoderBlock(d_model, d_ff, n_heads, dropout, dropout_shared_axes, mode, ff_activation) for _ in range(n_layers)]
# Assemble and return the model.
return tl.Serial(
# Encode
tl.Branch(
# Use `positional_encoder`
positional_encoder,
# Use trax padding mask
tl.PaddingMask(),
),
# Use `encoder_blocks`
encoder_blocks,
# Use select layer
tl.Select([0], n_in=2),
# Use trax layer normalization
tl.LayerNorm(),
# Map to output categories.
# Use trax mean. set axis to 1
tl.Mean(axis=1),
# Use trax Dense using `n_classes`
tl.Dense(n_classes),
# Use trax log softmax
tl.LogSoftmax(),
)
def RelformerLM(vocab_size,
d_model=512,
d_ff=2048,
vanilla_layers=(1, 1),
shorten_factor=3,
n_rel_layers=6,
n_heads=8,
dropout=0.1,
dropout_shared_axes=None,
vanilla_attn_type=tl.LSHSelfAttention,
pos_type='fixed-base',
max_len=3072,
n_raw_tokens_generated=1,
mode='train',
ff_activation=tl.FastGelu):
"""Returns a Transformer language model.
This model performs autoregressive language modeling:
- input: rank 2 tensor representing a batch of text strings via token IDs
plus padding markers; shape is (batch_size, sequence_length). The tensor
elements are integers in `range(vocab_size)`, and `0` values mark padding
positions.
- output: rank 3 tensor representing a batch of log-probability
distributions for each sequence position over possible token IDs;
shape is (batch_size, sequence_length, `vocab_size`).
This model uses only the decoder part of the overall Transformer.
Args:
vocab_size: Input vocabulary size -- each element of the input tensor
should be an integer in `range(vocab_size)`. These integers typically
represent token IDs from a vocabulary-based tokenizer.
d_model: Final dimension of tensors at most points in the model, including
the initial embedding output.
d_ff: Size of special dense layer in the feed-forward part of each encoder
block.
vanilla_layers: (pre_layers, post_layers) tuple - number of full token-level
Transformer decoder layers before and after shortening.
shorten_factor: by how much to shorten
n_rel_layers: number of Transformer blocks after the pooling. These blocks
use relative attention.
n_heads: Number of attention heads.
dropout: Stochastic rate (probability) for dropping an activation value
when applying dropout within an encoder block.
dropout_shared_axes: Tensor axes on which to share a dropout mask.
Sharing along batch and sequence axes (`dropout_shared_axes=(0,1)`) is
a useful way to save memory and apply consistent masks to activation
vectors at different sequence positions.
vanilla_attn_type: class: attention class such as SelfAttention to use in
the layers before and after shortening (vanilla layers).
pos_type: string, the type of positional embeddings to use.
max_len: int: maximum symbol length both for positional encoding and it is
also the maximum length of the possible inference in 'predict' mode
n_raw_tokens_generated: int: number of tokens generated with every pass
through model in 'predict' mode. Number of tokens should be smaller and
divisible by the first shorten factor we are using in the model.
It cannot be larger than one if we use vanilla layers because we would
lose autoregressive property of the model.
mode: str: 'train' or 'eval' or 'predict'.
ff_activation: Type of activation function at the end of each encoder
block; must be an activation-type subclass of `Layer`.
Returns:
A Transformer language model as a layer that maps from a tensor of tokens
to activations over a vocab set.
"""
token_encoder = [
tl.Embedding(vocab_size, d_model),
tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode)
]
positional_encoder = PositionalEncoder(mode, dropout, max_len, pos_type)
n_pre_decoder_blocks, n_post_decoder_blocks = vanilla_layers
def create_decoder_blocks(n_layers, total_pooling): # pylint: disable=invalid-name
context_bias_layer, location_bias_layer = _get_rel_att_inputs(
d_model, n_heads)
decoder_blocks = [
# pylint: disable=g-complex-comprehension
_RelativeDecoderBlock(d_model, d_ff, n_heads, dropout,
dropout_shared_axes, mode, ff_activation,
context_bias_layer, location_bias_layer,
total_pooling, max_len)
for _ in range(n_layers)
]
return decoder_blocks + [tl.LayerNorm()]
def create_reformer_blocks(n_layers, dense=True): # pylint: disable=invalid-name
if n_layers == 0:
return [tl.LayerNorm()]
d_per_head = d_model // n_heads
decoder_blocks = [
DecoderBlock(
d_model,
d_ff,
d_per_head,
d_per_head,
n_heads, # pylint: disable=g-complex-comprehension
vanilla_attn_type,
dropout,
ff_activation,
dropout,
ff_use_sru=0,
ff_chunk_size=0,
ff_sparsity=0,
attention_chunk_size=0,
mode=mode) for _ in range(n_layers)
]
return [
tl.Dup(),
tl.ReversibleSerial(decoder_blocks),
tl.Concatenate(),
tl.LayerNorm(),
tl.Dense(d_model) if dense else [],
]
pre_decoder_blocks = create_reformer_blocks(n_pre_decoder_blocks,
dense=True)
relative_decoder_blocks = create_decoder_blocks(n_rel_layers,
shorten_factor)
conv_layer = tl.Serial(tl.CausalConv(d_model, shorten_factor),
ff_activation())
post_decoder_blocks = create_reformer_blocks(n_post_decoder_blocks,
dense=False)
cacher = RelformerCacher(total_kv_pooling=shorten_factor,
n_raw_tokens_generated=n_raw_tokens_generated,
max_inference_length=max_len,
shift=shorten_factor - 1,
mode=mode)
picker = RelformerPicker(total_kv_pooling=shorten_factor,
n_raw_tokens_generated=n_raw_tokens_generated,
mode=mode)
cacher_conv = RelformerCacher(
total_kv_pooling=shorten_factor,
n_raw_tokens_generated=n_raw_tokens_generated,
max_inference_length=max_len,
shift=shorten_factor - 1,
sliding=True,
mode=mode)
picker_conv = PickLastTokenInPredict(mode=mode)
# Assemble and return the model.
return tl.Serial( # tokens (or chunked tuple of tokens)
tl.ShiftRight(mode=mode), # toks
token_encoder, # vecs
positional_encoder,
pre_decoder_blocks, # vecs
tl.Dup(),
cacher,
tl.ShiftRight(n_positions=shorten_factor - 1, mode=mode),
_DownsamplerLM(shorten_factor, d_model),
relative_decoder_blocks,
tl.Dropout(rate=dropout, shared_axes=[-2], mode=mode),
_UpsamplerLM(shorten_factor, d_model),
tl.LayerNorm(),
picker,
tl.Concatenate(),
cacher_conv,
conv_layer,
picker_conv,
post_decoder_blocks,
tl.Dense(vocab_size), # vecs
)
def ConfigurableTransformerEncoder(vocab_size,
n_classes=10,
d_model=512,
d_ff=2048,
n_layers=6,
n_heads=8,
max_len=2048,
dropout=0.1,
dropout_shared_axes=None,
mode='train',
ff_activation=tl.Relu,
ff_dropout=0.1,
ff_chunk_size=0,
ff_use_sru=0,
ff_sparsity=0,
ff_sparsity_type='1inN',
attention_chunk_size=0,
attention_type=tl.Attention,
pos_type=None,
pos_axial_shape=None,
pos_d_axial_embs=None):
"""Returns a Transformer encoder merged with an N-way categorization head.
This model performs text categorization:
- input: rank 2 tensor representing a batch of text strings via token IDs
plus padding markers; shape is (batch_size, sequence_length). The tensor
elements are integers in `range(vocab_size)`, and `0` values mark padding
positions.
- output: rank 2 tensor representing a batch of log-probability
distributions over N categories; shape is (batch_size, `n_classes`).
Args:
vocab_size: Input vocabulary size -- each element of the input tensor should
be an integer in `range(vocab_size)`. These integers typically represent
token IDs from a vocabulary-based tokenizer.
n_classes: Final dimension of the output tensors, representing N-way
classification.
d_model: Final dimension of tensors at most points in the model, including
the initial embedding output.
d_ff: Size of special dense layer in the feed-forward part of each encoder
block.
n_layers: Number of encoder blocks. Each block includes attention, dropout,
residual, feed-forward (`Dense`), and activation layers.
n_heads: Number of attention heads.
max_len: Maximum symbol length for positional encoding.
dropout: Stochastic rate (probability) for dropping an activation value when
applying dropout within an encoder block.
dropout_shared_axes: Tensor axes on which to share a dropout mask. Sharing
along batch and sequence axes (`dropout_shared_axes=(0,1)`) is a useful
way to save memory and apply consistent masks to activation vectors at
different sequence positions.
mode: If `'train'`, each encoder block will include dropout; else, it will
pass all values through unaltered.
ff_activation: Type of activation function at the end of each encoder block;
must be an activation-type subclass of `Layer`.
ff_dropout: Stochastic rate (probability) for dropping an activation value
when applying dropout after the FF dense layer.
ff_chunk_size: int; if > 0, chunk feed-forward into this-sized chunks
ff_use_sru: int or pair of ints; if > 0, we use this many SRU layers
in addition to the feed-forward block (second int specifies sru size)
ff_sparsity: int, if > 0 use sparse feed-forward block with this sparsity
ff_sparsity_type: string, if ff_sparsity >0,
use SparseFF if ff_sparsity_type=`'1inN'` and
use BlockSparseFF if ff_sparsity_type=`'Block'`
attention_chunk_size: int, if > 0 run attention chunked at this size
attention_type: The attention layer to use for the encoder part.
pos_type: string, the type of positional embeddings to use.
pos_axial_shape: tuple of ints: input shape to use for the axial position
encoding. If unset, axial position encoding is disabled.
pos_d_axial_embs: tuple of ints: depth of position embedding for each axis.
Tuple length must match pos_axial_shape, and values must sum to d_model.
Returns:
A Transformer model that maps strings (conveyed via token IDs) to
probability-like activations over a range of output classes.
"""
positional_encoder = [
tl.Embedding(vocab_size, d_model),
tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode),
PositionalEncoder(mode, dropout, max_len, pos_type, pos_axial_shape,
pos_d_axial_embs)
]
positional_encoder = tl.AssertFunction('...->...d', positional_encoder)
# pylint: disable=g-complex-comprehension
encoder_blocks = [
EncoderBlock(d_model, d_ff, n_heads, dropout, dropout_shared_axes,
mode, ff_activation, ff_dropout, ff_chunk_size,
ff_use_sru, ff_sparsity, ff_sparsity_type,
attention_chunk_size, attention_type)
for i in range(n_layers)
]
# pylint: enable=g-complex-comprehension
# Assemble and return the model.
return tl.Serial( # toks
# Encode.
tl.Branch(positional_encoder, tl.PaddingMask()), # vecs masks
encoder_blocks, # vecs masks
tl.Select([0], n_in=2), # vecs
tl.LayerNorm(), # vecs
# Map to output categories.
tl.Mean(axis=1), # vecs
tl.Dense(n_classes), # vecs
)
def TransformerDecoder(vocab_size=None,
d_model=512,
d_ff=2048,
n_layers=6,
n_heads=8,
d_attention_key=None,
d_attention_value=None,
attention_type=tl.DotProductCausalAttention,
dropout=0.1,
share_qk=False,
max_len=2048,
mode='train',
ff_activation=tl.Relu):
"""Returns a Transformer decoder model.
The input to the model is either continuous or discrete - controlled by
vocab_size. Does not shift the input to the right, i.e. the output for
timestep t is based on inputs up to timestep t inclusively.
Args:
vocab_size: int or None: vocab size if running on discrete input, None
otherwise.
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
d_attention_key: int: depth of key vector for each attention head (default
is d_model // n_heads)
d_attention_value: int: depth of value vector for each attention head
(default is d_model // n_heads)
attention_type: subclass of tl.BaseCausalAttention: attention class to use
dropout: float: dropout rate (how much to drop out)
share_qk: bool, whether to share queries and keys in decoder attention
max_len: int: maximum symbol length for positional encoding
mode: str: 'train' or 'eval'
ff_activation: the non-linearity in feed-forward layer
Returns:
A Transformer decoder as a layer that maps from a continuous or discrete
tensor to a continuous tensor.
"""
def DecoderBlocks(n_blocks): # vectors --> vectors
return [ # pylint: disable=g-complex-comprehension
_DecoderBlock(d_model, d_ff, n_heads, d_attention_key,
d_attention_value, attention_type, dropout, share_qk,
i, mode, ff_activation) for i in range(n_blocks)
]
embedding_or_dense = (tl.Embedding(d_model, vocab_size)
if vocab_size is not None else tl.Dense(d_model))
dropout_ = tl.Dropout(rate=dropout, mode=mode)
positional_encoding = tl.PositionalEncoding(max_len=max_len)
# Assemble and return the model.
return tl.Serial( # toks
embedding_or_dense, # vecs
dropout_, # vecs
positional_encoding, # vecs
DecoderBlocks(n_layers), # vecs
tl.LayerNorm(), # vecs
)
def ConfigurableTransformerLM(vocab_size,
d_model=512,
d_ff=2048,
n_layers=6,
n_heads=8,
max_len=2048,
dropout=0.1,
dropout_shared_axes=None,
mode='train',
ff_activation=tl.Relu,
ff_dropout=0.1,
ff_chunk_size=0,
ff_use_sru=0,
ff_sparsity=0,
ff_sparsity_type='1inN',
attention_chunk_size=0,
attention_type=tl.CausalAttention):
"""Returns a Transformer language model.
This model performs autoregressive language modeling:
- input: rank 2 tensor representing a batch of text strings via token IDs
plus padding markers; shape is (batch_size, sequence_length). The tensor
elements are integers in `range(vocab_size)`, and `0` values mark padding
positions.
- output: rank 3 tensor representing a batch of log-probability
distributions for each sequence position over possible token IDs;
shape is (batch_size, sequence_length, `vocab_size`).
This model uses only the decoder part of the overall Transformer.
Args:
vocab_size: Input vocabulary size -- each element of the input tensor should
be an integer in `range(vocab_size)`. These integers typically represent
token IDs from a vocabulary-based tokenizer.
d_model: Final dimension of tensors at most points in the model, including
the initial embedding output.
d_ff: Size of special dense layer in the feed-forward part of each encoder
block.
n_layers: Number of encoder blocks. Each block includes attention, dropout,
residual, feed-forward (`Dense`), and activation layers.
n_heads: Number of attention heads.
max_len: Maximum symbol length for positional encoding.
dropout: Stochastic rate (probability) for dropping an activation value when
applying dropout within an encoder block.
dropout_shared_axes: Tensor axes on which to share a dropout mask. Sharing
along batch and sequence axes (`dropout_shared_axes=(0,1)`) is a useful
way to save memory and apply consistent masks to activation vectors at
different sequence positions.
mode: If `'predict'`, use fast inference. If `'train'`, each encoder block
will include dropout; else, it will pass all values through unaltered.
ff_activation: Type of activation function at the end of each encoder block;
must be an activation-type subclass of `Layer`.
ff_dropout: Stochastic rate (probability) for dropping an activation value
when applying dropout after the FF dense layer.
ff_chunk_size: int; if > 0, chunk feed-forward into this-sized chunks
ff_use_sru: int; if > 0, we use this many SRU layers instead of feed-forward
ff_sparsity: int, if > 0 use sparse feed-forward block with this sparsity
ff_sparsity_type: string, if ff_sparsity >0,
use SparseFF if ff_sparsity_type=`'1inN'` and
use BlockSparseFF if ff_sparsity_type=`'Block'`
attention_chunk_size: int, if > 0 run attention chunked at this size
attention_type: The attention layer to use for the decoder part.
Returns:
A Transformer language model as a layer that maps from a tensor of tokens
to activations over a vocab set.
"""
positional_encoder = [
tl.Embedding(vocab_size, d_model),
tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode),
tl.PositionalEncoding(max_len=max_len, mode=mode)
]
# pylint: disable=g-complex-comprehension
decoder_blocks = [
_DecoderBlock(d_model, d_ff, n_heads, dropout, dropout_shared_axes,
mode, ff_activation, ff_dropout, ff_chunk_size,
ff_use_sru, ff_sparsity, ff_sparsity_type,
attention_chunk_size, attention_type)
for i in range(n_layers)
]
# pylint: enable=g-complex-comprehension
# Assemble and return the model.
return tl.Serial( # tokens (or chunked tuple of tokens)
tl.ShiftRight(mode=mode), # toks
positional_encoder, # vecs
decoder_blocks, # vecs
tl.LayerNorm(), # vecs
tl.Dense(vocab_size), # vecs
tl.LogSoftmax(), # vecs
)
def HourglassLM(vocab_size,
d_model=512,
d_ff=2048,
vanilla_layers=(1, 1),
hierarchy='6@3',
n_heads=8,
dropout=0.1,
dropout_shared_axes=None,
mode='train',
ff_activation=tl.FastGelu,
vanilla_attn_type=RelativeAttentionWrapper,
middle_attn_type=RelativeAttentionWrapper,
downsampling_fn=AttentionResampling,
upsampling_fn=AttentionResampling,
attention_downsampling_fn=AveragePooling,
attention_upsampling_fn=LinearUpsampling):
"""Returns a hierarchical Transformer language model.
This model performs autoregressive language modeling:
- input: rank 2 tensor representing a batch of text strings via token IDs
plus padding markers; shape is (batch_size, sequence_length). The tensor
elements are integers in `range(vocab_size)`, and `0` values mark padding
positions.
- output: rank 3 tensor representing a batch of log-probability
distributions for each sequence position over possible token IDs;
shape is (batch_size, sequence_length, `vocab_size`).
This model uses only the decoder part of the overall Transformer.
Args:
vocab_size: Input vocabulary size -- each element of the input tensor should
be an integer in `range(vocab_size)`. These integers typically represent
token IDs from a vocabulary-based tokenizer.
d_model: Final dimension of tensors at most points in the model, including
the initial embedding output.
d_ff: Size of special dense layer in the feed-forward part of each encoder
block.
vanilla_layers: (pre_layers, post_layers) tuple - number of full token-level
Transformer decoder layers before and after shortening.
hierarchy: string - shortening hierarchy, as described in the paper.
Hierarchy levels must form a palindrome, e.g. '1@2 2@6 1@2'.
n_heads: Number of attention heads.
dropout: Stochastic rate (probability) for dropping an activation value when
applying dropout within an encoder block.
dropout_shared_axes: Tensor axes on which to share a dropout mask. Sharing
along batch and sequence axes (`dropout_shared_axes=(0,1)`) is a useful
way to save memory and apply consistent masks to activation vectors at
different sequence positions.
mode: str: 'train' or 'eval'.
ff_activation: Type of activation function at the end of each encoder block;
must be an activation-type subclass of `Layer`.
vanilla_attn_type: class: attention class such as SelfAttention to use in
the layers before and after shortening (vanilla layers).
middle_attn_type: class: attention class to use in the middle layers (these
operating on the shortened sequence).
downsampling_fn: function that takes full token-level vectors of length `l`
and transforms them into `l` / `k` vectors, where `k` denotes
`shorten_factor` parameter.
upsampling_fn: function that takes shortened representations of a sequence,
consisting of `l` / `k` vectors and transforms them into full token-level
representations of length `l`.
attention_downsampling_fn: Downsampling function that transforms token-level
vectors into query vectors with reduced length. Necessary only when
AttentionResampling is used as `downsampling_fn`.
attention_upsampling_fn: Upsampling function for AttentionResampling. Valid
only when AttentionResampling is used as a `upsampling_fn`.
Returns:
A Transformer language model as a layer that maps from a tensor of tokens
to activations over a vocab set.
"""
assert mode != 'predict' # For now, 'predict' mode is unsupported.
hierarchy_n_layers, hierarchy_shorten_factors = _parse_hierarchy(hierarchy)
token_encoder = [
tl.Embedding(vocab_size, d_model),
tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode)
]
context_bias_layer, location_bias_layer = get_rel_att_inputs(d_model, n_heads)
n_pre_decoder_blocks, n_post_decoder_blocks = vanilla_layers
def create_decoder_blocks(n_layers, total_pooling, # pylint: disable = invalid-name
attention_type):
decoder_blocks = [
# pylint: disable=g-complex-comprehension
_RelativeDecoderBlock(attention_type, d_model, d_ff, n_heads, dropout,
dropout_shared_axes, mode, ff_activation,
context_bias_layer, location_bias_layer,
total_pooling) for _ in range(n_layers)
]
return decoder_blocks + [tl.LayerNorm()]
def create_hourglass_valley(rest_shorten_factors, rest_n_funnel_blocks, # pylint: disable = invalid-name
current_total_pooling):
assert rest_shorten_factors
assert len(rest_shorten_factors) == len(rest_n_funnel_blocks)
current_sf = rest_shorten_factors[0]
current_n_layers = rest_n_funnel_blocks[0]
shortening_layer = downsampling_fn(
current_sf,
d_model,
is_upsampling=False,
d_ff=d_ff,
n_heads=n_heads,
dropout=dropout,
dropout_shared_axes=dropout_shared_axes,
mode=mode,
ff_activation=ff_activation,
context_bias_layer=context_bias_layer,
location_bias_layer=location_bias_layer,
total_pooling=current_total_pooling,
resampling_fn=attention_downsampling_fn)
upsampling_layer = upsampling_fn(
current_sf,
d_model=d_model,
is_upsampling=True,
d_ff=d_ff,
n_heads=n_heads,
dropout=dropout,
dropout_shared_axes=dropout_shared_axes,
mode=mode,
ff_activation=ff_activation,
context_bias_layer=context_bias_layer,
location_bias_layer=location_bias_layer,
total_pooling=current_total_pooling,
resampling_fn=attention_upsampling_fn)
if len(rest_shorten_factors) > 1: # we need to go deeper again
pre_stage_blocks = create_decoder_blocks(
current_n_layers, current_total_pooling * current_sf,
middle_attn_type)
post_stage_blocks = create_decoder_blocks(
current_n_layers, current_total_pooling * current_sf,
middle_attn_type)
return [
tl.Dup(),
tl.ShiftRight(current_sf - 1, mode=mode), shortening_layer,
pre_stage_blocks, *create_hourglass_valley(
rest_shorten_factors[1:], rest_n_funnel_blocks[1:],
current_total_pooling * current_sf), post_stage_blocks,
upsampling_layer,
tl.LayerNorm(),
tl.Add()
]
else:
blocks = create_decoder_blocks(current_n_layers,
current_total_pooling * current_sf,
middle_attn_type)
return [
tl.Dup(),
tl.ShiftRight(current_sf - 1), shortening_layer, blocks,
upsampling_layer,
tl.LayerNorm(),
tl.Add()
]
pre_decoder_blocks = create_decoder_blocks(n_pre_decoder_blocks, 1,
vanilla_attn_type)
post_decoder_blocks = create_decoder_blocks(n_post_decoder_blocks, 1,
vanilla_attn_type)
valley = create_hourglass_valley(hierarchy_shorten_factors,
hierarchy_n_layers, 1)
# Assemble and return the model.
return tl.Serial( # tokens (or chunked tuple of tokens)
tl.ShiftRight(mode=mode), # toks
token_encoder, # vecs
pre_decoder_blocks, # vecs
valley, # shortened vecs
post_decoder_blocks, # vecs
tl.Dense(vocab_size), # vecs
)
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),
)
def TransformerLM(vocab_size,
d_model=512,
d_ff=2048,
n_layers=6,
n_heads=8,
d_attention_key=None,
d_attention_value=None,
attention_type=tl.DotProductCausalAttention,
dropout=0.1,
share_qk=False,
max_len=2048,
n_chunks=0,
mode='train',
ff_activation=tl.Relu):
"""Returns a 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
d_attention_key: int: depth of key vector for each attention head (default
is d_model // n_heads)
d_attention_value: int: depth of value vector for each attention head
(default is d_model // n_heads)
attention_type: subclass of tl.BaseCausalAttention: attention class to use
dropout: float: dropout rate (how much to drop out)
share_qk: bool, whether to share queries and keys in decoder attention
max_len: int: maximum symbol length for positional encoding
n_chunks: int: number of chunks (must match input pipeline)
mode: str: 'train', 'eval' or 'predict', predict mode is for fast inference
ff_activation: the non-linearity in feed-forward layer
Returns:
A Transformer language model as a layer that maps from a tensor of tokens
to activations over a vocab set.
"""
if n_chunks == 0:
concatenate_chunks = []
split_chunks = []
else:
concatenate_chunks = tl.Concatenate(n_items=n_chunks)
split_chunks = tl.Split(n_items=n_chunks, axis=-2)
positional_encoder = [
tl.Embedding(d_model, vocab_size),
tl.Dropout(rate=dropout, name='embedding', mode=mode),
tl.PositionalEncoding(max_len=max_len, mode=mode)
]
decoder_blocks = [
# pylint: disable=g-complex-comprehension
_DecoderBlock(d_model, d_ff, n_heads, d_attention_key,
d_attention_value, attention_type, dropout, share_qk, i,
mode, ff_activation) for i in range(n_layers)
]
# Assemble and return the model.
return tl.Serial( # tokens (or chunked tuple of tokens)
concatenate_chunks, # toks
tl.ShiftRight(mode=mode), # toks
positional_encoder, # vecs
decoder_blocks, # vecs
tl.LayerNorm(), # vecs
tl.Dense(vocab_size), # vecs
tl.LogSoftmax(), # vecs
split_chunks, # vecs (or chunked tuple of vecs)
)
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),
)
def PositionalEncoder(vocab_size): # tokens --> vectors
return [
tl.Embedding(d_model, vocab_size),
tl.Dropout(rate=dropout, mode=mode),
tl.PositionalEncoding(max_len=max_len),
]
def LSTMSeq2SeqAttn(input_vocab_size=256,
target_vocab_size=256,
d_model=512,
n_encoder_layers=2,
n_decoder_layers=2,
n_attention_heads=1,
attention_dropout=0.0,
mode='train'):
"""Returns an LSTM sequence-to-sequence model with attention.
The input to the model is a pair (input tokens, target tokens), e.g.,
an English sentence (tokenized) and its translation into German (tokenized).
The model works as follows:
* Input encoder runs on the input tokens and creates activations that
are used as both keys and values in attention.
* Pre-attention decoder runs on the targets and creates
activations that are used as queries in attention.
* Attention runs on the queries, keys and values masking out input padding.
* Decoder runs on the result, followed by a cross-entropy loss.
Args:
input_vocab_size: int: vocab size of the input
target_vocab_size: int: vocab size of the target
d_model: int: depth of embedding (n_units in the LSTM cell)
n_encoder_layers: int: number of LSTM layers in the encoder
n_decoder_layers: int: number of LSTM layers in the decoder after attention
n_attention_heads: int: number of attention heads
attention_dropout: float, dropout for the attention layer
mode: str: 'train', 'eval' or 'predict', predict mode is for fast inference
Returns:
An LSTM sequence-to-sequence model with attention.
"""
input_encoder = tl.Serial(
tl.Embedding(d_model, input_vocab_size),
[tl.LSTM(d_model) for _ in range(n_encoder_layers)],
)
pre_attention_decoder = tl.Serial(
tl.ShiftRight(mode=mode),
tl.Embedding(d_model, target_vocab_size),
tl.LSTM(d_model),
)
def PrepareAttentionInputs():
"""Layer that prepares queries, keys, values and mask for attention."""
def F(encoder_activations, decoder_activations, input_tokens):
keys = values = encoder_activations
queries = decoder_activations
# Mask is 1 where inputs are not padding (0) and 0 where they are padding.
mask = (input_tokens != 0)
# We need to add axes to the mask for attention heads and decoder length.
mask = jnp.reshape(mask, (mask.shape[0], 1, 1, mask.shape[1]))
# Broadcast so mask is [batch, 1 for heads, decoder-len, encoder-len].
mask = mask + jnp.zeros((1, 1, decoder_activations.shape[1], 1))
return queries, keys, values, mask
return tl.Fn('PrepareAttentionInputs', F, n_out=4)
return tl.Serial( # in-toks, target-toks
tl.Select([0, 1, 0, 1]), # in-toks, target-toks, in-toks, target-toks
tl.Parallel(input_encoder, pre_attention_decoder),
PrepareAttentionInputs(), # q, k, v, mask, target-toks
tl.Residual(
tl.AttentionQKV(d_model, n_heads=n_attention_heads,
dropout=attention_dropout, mode=mode)
), # decoder-vecs, mask, target-toks
tl.Select([0, 2]), # decoder-vecs, target-toks
[tl.LSTM(d_model) for _ in range(n_decoder_layers)],
tl.Dense(target_vocab_size),
tl.LogSoftmax()
)
def TransformerDecoder(vocab_size=None,
d_model=512,
d_ff=2048,
n_layers=6,
n_heads=8,
max_len=2048,
dropout=0.1,
dropout_shared_axes=None,
mode='train',
ff_activation=tl.Relu):
"""Returns a Transformer decoder.
This model maps sequential inputs to sequential outputs:
- input if `vocab_size` is specified: rank 2 tensor representing a batch
of text strings via token IDs plus padding markers; shape is
(batch_size, sequence_length). The tensor elements are integers in
`range(vocab_size)`, and `0` values mark padding positions.
- input if `vocab_size` is None: rank 2 tensor representing a batch
of activation vectors; shape is (batch_size, sequence_length, `d_model`).
- output: rank 3 tensor with shape (batch_size, sequence_length, `d_model`).
The model uses causal attention and does *not* shift the input to the right.
Thus, the output for position `t` is based on inputs up to and including
position `t`.
Args:
vocab_size: If specified, gives the input vocabulary size -- each element
of the input tensor should be an integer in `range(vocab_size)`.
If None, indicates that the model expects as input floating point
vectors, each with `d_model` components.
d_model: Final dimension of tensors at most points in the model, including
the initial embedding output.
d_ff: Size of special dense layer in the feed-forward part of each decoder
block.
n_layers: Number of decoder blocks. Each block includes attention, dropout,
residual, feed-forward (`Dense`), and activation layers.
n_heads: Number of attention heads.
max_len: Maximum symbol length for positional encoding.
dropout: Stochastic rate (probability) for dropping an activation value
when applying dropout within a decoder block.
dropout_shared_axes: Tensor axes on which to share a dropout mask.
Sharing along batch and sequence axes (`dropout_shared_axes=(0,1)`) is
a useful way to save memory and apply consistent masks to activation
vectors at different sequence positions.
mode: If `'train'`, each decoder block will include dropout; else, it will
pass all values through unaltered.
ff_activation: Type of activation function at the end of each decoder
block; must be an activation-type subclass of `Layer`.
Returns:
If `vocab_size` is defined: a Transformer model that maps strings (conveyed
via token IDs) to sequences of activation vectors.
If `vocab_size` is None: a Transformer model that maps sequences of
activation vectors to sequences of activation vectors.
"""
positional_encoder = [(tl.Embedding(vocab_size, d_model)
if vocab_size is not None else tl.Dense(d_model)),
tl.Dropout(rate=dropout,
shared_axes=dropout_shared_axes,
mode=mode),
tl.PositionalEncoding(max_len=max_len)]
decoder_blocks = [
# pylint: disable=g-complex-comprehension
_DecoderBlock(d_model, d_ff, n_heads, dropout, dropout_shared_axes,
mode, ff_activation) for i in range(n_layers)
]
# Assemble and return the model.
return tl.Serial( # toks
positional_encoder, # vecs
decoder_blocks, # vecs
tl.LayerNorm(), # vecs
)
def test_run_reversible_same_as_default_extended(self):
"""Runs the reversible trainer, check results are the same as default."""
inputs_batch = np.arange(8).reshape((2, 4))
targets_batch = 2 * inputs_batch
labeled_batch = (inputs_batch, targets_batch,
np.ones_like(targets_batch))
# We want to test rng propagation too, so adding some dropout layers.
first_layer = tl.Serial(tl.Embedding(9, 4), tl.Dropout(0.5), tl.Dup())
rev_layers1 = [
tl.ReversibleHalfResidual(tl.Dense(4), tl.Dropout(0.2)),
tl.ReversibleSwap(),
tl.ReversibleHalfResidual(tl.Dropout(0.5), tl.Dense(4)),
tl.ReversibleSwap()
]
mid_layer = tl.Serial(tl.Add(), tl.Dense(4), tl.Dup())
rev_layers2 = [
tl.ReversibleHalfResidual(tl.Dense(4), tl.Dropout(0.3)),
tl.ReversibleSwap()
]
loss_layer = tl.Serial(tl.Concatenate(), tl.Dense(19), tl.Dropout(0.3),
tl.LogSoftmax(), tl.CrossEntropyLoss())
model = tl.Serial([first_layer] + rev_layers1 + [mid_layer] +
rev_layers2 + [loss_layer])
rng_init = fastmath.random.get_prng(12)
model.init(labeled_batch, rng=rng_init)
optimizer_fn = optimizers.Adam # to test slots
# Make 3 steps with the original trainer.
optimizer = optimizer_fn()
optimizer.tree_init(model.weights)
trainer = optimizers.Trainer(model, optimizer)
rng_step1 = fastmath.random.get_prng(7)
rng_step2 = fastmath.random.get_prng(8)
rng_step3 = fastmath.random.get_prng(9)
trainer.one_step(labeled_batch, rng_step1)
trainer.one_step(labeled_batch, rng_step2, learning_rate=0.02)
trainer.one_step(labeled_batch, rng_step3, learning_rate=0.03)
first_layer_weights1 = first_layer.weights
rev_layer12_weights1 = rev_layers1[2].weights
mid_layer_weights1 = mid_layer.weights
rev_layer20_weights1 = rev_layers2[0].weights
loss_layer_weights1 = loss_layer.weights
# Now make 3 steps with reversible trainer.
model.init(labeled_batch, rng=rng_init)
# TODO(lukaszkaiser): this test seems to fail with memoize_jit, why?
trainer = optimizers.ReversibleSerialTrainer(
[(first_layer.sublayers, rev_layers1),
(mid_layer.sublayers, rev_layers2)],
loss_layer,
optimizer_fn,
memoize_jit=False)
trainer.one_step(labeled_batch, rng_step1)
trainer.one_step(labeled_batch, rng_step2, learning_rate=0.02)
trainer.one_step(labeled_batch, rng_step3, learning_rate=0.03)
# Check that weights end up the same.
self._assert_all_equal(loss_layer_weights1, loss_layer.weights)
self._assert_all_equal(rev_layer20_weights1, rev_layers2[0].weights)
self._assert_all_equal(mid_layer_weights1, mid_layer.weights)
self._assert_all_equal(rev_layer12_weights1, rev_layers1[2].weights)
self._assert_all_equal(first_layer_weights1, first_layer.weights)
def TransformerEncoder(vocab_size,
n_classes=10,
d_model=512,
d_ff=2048,
n_layers=6,
n_heads=8,
max_len=2048,
dropout=0.1,
dropout_shared_axes=None,
mode='train',
ff_activation=tl.Relu):
"""Returns a Transformer encoder merged with an N-way categorization head.
This model performs text categorization:
- input: rank 2 tensor representing a batch of text strings via token IDs
plus padding markers; shape is (batch_size, sequence_length). The tensor
elements are integers in `range(vocab_size)`, and `0` values mark padding
positions.
- output: rank 2 tensor representing a batch of log-probability
distributions over N categories; shape is (batch_size, `n_classes`).
Args:
vocab_size: Input vocabulary size -- each element of the input tensor
should be an integer in `range(vocab_size)`. These integers typically
represent token IDs from a vocabulary-based tokenizer.
n_classes: Final dimension of the output tensors, representing N-way
classification.
d_model: Final dimension of tensors at most points in the model, including
the initial embedding output.
d_ff: Size of special dense layer in the feed-forward part of each encoder
block.
n_layers: Number of encoder blocks. Each block includes attention, dropout,
residual, feed-forward (`Dense`), and activation layers.
n_heads: Number of attention heads.
max_len: Maximum symbol length for positional encoding.
dropout: Stochastic rate (probability) for dropping an activation value
when applying dropout within an encoder block.
dropout_shared_axes: Tensor axes on which to share a dropout mask.
Sharing along batch and sequence axes (`dropout_shared_axes=(0,1)`) is
a useful way to save memory and apply consistent masks to activation
vectors at different sequence positions.
mode: If `'train'`, each encoder block will include dropout; else, it will
pass all values through unaltered.
ff_activation: Type of activation function at the end of each encoder
block; must be an activation-type subclass of `Layer`.
Returns:
A Transformer model that maps strings (conveyed via token IDs) to
probability-like activations over a range of output classes.
"""
positional_encoder = [
tl.Embedding(vocab_size, d_model),
tl.Dropout(rate=dropout, shared_axes=dropout_shared_axes, mode=mode),
tl.PositionalEncoding(max_len=max_len)
]
encoder_blocks = [
_EncoderBlock(d_model, d_ff, n_heads, dropout, dropout_shared_axes,
mode, ff_activation) for i in range(n_layers)
]
# Assemble and return the model.
return tl.Serial( # toks
# Encode.
tl.Branch(positional_encoder, tl.PaddingMask()), # vecs masks
encoder_blocks, # vecs masks
tl.Select([0], n_in=2), # vecs
tl.LayerNorm(), # vecs
# Map to output categories.
tl.Mean(axis=1), # vecs
tl.Dense(n_classes), # vecs
tl.LogSoftmax(), # vecs
)
def ReformerLM(vocab_size,
d_model=512,
d_ff=2048,
d_attention_key=64,
d_attention_value=64,
n_layers=6,
n_heads=8,
dropout=0.1,
max_len=2048,
attention_type=tl.SelfAttention,
axial_pos_shape=(),
d_axial_pos_embs=None,
ff_activation=tl.FastGelu,
ff_use_sru=0,
ff_chunk_size=0,
ff_sparsity=0,
attention_chunk_size=0,
mode='train'):
"""Reversible transformer language model (only uses a decoder, no encoder).
Args:
vocab_size: int: vocab size
d_model: int: depth of *each half* of the two-part features
d_ff: int: depth of feed-forward layer
d_attention_key: int: depth of key vector for each attention head
d_attention_value: int: depth of value vector for each attention head
n_layers: int: number of 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
attention_type: class: attention class to use, such as SelfAttention.
axial_pos_shape: tuple of ints: input shape to use for the axial position
encoding. If unset, axial position encoding is disabled.
d_axial_pos_embs: tuple of ints: depth of position embedding for each axis.
Tuple length must match axial_pos_shape, and values must sum to d_model.
ff_activation: the non-linearity in feed-forward layer
ff_use_sru: int; if > 0, we use this many SRU layers instead of feed-forward
ff_chunk_size: int; if > 0, chunk feed-forward into this-sized chunks
ff_sparsity: int, if > 0 use sparse feed-forward block with this sparsity
attention_chunk_size: int, if > 0 run attention chunked at this size
mode: str: 'train', 'eval', or 'predict'
Returns:
the layer.
"""
positional_encoding = ct.PositionalEncoder(mode, dropout, max_len,
axial_pos_shape,
d_axial_pos_embs)
positional_embedder = [
tl.Embedding(vocab_size, d_model),
tl.Dropout(rate=dropout, shared_axes=[-2], mode=mode), # pylint: disable=no-value-for-parameter
positional_encoding,
]
decoder_blocks = []
if isinstance(attention_type, (tuple, list)):
assert n_layers % len(attention_type) == 0
else:
attention_type = [attention_type]
for layer_idx in range(n_layers):
layer_attention_type = attention_type[layer_idx % len(attention_type)]
decoder_block = DecoderBlock(d_model,
d_ff,
d_attention_key,
d_attention_value,
n_heads,
attention_type=layer_attention_type,
dropout=dropout,
ff_activation=ff_activation,
ff_dropout=dropout,
ff_use_sru=ff_use_sru,
ff_chunk_size=ff_chunk_size,
ff_sparsity=ff_sparsity,
attention_chunk_size=attention_chunk_size,
mode=mode)
decoder_blocks.append(decoder_block)
return tl.Serial(
tl.ShiftRight(mode=mode),
positional_embedder,
tl.Dup(),
tl.ReversibleSerial(decoder_blocks),
tl.Concatenate(),
# TODO(kitaev): Test whether dropout should go before or after the
# LayerNorm, and whether dropout broadcasting is needed here.
tl.LayerNorm(),
tl.Dropout(rate=dropout, shared_axes=[-2], mode=mode), # pylint: disable=no-value-for-parameter
tl.Dense(vocab_size),
tl.LogSoftmax(),
)
def TransformerLM(vocab_size,
d_model=D_MODEL,
d_ff=D_FF,
n_layers=N_LAYERS,
n_heads=N_HEADS,
max_len=MAX_SEQUENCE_LENGTH,
dropout=DROPOUT_RATE,
dropout_shared_axes=DROPOUT_SHARED_AXES,
mode=MODE,
ff_activation=FF_ACTIVATION_TYPE):
"""Returns a Transformer language model.
This model performs autoregressive language modeling:
- input: Array representing a batch of text strings via token IDs
plus padding markers; shape is (batch_size, sequence_length). Array
elements are integers in ``range(vocab_size)``, and 0 values mark padding
positions.
- output: 3-D array of raw activations with last/innermost dimension of
``vocab_size``, suitable for decoding into a batch of token strings;
shape is (batch_size, sequence_length, ``vocab_size``).
This model uses only the decoder part of the overall Transformer.
Args:
vocab_size: Input vocabulary size -- each element of the input array
should be an integer in ``range(vocab_size)``. These integers typically
represent token IDs from a vocabulary-based tokenizer.
d_model: Last/innermost dimension of activation arrays at most points in
the model, including the initial embedding output.
d_ff: Last/innermost dimension of special (typically wider)
:py:class:`Dense` layer in the feedforward part of each encoder block.
n_layers: Number of decoder blocks. Each block includes attention, dropout,
residual, layer-norm, feedforward (:py:class:`Dense`), and activation
layers.
n_heads: Number of attention heads.
max_len: Maximum symbol length for positional encoding.
dropout: Stochastic rate (probability) for dropping an activation value
when applying dropout within decoder blocks. The same rate is also
used for attention dropout in decoder blocks.
dropout_shared_axes: Tensor axes on which to share a dropout mask.
Sharing along batch and sequence axes (``dropout_shared_axes=(0,1)``)
is a useful way to save memory and apply consistent masks to activation
vectors at different sequence positions.
mode: If ``'predict'``, use fast inference. If ``'train'``, each decoder
block will include dropout; else, it will pass all values through
unaltered.
ff_activation: Type of activation function at the end of each encoder
block; must be an activation-type subclass of :py:class:`Layer`.
Returns:
A Transformer language model that maps strings (represented as token ID
sequences) to sequences of raw (non-normalized) activation vectors; each
vector in the sequence can be mapped (e.g., by `argmax`) to a token ID.
"""
def _Dropout():
return tl.Dropout(rate=dropout,
shared_axes=dropout_shared_axes,
mode=mode)
def _DecBlock():
return _DecoderBlock(d_model, d_ff, n_heads, dropout,
dropout_shared_axes, mode, ff_activation)
return tl.Serial(
tl.ShiftRight(mode=mode),
tl.Embedding(vocab_size, d_model),
_Dropout(),
tl.PositionalEncoding(max_len=max_len, mode=mode),
[_DecBlock() for _ in range(n_layers)],
tl.LayerNorm(),
tl.Dense(vocab_size),
)
