def build(self, input_shape):
"""Builds the layer."""
# Layers for linearly projecting the queries, keys, and values.
size_per_head = self.hidden_size // self.num_heads
self.query_dense_layer = layers.DenseEinsum(
output_shape=(self.num_heads, size_per_head),
kernel_initializer="glorot_uniform",
use_bias=False,
name="query")
self.key_dense_layer = layers.DenseEinsum(
output_shape=(self.num_heads, size_per_head),
kernel_initializer="glorot_uniform",
use_bias=False,
name="key")
self.value_dense_layer = layers.DenseEinsum(
output_shape=(self.num_heads, size_per_head),
kernel_initializer="glorot_uniform",
use_bias=False,
name="value")
self.output_dense_layer = layers.DenseEinsum(
output_shape=self.hidden_size,
num_summed_dimensions=2,
kernel_initializer="glorot_uniform",
use_bias=False,
name="output_transform")
super(Attention, self).build(input_shape)
def build(self, input_shape):
# Self attention.
self.self_attention = layers.CachedAttention(
num_heads=self.num_attention_heads,
key_size=self.attention_head_size,
dropout_rate=self.attention_probs_dropout_prob,
kernel_initializer=self._kernel_initializer,
name="self_attention")
self.self_attention_output_dense = layers.DenseEinsum(
output_shape=self.hidden_size,
num_summed_dimensions=2,
kernel_initializer=self._kernel_initializer,
bias_initializer=self._bias_initializer,
name="self_attention_output")
self.self_attention_dropout = tf.keras.layers.Dropout(
rate=self.hidden_dropout_prob)
self.self_attention_layer_norm = (tf.keras.layers.LayerNormalization(
name="self_attention_layer_norm", axis=-1, epsilon=1e-12))
# Encoder-decoder attention.
self.encdec_attention = self._cross_attention_cls(
num_heads=self.num_attention_heads,
key_size=self.attention_head_size,
dropout_rate=self.attention_probs_dropout_prob,
output_shape=self.hidden_size,
kernel_initializer=self._kernel_initializer,
name="attention/encdec")
# TODO(hongkuny): Remove when checkpoint backward compatibility is resolved.
# pylint: disable=protected-access
self.self_attention.build(input_shape)
self.self_attention_output_dense = self.self_attention._output_dense
self.encdec_attention.build(input_shape)
self.encdec_attention_output_dense = self.encdec_attention._output_dense
self.encdec_attention_dropout = tf.keras.layers.Dropout(
rate=self.hidden_dropout_prob)
self.encdec_attention_layer_norm = (tf.keras.layers.LayerNormalization(
name="attention/encdec_output_layer_norm", axis=-1, epsilon=1e-12))
# Feed-forward projection.
self.intermediate_dense = layers.DenseEinsum(
output_shape=self.intermediate_size,
activation=None,
kernel_initializer=self._kernel_initializer,
bias_initializer=self._bias_initializer,
name="intermediate")
self.intermediate_activation_layer = tf.keras.layers.Activation(
self.intermediate_activation)
self.output_dense = layers.DenseEinsum(
output_shape=self.hidden_size,
kernel_initializer=self._kernel_initializer,
bias_initializer=self._bias_initializer,
name="output")
self.output_dropout = tf.keras.layers.Dropout(
rate=self.hidden_dropout_prob)
self.output_layer_norm = tf.keras.layers.LayerNormalization(
name="output_layer_norm", axis=-1, epsilon=1e-12)
super(TransformerDecoderBlock, self).build(input_shape)
def build(self, unused_input_shapes):
# Self attention.
self.self_attention = layers.CachedAttention(
num_heads=self.num_attention_heads,
head_size=self.attention_head_size,
dropout_rate=self.attention_probs_dropout_prob,
kernel_initializer=self._kernel_initializer,
name="self_attention")
self.self_attention_output_dense = layers.DenseEinsum(
output_shape=self.hidden_size,
num_summed_dimensions=2,
kernel_initializer=self._kernel_initializer,
bias_initializer=self._bias_initializer,
name="self_attention_output")
self.self_attention_dropout = tf.keras.layers.Dropout(
rate=self.hidden_dropout_prob)
self.self_attention_layer_norm = (tf.keras.layers.LayerNormalization(
name="self_attention_layer_norm", axis=-1, epsilon=1e-12))
# Encoder-decoder attention.
self.encdec_attention = self._cross_attention_cls(
num_heads=self.num_attention_heads,
head_size=self.attention_head_size,
dropout_rate=self.attention_probs_dropout_prob,
kernel_initializer=self._kernel_initializer,
name="attention/encdec")
self.encdec_attention_output_dense = layers.DenseEinsum(
output_shape=self.hidden_size,
num_summed_dimensions=2,
kernel_initializer=self._kernel_initializer,
bias_initializer=self._bias_initializer,
name="attention/encdec_output")
self.encdec_attention_dropout = tf.keras.layers.Dropout(
rate=self.hidden_dropout_prob)
self.encdec_attention_layer_norm = (tf.keras.layers.LayerNormalization(
name="attention/encdec_output_layer_norm", axis=-1, epsilon=1e-12))
# Feed-forward projection.
self.intermediate_dense = layers.DenseEinsum(
output_shape=self.intermediate_size,
activation=None,
kernel_initializer=self._kernel_initializer,
bias_initializer=self._bias_initializer,
name="intermediate")
self.intermediate_activation_layer = tf.keras.layers.Activation(
self.intermediate_activation)
self.output_dense = layers.DenseEinsum(
output_shape=self.hidden_size,
kernel_initializer=self._kernel_initializer,
bias_initializer=self._bias_initializer,
name="output")
self.output_dropout = tf.keras.layers.Dropout(
rate=self.hidden_dropout_prob)
self.output_layer_norm = tf.keras.layers.LayerNormalization(
name="output_layer_norm", axis=-1, epsilon=1e-12)
super(TransformerDecoderBlock, self).build(unused_input_shapes)
def build(self, input_shape):
"""Builds the layer."""
# Layers for linearly projecting the queries, keys, and values.
size_per_head = self.hidden_size // self.num_heads
def _glorot_initializer(fan_in, fan_out):
limit = math.sqrt(6.0 / (fan_in + fan_out))
return tf.keras.initializers.RandomUniform(minval=-limit,
maxval=limit)
attention_initializer = _glorot_initializer(input_shape.as_list()[-1],
self.hidden_size)
self.query_dense_layer = layers.DenseEinsum(
output_shape=(self.num_heads, size_per_head),
kernel_initializer=attention_initializer,
use_bias=False,
name="query")
self.key_dense_layer = layers.DenseEinsum(
output_shape=(self.num_heads, size_per_head),
kernel_initializer=attention_initializer,
use_bias=False,
name="key")
self.value_dense_layer = layers.DenseEinsum(
output_shape=(self.num_heads, size_per_head),
kernel_initializer=attention_initializer,
use_bias=False,
name="value")
output_initializer = _glorot_initializer(self.hidden_size,
self.hidden_size)
self.output_dense_layer = layers.DenseEinsum(
output_shape=self.hidden_size,
num_summed_dimensions=2,
kernel_initializer=output_initializer,
use_bias=False,
name="output_transform")
super(Attention, self).build(input_shape)
def build(self, unused_input_shapes):
self._query_dense = layers.DenseEinsum(
output_shape=(self._num_heads, self._head_size),
kernel_initializer=self._kernel_initializer,
bias_initializer=self._bias_initializer,
kernel_regularizer=self._kernel_regularizer,
bias_regularizer=self._bias_regularizer,
activity_regularizer=self._activity_regularizer,
kernel_constraint=self._kernel_constraint,
bias_constraint=self._bias_constraint,
dtype=self.dtype,
name="encdocatt_query")
self._key_dense = layers.DenseEinsum(
output_shape=(self._num_heads, self._head_size),
kernel_initializer=self._kernel_initializer,
bias_initializer=self._bias_initializer,
kernel_regularizer=self._kernel_regularizer,
bias_regularizer=self._bias_regularizer,
activity_regularizer=self._activity_regularizer,
kernel_constraint=self._kernel_constraint,
bias_constraint=self._bias_constraint,
dtype=self.dtype,
name="encdocatt_key")
super(DocAttention, self).build(unused_input_shapes)
def __init__(
self,
vocab_size,
embedding_width=128,
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),
**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,
'embedding_width': embedding_width,
'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),
}
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=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)
position_embeddings = self._position_embedding_layer(word_embeddings)
type_embeddings = (layers.OnDeviceEmbedding(
vocab_size=type_vocab_size,
embedding_width=embedding_width,
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)(embeddings))
# We project the 'embedding' output to 'hidden_size' if it is not already
# 'hidden_size'.
if embedding_width != hidden_size:
embeddings = layers.DenseEinsum(
output_shape=hidden_size,
kernel_initializer=initializer,
name='embedding_projection')(embeddings)
data = embeddings
attention_mask = layers.SelfAttentionMask()([data, mask])
shared_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='transformer')
for _ in range(num_layers):
data = shared_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(AlbertTransformerEncoder,
self).__init__(inputs=[word_ids, mask, type_ids],
outputs=[data, cls_output],
**kwargs)
