juancopi81's picture
Add t5x and mt3 models
b100e1c
raw
history blame
13.8 kB
# Copyright 2022 The MT3 Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""T5.1.1 Transformer model."""
from typing import Any, Sequence
from flax import linen as nn
from flax import struct
import jax.numpy as jnp
from mt3 import layers
@struct.dataclass
class T5Config:
"""Global hyperparameters used to minimize obnoxious kwarg plumbing."""
vocab_size: int
# Activation dtypes.
dtype: Any = jnp.float32
emb_dim: int = 512
num_heads: int = 8
num_encoder_layers: int = 6
num_decoder_layers: int = 6
head_dim: int = 64
mlp_dim: int = 2048
# Activation functions are retrieved from Flax.
mlp_activations: Sequence[str] = ('relu',)
dropout_rate: float = 0.1
# If `True`, the embedding weights are used in the decoder output layer.
logits_via_embedding: bool = False
class EncoderLayer(nn.Module):
"""Transformer encoder layer."""
config: T5Config
@nn.compact
def __call__(self, inputs, encoder_mask=None, deterministic=False):
cfg = self.config
# Attention block.
assert inputs.ndim == 3
x = layers.LayerNorm(
dtype=cfg.dtype, name='pre_attention_layer_norm')(
inputs)
# [batch, length, emb_dim] -> [batch, length, emb_dim]
x = layers.MultiHeadDotProductAttention(
num_heads=cfg.num_heads,
dtype=cfg.dtype,
head_dim=cfg.head_dim,
dropout_rate=cfg.dropout_rate,
name='attention')(
x, x, encoder_mask, deterministic=deterministic)
x = nn.Dropout(
rate=cfg.dropout_rate, broadcast_dims=(-2,))(
x, deterministic=deterministic)
x = x + inputs
# MLP block.
y = layers.LayerNorm(dtype=cfg.dtype, name='pre_mlp_layer_norm')(x)
# [batch, length, emb_dim] -> [batch, length, emb_dim]
y = layers.MlpBlock(
intermediate_dim=cfg.mlp_dim,
activations=cfg.mlp_activations,
intermediate_dropout_rate=cfg.dropout_rate,
dtype=cfg.dtype,
name='mlp',
)(y, deterministic=deterministic)
y = nn.Dropout(
rate=cfg.dropout_rate, broadcast_dims=(-2,))(
y, deterministic=deterministic)
y = y + x
return y
class DecoderLayer(nn.Module):
"""Transformer decoder layer that attends to the encoder."""
config: T5Config
@nn.compact
def __call__(self,
inputs,
encoded,
decoder_mask=None,
encoder_decoder_mask=None,
deterministic=False,
decode=False,
max_decode_length=None):
cfg = self.config
# inputs: embedded inputs to the decoder with shape [batch, length, emb_dim]
x = layers.LayerNorm(
dtype=cfg.dtype, name='pre_self_attention_layer_norm')(
inputs)
# Self-attention block
x = layers.MultiHeadDotProductAttention(
num_heads=cfg.num_heads,
dtype=cfg.dtype,
head_dim=cfg.head_dim,
dropout_rate=cfg.dropout_rate,
name='self_attention')(
x,
x,
decoder_mask,
deterministic=deterministic,
decode=decode)
x = nn.Dropout(
rate=cfg.dropout_rate, broadcast_dims=(-2,))(
x, deterministic=deterministic)
x = x + inputs
# Encoder-Decoder block.
y = layers.LayerNorm(
dtype=cfg.dtype, name='pre_cross_attention_layer_norm')(
x)
y = layers.MultiHeadDotProductAttention(
num_heads=cfg.num_heads,
dtype=cfg.dtype,
head_dim=cfg.head_dim,
dropout_rate=cfg.dropout_rate,
name='encoder_decoder_attention')(
y, encoded, encoder_decoder_mask, deterministic=deterministic)
y = nn.Dropout(
rate=cfg.dropout_rate, broadcast_dims=(-2,))(
y, deterministic=deterministic)
y = y + x
# MLP block.
z = layers.LayerNorm(dtype=cfg.dtype, name='pre_mlp_layer_norm')(y)
z = layers.MlpBlock(
intermediate_dim=cfg.mlp_dim,
activations=cfg.mlp_activations,
intermediate_dropout_rate=cfg.dropout_rate,
dtype=cfg.dtype,
name='mlp',
)(z, deterministic=deterministic)
z = nn.Dropout(
rate=cfg.dropout_rate, broadcast_dims=(-2,))(
z, deterministic=deterministic)
z = z + y
return z
class Encoder(nn.Module):
"""A stack of encoder layers."""
config: T5Config
@nn.compact
def __call__(self,
encoder_input_tokens,
encoder_mask=None,
deterministic=False):
cfg = self.config
assert encoder_input_tokens.ndim == 3 # [batch, length, depth]
seq_length = encoder_input_tokens.shape[-2]
inputs_positions = jnp.arange(seq_length)[None, :]
# [batch, length, depth] -> [batch, length, emb_dim]
x = layers.DenseGeneral(
cfg.emb_dim,
dtype=cfg.dtype,
kernel_init=nn.linear.default_kernel_init,
kernel_axes=('vocab', 'embed'),
name='continuous_inputs_projection')(encoder_input_tokens)
x = x + layers.FixedEmbed(features=cfg.emb_dim)(inputs_positions)
x = nn.Dropout(
rate=cfg.dropout_rate, broadcast_dims=(-2,))(
x, deterministic=deterministic)
x = x.astype(cfg.dtype)
for lyr in range(cfg.num_encoder_layers):
# [batch, length, emb_dim] -> [batch, length, emb_dim]
x = EncoderLayer(
config=cfg,
name=f'layers_{lyr}')(x, encoder_mask, deterministic)
x = layers.LayerNorm(dtype=cfg.dtype, name='encoder_norm')(x)
return nn.Dropout(rate=cfg.dropout_rate)(x, deterministic=deterministic)
class Decoder(nn.Module):
"""A stack of decoder layers as a part of an encoder-decoder architecture."""
config: T5Config
@nn.compact
def __call__(self,
encoded,
decoder_input_tokens,
decoder_positions=None,
decoder_mask=None,
encoder_decoder_mask=None,
deterministic=False,
decode=False,
max_decode_length=None):
cfg = self.config
assert decoder_input_tokens.ndim == 2 # [batch, len]
seq_length = decoder_input_tokens.shape[-1]
decoder_positions = jnp.arange(seq_length)[None, :]
# [batch, length] -> [batch, length, emb_dim]
y = layers.Embed(
num_embeddings=cfg.vocab_size,
features=cfg.emb_dim,
dtype=cfg.dtype,
attend_dtype=jnp.float32, # for logit training stability
embedding_init=nn.initializers.normal(stddev=1.0),
one_hot=True,
name='token_embedder')(decoder_input_tokens.astype('int32'))
y = y + layers.FixedEmbed(features=cfg.emb_dim)(
decoder_positions, decode=decode)
y = nn.Dropout(
rate=cfg.dropout_rate, broadcast_dims=(-2,))(
y, deterministic=deterministic)
y = y.astype(cfg.dtype)
for lyr in range(cfg.num_decoder_layers):
# [batch, length, emb_dim] -> [batch, length, emb_dim]
y = DecoderLayer(
config=cfg, name=f'layers_{lyr}')(
y,
encoded,
decoder_mask=decoder_mask,
encoder_decoder_mask=encoder_decoder_mask,
deterministic=deterministic,
decode=decode,
max_decode_length=max_decode_length)
y = layers.LayerNorm(dtype=cfg.dtype, name='decoder_norm')(y)
y = nn.Dropout(
rate=cfg.dropout_rate, broadcast_dims=(-2,))(
y, deterministic=deterministic)
# [batch, length, emb_dim] -> [batch, length, vocab_size]
if cfg.logits_via_embedding:
# Use the transpose of embedding matrix for logit transform.
logits = self.shared_embedding.attend(y)
# Correctly normalize pre-softmax logits for this shared case.
logits = logits / jnp.sqrt(y.shape[-1])
else:
logits = layers.DenseGeneral(
cfg.vocab_size,
dtype=jnp.float32, # Use float32 for stabiliity.
kernel_axes=('embed', 'vocab'),
name='logits_dense')(
y)
return logits
class Transformer(nn.Module):
"""An encoder-decoder Transformer model."""
config: T5Config
def setup(self):
cfg = self.config
self.encoder = Encoder(config=cfg)
self.decoder = Decoder(config=cfg)
def encode(self,
encoder_input_tokens,
encoder_segment_ids=None,
enable_dropout=True):
"""Applies Transformer encoder-branch on the inputs."""
cfg = self.config
assert encoder_input_tokens.ndim == 3 # (batch, length, depth)
# Make padding attention mask; we don't actually mask out any input
# positions, letting the model potentially attend to the zero vector used as
# padding.
encoder_mask = layers.make_attention_mask(
jnp.ones(encoder_input_tokens.shape[:-1]),
jnp.ones(encoder_input_tokens.shape[:-1]),
dtype=cfg.dtype)
# Add segmentation block-diagonal attention mask if using segmented data.
if encoder_segment_ids is not None:
encoder_mask = layers.combine_masks(
encoder_mask,
layers.make_attention_mask(
encoder_segment_ids,
encoder_segment_ids,
jnp.equal,
dtype=cfg.dtype))
return self.encoder(
encoder_input_tokens, encoder_mask, deterministic=not enable_dropout)
def decode(
self,
encoded,
encoder_input_tokens, # only needed for masks
decoder_input_tokens,
decoder_target_tokens,
encoder_segment_ids=None,
decoder_segment_ids=None,
decoder_positions=None,
enable_dropout=True,
decode=False,
max_decode_length=None):
"""Applies Transformer decoder-branch on encoded-input and target."""
cfg = self.config
# Make padding attention masks.
if decode:
# Do not mask decoder attention based on targets padding at
# decoding/inference time.
decoder_mask = None
encoder_decoder_mask = layers.make_attention_mask(
jnp.ones_like(decoder_target_tokens),
jnp.ones(encoder_input_tokens.shape[:-1]),
dtype=cfg.dtype)
else:
decoder_mask = layers.make_decoder_mask(
decoder_target_tokens=decoder_target_tokens,
dtype=cfg.dtype,
decoder_segment_ids=decoder_segment_ids)
encoder_decoder_mask = layers.make_attention_mask(
decoder_target_tokens > 0,
jnp.ones(encoder_input_tokens.shape[:-1]),
dtype=cfg.dtype)
# Add segmentation block-diagonal attention masks if using segmented data.
if encoder_segment_ids is not None:
if decode:
raise ValueError(
'During decoding, packing should not be used but '
'`encoder_segment_ids` was passed to `Transformer.decode`.')
encoder_decoder_mask = layers.combine_masks(
encoder_decoder_mask,
layers.make_attention_mask(
decoder_segment_ids,
encoder_segment_ids,
jnp.equal,
dtype=cfg.dtype))
logits = self.decoder(
encoded,
decoder_input_tokens=decoder_input_tokens,
decoder_positions=decoder_positions,
decoder_mask=decoder_mask,
encoder_decoder_mask=encoder_decoder_mask,
deterministic=not enable_dropout,
decode=decode,
max_decode_length=max_decode_length)
return logits.astype(self.config.dtype)
def __call__(self,
encoder_input_tokens,
decoder_input_tokens,
decoder_target_tokens,
encoder_segment_ids=None,
decoder_segment_ids=None,
encoder_positions=None,
decoder_positions=None,
*,
enable_dropout: bool = True,
decode: bool = False):
"""Applies Transformer model on the inputs.
This method requires both decoder_target_tokens and decoder_input_tokens,
which is a shifted version of the former. For a packed dataset, it usually
has additional processing applied. For example, the first element of each
sequence has id 0 instead of the shifted EOS id from the previous sequence.
Args:
encoder_input_tokens: input data to the encoder.
decoder_input_tokens: input token to the decoder.
decoder_target_tokens: target token to the decoder.
encoder_segment_ids: encoder segmentation info for packed examples.
decoder_segment_ids: decoder segmentation info for packed examples.
encoder_positions: encoder subsequence positions for packed examples.
decoder_positions: decoder subsequence positions for packed examples.
enable_dropout: Ensables dropout if set to True.
decode: Whether to prepare and use an autoregressive cache.
Returns:
logits array from full transformer.
"""
encoded = self.encode(
encoder_input_tokens,
encoder_segment_ids=encoder_segment_ids,
enable_dropout=enable_dropout)
return self.decode(
encoded,
encoder_input_tokens, # only used for masks
decoder_input_tokens,
decoder_target_tokens,
encoder_segment_ids=encoder_segment_ids,
decoder_segment_ids=decoder_segment_ids,
decoder_positions=decoder_positions,
enable_dropout=enable_dropout,
decode=decode)