text: #!/usr/bin/python3
  
  import os
  import time
  import numpy as np
  import tensorflow as tf
  from util import *
  from models import *
  
  
  def scaled_dot_product_attention(query, key, value, mask):
      matmul_qk = tf.matmul(query, key, transpose_b=True)
      depth = tf.cast(tf.shape(key)[-1], tf.float32)
      logits = matmul_qk / tf.math.sqrt(depth)
  
      if mask is not None:
          logits += (mask * -1e9)
  
      attention_weights = tf.nn.softmax(logits, axis=-1)
      output = tf.matmul(attention_weights, value)
  
      return output, attention_weights
  
  
  def create_padding_mask(x):
      mask = tf.cast(tf.math.equal(x, 0), tf.float32)
      return mask[:, tf.newaxis, tf.newaxis, :]
  
  
  def create_look_ahead_mask(x):
      seq_len = tf.shape(x)[1]
      look_ahead_mask = 1 - tf.linalg.band_part(tf.ones((seq_len, seq_len)), -1, 0)
      padding_mask = create_padding_mask(x)
      return tf.maximum(look_ahead_mask, padding_mask)
  
  
  class PositionalEncoding(tf.keras.layers.Layer):
      def __init__(self, position, d_model):
          super(PositionalEncoding, self).__init__()
          self.pos_encoding = self.positional_encoding(position, d_model)
  
      def get_angles(self, position, i, d_model):
          angles = 1 / tf.pow(10000, (2 * (i // 2)) / tf.cast(d_model, tf.float32))
          return position * angles
  
      def positional_encoding(self, position, d_model):
          angle_rads = self.get_angles(
              position=tf.range(position, dtype=tf.float32)[:, tf.newaxis],
              i=tf.range(d_model, dtype=tf.float32)[tf.newaxis, :],
              d_model=d_model)
  
          sines = tf.math.sin(angle_rads[:, 0::2])
          cosines = tf.math.cos(angle_rads[:, 1::2])
  
          angle_rads = np.zeros(angle_rads.shape)
          angle_rads[:, 0::2] = sines
          angle_rads[:, 1::2] = cosines
          pos_encoding = tf.constant(angle_rads)
          pos_encoding = pos_encoding[tf.newaxis, ...]
          # print(pos_encoding.shape)
          return tf.cast(pos_encoding, tf.float32)
  
      def call(self, inputs):
          return inputs + self.pos_encoding[:, :tf.shape(inputs)[1], :]
  
  
  class MultiHeadAttention(tf.keras.layers.Layer):
      def __init__(self, d_model, num_heads, name="multi_head_attention"):
          super(MultiHeadAttention, self).__init__(name=name)
          self.num_heads = num_heads
          self.d_model = d_model
          assert d_model % self.num_heads == 0
          self.depth = d_model // self.num_heads
          self.query_dense = tf.keras.layers.Dense(units=d_model)
          self.key_dense = tf.keras.layers.Dense(units=d_model)
          self.value_dense = tf.keras.layers.Dense(units=d_model)
          self.dense = tf.keras.layers.Dense(units=d_model)
  
      def split_heads(self, inputs, batch_size):
          inputs = tf.reshape(inputs, shape=(batch_size, -1, self.num_heads, self.depth))
          return tf.transpose(inputs, perm=[0, 2, 1, 3])
  
      def call(self, inputs):
          query, key, value, mask = inputs['query'], inputs['key'], inputs['value'], inputs['mask']
          batch_size = tf.shape(query)[0]
  
          query = self.query_dense(query)
          key = self.key_dense(key)
          value = self.value_dense(value)
  
          query = self.split_heads(query, batch_size)
          key = self.split_heads(key, batch_size)
          value = self.split_heads(value, batch_size)
  
          scaled_attention, _ = scaled_dot_product_attention(query, key, value, mask)
          scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])
          concat_attention = tf.reshape(scaled_attention, (batch_size, -1, self.d_model))
          outputs = self.dense(concat_attention)
  
          return outputs
  
  
  def encoder_layer(dff, d_model, num_heads, dropout, name="encoder_layer"):
      inputs = tf.keras.Input(shape=(None, d_model), name="inputs")
  
      padding_mask = tf.keras.Input(shape=(1, 1, None), name="padding_mask")
  
      attention = MultiHeadAttention(d_model, num_heads, name="attention")(
          {'query': inputs, 'key': inputs, 'value': inputs, 'mask': padding_mask})
  
      attention = tf.keras.layers.Dropout(rate=dropout)(attention)
      attention = tf.keras.layers.LayerNormalization(epsilon=1e-6)(inputs + attention)
  
      outputs = tf.keras.layers.Dense(units=dff, activation='relu')(attention)
      outputs = tf.keras.layers.Dense(units=d_model)(outputs)
  
      outputs = tf.keras.layers.Dropout(rate=dropout)(outputs)
      outputs = tf.keras.layers.LayerNormalization(epsilon=1e-6)(attention + outputs)
  
      return tf.keras.Model(inputs=[inputs, padding_mask], outputs=outputs, name=name)
  
  
  def encoder(vocab_size, num_layers, dff, d_model, num_heads, dropout, name="encoder"):
      inputs = tf.keras.Input(shape=(None,), name="inputs")
      padding_mask = tf.keras.Input(shape=(1, 1, None), name="padding_mask")
  
      embeddings = tf.keras.layers.Embedding(vocab_size, d_model)(inputs)
      embeddings *= tf.math.sqrt(tf.cast(d_model, tf.float32))
      embeddings = PositionalEncoding(vocab_size, d_model)(embeddings)
      outputs = tf.keras.layers.Dropout(rate=dropout)(embeddings)
  
      for i in range(num_layers):
          outputs = encoder_layer(dff=dff, d_model=d_model, num_heads=num_heads,
                                  dropout=dropout, name="encoder_layer_{}".format(i),
                                  )([outputs, padding_mask])
  
      return tf.keras.Model(inputs=[inputs, padding_mask], outputs=outputs, name=name)
  
  
  def decoder_layer(dff, d_model, num_heads, dropout, name="decoder_layer"):
      inputs = tf.keras.Input(shape=(None, d_model), name="inputs")
      enc_outputs = tf.keras.Input(shape=(None, d_model), name="encoder_outputs")
  
      look_ahead_mask = tf.keras.Input(shape=(1, None, None), name="look_ahead_mask")
      padding_mask = tf.keras.Input(shape=(1, 1, None), name='padding_mask')
  
      attention1 = MultiHeadAttention(d_model, num_heads, name="attention_1")(
          inputs={'query': inputs, 'key': inputs, 'value': inputs, 'mask': look_ahead_mask})
  
      attention1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)(attention1 + inputs)
  
      attention2 = MultiHeadAttention(d_model, num_heads, name="attention_2")(
          inputs={'query': attention1, 'key': enc_outputs, 'value': enc_outputs, 'mask': padding_mask})
  
      attention2 = tf.keras.layers.Dropout(rate=dropout)(attention2)
      attention2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)(attention2 + attention1)
  
      outputs = tf.keras.layers.Dense(units=dff, activation='relu')(attention2)
      outputs = tf.keras.layers.Dense(units=d_model)(outputs)
  
      outputs = tf.keras.layers.Dropout(rate=dropout)(outputs)
      outputs = tf.keras.layers.LayerNormalization(epsilon=1e-6)(outputs + attention2)
  
      return tf.keras.Model(inputs=[inputs, enc_outputs, look_ahead_mask, padding_mask],
                            outputs=outputs, name=name)
  
  
  def decoder(vocab_size, num_layers, dff, d_model, num_heads, dropout, name='decoder'):
      inputs = tf.keras.Input(shape=(None,), name='inputs')
      enc_outputs = tf.keras.Input(shape=(None, d_model), name='encoder_outputs')
  
      look_ahead_mask = tf.keras.Input(shape=(1, None, None), name='look_ahead_mask')
      padding_mask = tf.keras.Input(shape=(1, 1, None), name='padding_mask')
  
      embeddings = tf.keras.layers.Embedding(vocab_size, d_model)(inputs)
      embeddings *= tf.math.sqrt(tf.cast(d_model, tf.float32))
      embeddings = PositionalEncoding(vocab_size, d_model)(embeddings)
      outputs = tf.keras.layers.Dropout(rate=dropout)(embeddings)
  
      for i in range(num_layers):
          outputs = decoder_layer(dff=dff, d_model=d_model, num_heads=num_heads,
                                  dropout=dropout, name='decoder_layer_{}'.format(i),)(
              inputs=[outputs, enc_outputs, look_ahead_mask, padding_mask])
  
      return tf.keras.Model(inputs=[inputs, enc_outputs, look_ahead_mask, padding_mask],
                            outputs=outputs, name=name)
  
  
  def transformer(input_dim, num_layers, dff, d_model, num_heads, dropout, name="transformer"):
      inputs = tf.keras.Input(shape=(None,), name="inputs")
      dec_inputs = tf.keras.Input(shape=(None,), name="dec_inputs")
  
      enc_padding_mask = tf.keras.layers.Lambda(create_padding_mask,
                                                output_shape=(1, 1, None),
                                                name='enc_padding_mask')(inputs)
  
      look_ahead_mask = tf.keras.layers.Lambda(create_look_ahead_mask,
                                               output_shape=(1, None, None),
                                               name='look_ahead_mask')(dec_inputs)
  
      dec_padding_mask = tf.keras.layers.Lambda(create_padding_mask,
                                                output_shape=(1, 1, None),
                                                name='dec_padding_mask')(inputs)
  
      enc_outputs = encoder(vocab_size=input_dim, num_layers=num_layers, dff=dff,
                            d_model=d_model, num_heads=num_heads, dropout=dropout,)(
          inputs=[inputs, enc_padding_mask])
  
      dec_outputs = decoder(vocab_size=input_dim, num_layers=num_layers, dff=dff,
                            d_model=d_model, num_heads=num_heads, dropout=dropout,)(
          inputs=[dec_inputs, enc_outputs, look_ahead_mask, dec_padding_mask])
  
      outputs = tf.keras.layers.Dense(units=OUT_DIM, name="outputs")(dec_outputs)
  
      return tf.keras.Model(inputs=[inputs, dec_inputs], outputs=outputs, name=name)
  
  
  if __name__=='__main__':
      MAXLEN = 1000
      EMBEDDING_DIM = 256
      BATCH_SIZE = 16
      EPOCHS = 500
      TEST_RATIO = 0.0
      #MODE = 'train'  # train | test
      MODE = 'test'  # train | test
  
      input_tensor, target_tensor = load_dataset('../seq2seq/libtiff', pad_maxlen=MAXLEN)
      #input_tensor, target_tensor = load_dataset('../seq2seq/libpng', pad_maxlen=MAXLEN)
      #input_tensor, target_tensor = load_dataset('../seq2seq/init_dataset/PNG/path', pad_maxlen=MAXLEN)
      #input_tensor, target_tensor = load_dataset('../seq2seq/TIFF', pad_maxlen=MAXLEN)
      train_ds, test_ds = split_tensor(input_tensor, target_tensor,
                                       batch_size=BATCH_SIZE, algo='transformer', test_ratio=TEST_RATIO)
  
      model = transformer(input_dim=MAXLEN, num_layers=4, dff=512,
                          d_model=EMBEDDING_DIM, num_heads=4, dropout=0.3, name="transformer")
  
      dir_path = './model/transformer'
      filename = 'transformer_model_latest'
      fullname = os.path.join(dir_path, filename)
  
      if MODE == 'train':
          model = train_transformer_model(model, train_ds, EPOCHS, MAXLEN, early_stop_patience=2)
          if (os.path.isdir(dir_path) == False):
              os.mkdir(dir_path)
  
          cpfile = os.path.join(dir_path, 'checkpoint')
          if (os.path.isfile(cpfile) == True):
              filetime = time.strftime("_%Y%m%d-%H%M%S")
              newcp = cpfile + filetime
              fullname1 = fullname + '.index'
              fullname2 = fullname + '.data-00000-of-00001'
              newname1 = fullname + filetime + '.index'
              newname2 = fullname + filetime + '.data-00000-of-00001'
  
              os.rename(cpfile, newcp)
              os.rename(fullname1, newname1)
              os.rename(fullname2, newname2)
  
          model.save_weights(fullname)
      elif MODE == 'test':
          model.load_weights(fullname)
          #test_transformer_model(model, test_ds, MAXLEN, verbose=False, save='./output/PNG/transformer/')
          test_transformer_model(model, test_ds, MAXLEN, verbose=False, save='./output/transformer')
  