def __init__(self, vocab_size, hidden_size, filter_size, num_heads,
num_encoder_layers, num_decoder_layers, label_smoothing,
dropout):
self._dtype = tf.float32
self._embedding_layer = embedding.Embedding(vocab_size, hidden_size,
"weights", self._dtype)
block_fn = lambda: transformer_block.TransformerBlock(
hidden_size, filter_size, num_heads, dropout)
self._encoder_layers = [block_fn() for _ in range(num_encoder_layers)]
self._decoder_layers = [block_fn() for _ in range(num_decoder_layers)]
self._dropout_fn = lambda x, training: tf.compat.v2.nn.dropout(
x, dropout, noise_shape=[x.shape[0], 1, x.shape[2]]) if training else x
self._vocab_size = vocab_size
self._num_heads = num_heads
self._label_smoothing = label_smoothing
self._decoder_scope_name = "decoder"
def test_embedding_layer_output(self):
embedding_layer = embedding.Embedding(12, 64, "test", tf.float32)
logits = embedding_layer(tf.ones((5, 7, 64)), False)
self.assertEqual(logits.shape, [5, 7, 12])
def test_embedding_layer_input(self):
embedding_layer = embedding.Embedding(12, 64, "test", tf.float32)
outputs = embedding_layer(tf.ones((5, 7), tf.int64), True)
self.assertEqual(outputs.shape, [5, 7, 64])