def __init__(self, **kwargs):
super().__init__(**kwargs)
embedding_width = 768
dropout_rate = 0.1
initializer = tf.keras.initializers.TruncatedNormal(stddev=0.02)
self._embedding_layer = layers.OnDeviceEmbedding(
vocab_size=30522,
embedding_width=embedding_width,
initializer=initializer,
name="word_embeddings",
)
# Always uses dynamic slicing for simplicity.
self._position_embedding_layer = layers.PositionEmbedding(
initializer=initializer,
use_dynamic_slicing=True,
max_sequence_length=512,
name="position_embedding",
)
self._type_embedding_layer = layers.OnDeviceEmbedding(
vocab_size=2,
embedding_width=embedding_width,
initializer=initializer,
use_one_hot=True,
name="type_embeddings",
)
self._add = tf.keras.layers.Add()
self._layer_norm = tf.keras.layers.LayerNormalization(
name="embeddings/layer_norm",
axis=-1,
epsilon=1e-12,
dtype=tf.float32)
self._dropout = tf.keras.layers.Dropout(rate=dropout_rate)
self._attention_mask = layers.SelfAttentionMask()
self._transformer_layers = []
for i in range(12):
layer = layers.Transformer(
num_attention_heads=12,
intermediate_size=3072,
intermediate_activation=activations.gelu,
dropout_rate=dropout_rate,
attention_dropout_rate=0.1,
output_range=None,
kernel_initializer=initializer,
name="transformer/layer_%d" % i,
)
self._transformer_layers.append(layer)
self._lambda = tf.keras.layers.Lambda(
lambda x: tf.squeeze(x[:, 0:1, :], axis=1))
self._pooler_layer = tf.keras.layers.Dense(
units=embedding_width,
activation="tanh",
kernel_initializer=initializer,
name="pooler_transform",
)
def build(self, input_shape):
"""Implements build() for the layer."""
self.encoder_layers = []
for i in range(self._num_layers):
self.encoder_layers.append(
layers.Transformer(
num_attention_heads=self._num_attention_heads,
intermediate_size=self._intermediate_size,
intermediate_activation=self._activation,
dropout_rate=self._dropout_rate,
attention_dropout_rate=self._attention_dropout_rate,
use_bias=self._use_bias,
norm_first=self._norm_first,
norm_epsilon=self._norm_epsilon,
intermediate_dropout=self._intermediate_dropout,
attention_initializer=attention_initializer(input_shape[2]),
name=("layer_%d" % i)))
self.output_normalization = tf.keras.layers.LayerNormalization(
epsilon=self._norm_epsilon, dtype="float32")
super(TransformerEncoder, self).build(input_shape)
def __init__(self,
encoder_cfg,
num_task_agnostic_layers,
output='logits',
name='rtd',
**kwargs):
super(ReplacedTokenDetectionHead, self).__init__(name=name, **kwargs)
self.num_task_agnostic_layers = num_task_agnostic_layers
self.hidden_size = encoder_cfg['embedding_cfg']['hidden_size']
self.num_hidden_instances = encoder_cfg['num_hidden_instances']
self.hidden_cfg = encoder_cfg['hidden_cfg']
self.activation = self.hidden_cfg['intermediate_activation']
self.initializer = self.hidden_cfg['kernel_initializer']
self.hidden_layers = []
for i in range(self.num_task_agnostic_layers,
self.num_hidden_instances):
self.hidden_layers.append(
layers.Transformer(
num_attention_heads=self.hidden_cfg['num_attention_heads'],
intermediate_size=self.hidden_cfg['intermediate_size'],
intermediate_activation=self.activation,
dropout_rate=self.hidden_cfg['dropout_rate'],
attention_dropout_rate=self.
hidden_cfg['attention_dropout_rate'],
kernel_initializer=self.initializer,
name='transformer/layer_%d_rtd' % i))
self.dense = tf.keras.layers.Dense(self.hidden_size,
activation=self.activation,
kernel_initializer=self.initializer,
name='transform/rtd_dense')
self.rtd_head = tf.keras.layers.Dense(
units=1,
kernel_initializer=self.initializer,
name='transform/rtd_head')
if output not in ('predictions', 'logits'):
raise ValueError((
'Unknown `output` value "%s". `output` can be either "logits" or '
'"predictions"') % output)
self._output_type = output
def build(self, input_shape):
self.hidden_layers = []
for i in range(self.num_task_agnostic_layers,
self.num_hidden_instances):
self.hidden_layers.append(
layers.Transformer(
num_attention_heads=self.hidden_cfg['num_attention_heads'],
intermediate_size=self.hidden_cfg['intermediate_size'],
intermediate_activation=self.activation,
dropout_rate=self.hidden_cfg['dropout_rate'],
attention_dropout_rate=self.
hidden_cfg['attention_dropout_rate'],
kernel_initializer=self.initializer,
name='transformer/layer_%d_rtd' % i))
self.dense = tf.keras.layers.Dense(self.hidden_size,
activation=self.activation,
kernel_initializer=self.initializer,
name='transform/rtd_dense')
self.rtd_head = tf.keras.layers.Dense(
units=1,
kernel_initializer=self.initializer,
name='transform/rtd_head')
def __init__(
self,
vocab_size,
hidden_size=768,
num_layers=12,
num_attention_heads=12,
sequence_length=512,
max_sequence_length=None,
type_vocab_size=16,
intermediate_size=3072,
activation=activations.gelu,
dropout_rate=0.1,
attention_dropout_rate=0.1,
initializer=tf.keras.initializers.TruncatedNormal(stddev=0.02),
float_dtype='float32',
**kwargs):
activation = tf.keras.activations.get(activation)
initializer = tf.keras.initializers.get(initializer)
if not max_sequence_length:
max_sequence_length = sequence_length
self._self_setattr_tracking = False
self._config_dict = {
'vocab_size': vocab_size,
'hidden_size': hidden_size,
'num_layers': num_layers,
'num_attention_heads': num_attention_heads,
'sequence_length': sequence_length,
'max_sequence_length': max_sequence_length,
'type_vocab_size': type_vocab_size,
'intermediate_size': intermediate_size,
'activation': tf.keras.activations.serialize(activation),
'dropout_rate': dropout_rate,
'attention_dropout_rate': attention_dropout_rate,
'initializer': tf.keras.initializers.serialize(initializer),
'float_dtype': float_dtype,
}
word_ids = tf.keras.layers.Input(shape=(sequence_length, ),
dtype=tf.int32,
name='input_word_ids')
mask = tf.keras.layers.Input(shape=(sequence_length, ),
dtype=tf.int32,
name='input_mask')
type_ids = tf.keras.layers.Input(shape=(sequence_length, ),
dtype=tf.int32,
name='input_type_ids')
self._embedding_layer = layers.OnDeviceEmbedding(
vocab_size=vocab_size,
embedding_width=hidden_size,
initializer=initializer,
name='word_embeddings')
word_embeddings = self._embedding_layer(word_ids)
# Always uses dynamic slicing for simplicity.
self._position_embedding_layer = layers.PositionEmbedding(
initializer=initializer,
use_dynamic_slicing=True,
max_sequence_length=max_sequence_length)
position_embeddings = self._position_embedding_layer(word_embeddings)
type_embeddings = (layers.OnDeviceEmbedding(
vocab_size=type_vocab_size,
embedding_width=hidden_size,
initializer=initializer,
use_one_hot=True,
name='type_embeddings')(type_ids))
embeddings = tf.keras.layers.Add()(
[word_embeddings, position_embeddings, type_embeddings])
embeddings = (tf.keras.layers.LayerNormalization(
name='embeddings/layer_norm',
axis=-1,
epsilon=1e-12,
dtype=tf.float32)(embeddings))
embeddings = (tf.keras.layers.Dropout(rate=dropout_rate,
dtype=tf.float32)(embeddings))
if float_dtype == 'float16':
embeddings = tf.cast(embeddings, tf.float16)
data = embeddings
attention_mask = MakeAttentionMaskLayer()([data, mask])
for i in range(num_layers):
layer = layers.Transformer(
num_attention_heads=num_attention_heads,
intermediate_size=intermediate_size,
intermediate_activation=activation,
dropout_rate=dropout_rate,
attention_dropout_rate=attention_dropout_rate,
kernel_initializer=initializer,
dtype=float_dtype,
name='transformer/layer_%d' % i)
data = layer([data, attention_mask])
first_token_tensor = (tf.keras.layers.Lambda(
lambda x: tf.squeeze(x[:, 0:1, :], axis=1))(data))
cls_output = tf.keras.layers.Dense(
units=hidden_size,
activation='tanh',
kernel_initializer=initializer,
name='pooler_transform')(first_token_tensor)
super(TransformerEncoder,
self).__init__(inputs=[word_ids, mask, type_ids],
outputs=[data, cls_output],
**kwargs)
def __init__(
self,
vocab_size,
hidden_size=768,
num_layers=12,
num_attention_heads=12,
sequence_length=512,
max_sequence_length=None,
type_vocab_sizes=(16, ),
num_float_features=0,
intermediate_size=3072,
activation=activations.gelu,
dropout_rate=0.1,
attention_dropout_rate=0.1,
initializer=tf.keras.initializers.TruncatedNormal(stddev=0.02),
return_all_encoder_outputs=False,
bert_init_ckpt=None,
**kwargs):
activation = tf.keras.activations.get(activation)
initializer = tf.keras.initializers.get(initializer)
if not max_sequence_length:
max_sequence_length = sequence_length
self._self_setattr_tracking = False
num_type_features = len(type_vocab_sizes)
self._config_dict = {
'vocab_size': vocab_size,
'hidden_size': hidden_size,
'num_layers': num_layers,
'num_attention_heads': num_attention_heads,
'sequence_length': sequence_length,
'max_sequence_length': max_sequence_length,
'type_vocab_sizes': type_vocab_sizes,
'num_type_features': num_type_features,
'num_float_features': num_float_features,
'intermediate_size': intermediate_size,
'activation': tf.keras.activations.serialize(activation),
'dropout_rate': dropout_rate,
'attention_dropout_rate': attention_dropout_rate,
'initializer': tf.keras.initializers.serialize(initializer),
'return_all_encoder_outputs': return_all_encoder_outputs,
}
word_ids = tf.keras.layers.Input(shape=(sequence_length, ),
dtype=tf.int32,
name='input_word_ids')
mask = tf.keras.layers.Input(shape=(sequence_length, ),
dtype=tf.int32,
name='input_mask')
all_inputs = [word_ids, mask]
if num_type_features:
type_ids = tf.keras.layers.Input(shape=(sequence_length,
num_type_features),
dtype=tf.int32,
name='input_type_ids')
all_inputs.append(type_ids)
if num_float_features:
float_features = tf.keras.layers.Input(shape=(sequence_length,
num_float_features),
dtype=tf.float32,
name='float_features')
all_inputs.append(float_features)
self._embedding_layer = layers.OnDeviceEmbedding(
vocab_size=vocab_size,
embedding_width=hidden_size,
initializer=initializer,
name='word_embeddings')
word_embeddings = self._embedding_layer(word_ids)
# Always uses dynamic slicing for simplicity.
self._position_embedding_layer = modeling.layers.PositionEmbedding(
initializer=initializer, max_length=max_sequence_length)
position_embeddings = self._position_embedding_layer(word_embeddings)
all_embeddings = [word_embeddings, position_embeddings]
if num_type_features:
type_embeddings = [(layers.OnDeviceEmbedding(
vocab_size=type_vocab_sizes[idx],
embedding_width=hidden_size,
initializer=initializer,
use_one_hot=True,
name='type_embeddings_{}'.format(idx))(type_ids[..., idx]))
for idx in range(num_type_features)]
all_embeddings += type_embeddings
if num_float_features:
float_embeddings = [
(
tf.keras.layers.Dense(
hidden_size, name='float_features_{}'.format(idx))(
# Expanding the last dim here is important.
float_features[..., idx, None]))
for idx in range(num_float_features)
]
all_embeddings += float_embeddings
embeddings = tf.keras.layers.Add()(all_embeddings)
embeddings = (tf.keras.layers.LayerNormalization(
name='embeddings/layer_norm',
axis=-1,
epsilon=1e-12,
dtype=tf.float32)(embeddings))
embeddings = (tf.keras.layers.Dropout(rate=dropout_rate)(embeddings))
self._transformer_layers = []
data = embeddings
attention_mask = layers.SelfAttentionMask()([data, mask])
encoder_outputs = []
for i in range(num_layers):
layer = layers.Transformer(
num_attention_heads=num_attention_heads,
intermediate_size=intermediate_size,
intermediate_activation=activation,
dropout_rate=dropout_rate,
attention_dropout_rate=attention_dropout_rate,
kernel_initializer=initializer,
name='model/layer_with_weights-%d' % (i + 4))
self._transformer_layers.append(layer)
data = layer([data, attention_mask])
encoder_outputs.append(data)
first_token_tensor = (
tf.keras.layers.Lambda(lambda x: tf.squeeze(x[:, 0:1, :], axis=1))(
encoder_outputs[-1]))
cls_output = tf.keras.layers.Dense(
units=hidden_size,
activation='tanh',
kernel_initializer=initializer,
name='pooler_transform')(first_token_tensor)
if return_all_encoder_outputs:
outputs = [encoder_outputs, cls_output]
else:
outputs = [encoder_outputs[-1], cls_output]
super(TransformerEncoder, self).__init__(inputs=all_inputs,
outputs=outputs,
**kwargs)
if bert_init_ckpt and learner_flags.INIT_CHECKPOINT.value is None:
self.init_weights(bert_init_ckpt)
def __init__(
self,
vocab_size,
hidden_size=768,
num_layers=12,
num_attention_heads=12,
sequence_length=512,
max_sequence_length=None,
type_vocab_size=16,
intermediate_size=3072,
activation=activations.gelu,
dropout_rate=0.1,
attention_dropout_rate=0.1,
initializer=tf.keras.initializers.TruncatedNormal(stddev=0.02),
return_all_encoder_outputs=False,
output_range=None,
embedding_width=None,
**kwargs):
activation = tf.keras.activations.get(activation)
initializer = tf.keras.initializers.get(initializer)
if not max_sequence_length:
max_sequence_length = sequence_length
self._self_setattr_tracking = False
self._config_dict = {
'vocab_size': vocab_size,
'hidden_size': hidden_size,
'num_layers': num_layers,
'num_attention_heads': num_attention_heads,
'sequence_length': sequence_length,
'max_sequence_length': max_sequence_length,
'type_vocab_size': type_vocab_size,
'intermediate_size': intermediate_size,
'activation': tf.keras.activations.serialize(activation),
'dropout_rate': dropout_rate,
'attention_dropout_rate': attention_dropout_rate,
'initializer': tf.keras.initializers.serialize(initializer),
'return_all_encoder_outputs': return_all_encoder_outputs,
'output_range': output_range,
'embedding_width': embedding_width,
}
word_ids = tf.keras.layers.Input(shape=(sequence_length, ),
dtype=tf.int32,
name='input_word_ids')
mask = tf.keras.layers.Input(shape=(sequence_length, ),
dtype=tf.int32,
name='input_mask')
type_ids = tf.keras.layers.Input(shape=(sequence_length, ),
dtype=tf.int32,
name='input_type_ids')
if embedding_width is None:
embedding_width = hidden_size
self._embedding_layer = layers.OnDeviceEmbedding(
vocab_size=vocab_size,
embedding_width=embedding_width,
initializer=initializer,
name='word_embeddings')
word_embeddings = self._embedding_layer(word_ids)
# Always uses dynamic slicing for simplicity.
self._position_embedding_layer = layers.PositionEmbedding(
initializer=initializer,
use_dynamic_slicing=True,
max_sequence_length=max_sequence_length,
name='position_embedding')
position_embeddings = self._position_embedding_layer(word_embeddings)
self._type_embedding_layer = layers.OnDeviceEmbedding(
vocab_size=type_vocab_size,
embedding_width=embedding_width,
initializer=initializer,
use_one_hot=True,
name='type_embeddings')
type_embeddings = self._type_embedding_layer(type_ids)
embeddings = tf.keras.layers.Add()(
[word_embeddings, position_embeddings, type_embeddings])
embeddings = (tf.keras.layers.LayerNormalization(
name='embeddings/layer_norm',
axis=-1,
epsilon=1e-12,
dtype=tf.float32)(embeddings))
embeddings = (tf.keras.layers.Dropout(rate=dropout_rate)(embeddings))
# We project the 'embedding' output to 'hidden_size' if it is not already
# 'hidden_size'.
if embedding_width != hidden_size:
self._embedding_projection = tf.keras.layers.experimental.EinsumDense(
'...x,xy->...y',
output_shape=hidden_size,
bias_axes='y',
kernel_initializer=initializer,
name='embedding_projection')
embeddings = self._embedding_projection(embeddings)
self._transformer_layers = []
data = embeddings
attention_mask = layers.SelfAttentionMask()([data, mask])
encoder_outputs = []
for i in range(num_layers):
if i == num_layers - 1 and output_range is not None:
transformer_output_range = output_range
else:
transformer_output_range = None
layer = layers.Transformer(
num_attention_heads=num_attention_heads,
intermediate_size=intermediate_size,
intermediate_activation=activation,
dropout_rate=dropout_rate,
attention_dropout_rate=attention_dropout_rate,
output_range=transformer_output_range,
kernel_initializer=initializer,
name='transformer/layer_%d' % i)
self._transformer_layers.append(layer)
data = layer([data, attention_mask])
encoder_outputs.append(data)
first_token_tensor = (
tf.keras.layers.Lambda(lambda x: tf.squeeze(x[:, 0:1, :], axis=1))(
encoder_outputs[-1]))
self._pooler_layer = tf.keras.layers.Dense(
units=hidden_size,
activation='tanh',
kernel_initializer=initializer,
name='pooler_transform')
cls_output = self._pooler_layer(first_token_tensor)
if return_all_encoder_outputs:
outputs = [encoder_outputs, cls_output]
else:
outputs = [encoder_outputs[-1], cls_output]
super(TransformerEncoder,
self).__init__(inputs=[word_ids, mask, type_ids],
outputs=outputs,
**kwargs)
