| import tensorflow as tf
|
| from .encoder import Encoder
|
| from .decoder import Decoder
|
| from tensorflow.keras.layers import Dense
|
|
|
| @tf.keras.utils.register_keras_serializable()
|
| class Transformer(tf.keras.Model):
|
| def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size,
|
| target_vocab_size, max_tokens, dropout_rate=0.1, **kwargs):
|
| super(Transformer, self).__init__(**kwargs)
|
| self.num_layers = num_layers
|
| self.d_model = d_model
|
| self.num_heads = num_heads
|
| self.dff = dff
|
| self.input_vocab_size = input_vocab_size
|
| self.target_vocab_size = target_vocab_size
|
| self.max_tokens = max_tokens
|
| self.dropout_rate = dropout_rate
|
|
|
| self.encoder = Encoder(num_layers, d_model, num_heads, dff,
|
| input_vocab_size, max_tokens, dropout_rate)
|
| self.decoder = Decoder(num_layers, d_model, num_heads, dff,
|
| target_vocab_size, max_tokens, dropout_rate)
|
| self.final_layer = Dense(target_vocab_size)
|
|
|
| def call(self, inputs, training=None):
|
| enc_input, dec_input = inputs
|
| enc_padding_mask = self.create_padding_mask(enc_input)
|
| look_ahead_mask = self.create_look_ahead_mask(tf.shape(dec_input)[1])
|
| dec_padding_mask = self.create_padding_mask(enc_input)
|
| enc_output = self.encoder(enc_input, training=training, mask=enc_padding_mask)
|
| dec_output = self.decoder(dec_input, enc_output, training=training,
|
| look_ahead_mask=look_ahead_mask,
|
| padding_mask=dec_padding_mask)
|
| final_output = self.final_layer(dec_output)
|
| return final_output
|
|
|
| def create_padding_mask(self, seq):
|
| mask = tf.cast(tf.math.equal(seq, 0), tf.float32)
|
| return mask[:, tf.newaxis, tf.newaxis, :]
|
|
|
| def create_look_ahead_mask(self, size):
|
| mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)
|
| return mask
|
|
|
| def get_config(self):
|
| config = super().get_config()
|
| config.update({
|
| 'num_layers': self.num_layers,
|
| 'd_model': self.d_model,
|
| 'num_heads': self.num_heads,
|
| 'dff': self.dff,
|
| 'input_vocab_size': self.input_vocab_size,
|
| 'target_vocab_size': self.target_vocab_size,
|
| 'max_tokens': self.max_tokens,
|
| 'dropout_rate': self.dropout_rate
|
| })
|
| return config |
