text: import tensorflow as tf
  
  
  # Encoder
  
  class Encoder(tf.keras.Model):
      def __init__(self, vocab_szie, embedding_dim, enc_units, batch_sz):
          super(Encoder, self).__init__()
          self.batch_sz = batch_sz
          self.enc_units = enc_units
          self.embedding = tf.keras.layers.Embedding(vocab_szie, embedding_dim)
          self.gru = tf.keras.layers.GRU(self.enc_units,
                                         return_sequences=True,
                                         return_state=True,
                                         recurrent_initializer='glorot_uniform')
  
      def call(self, inputs):
          input, hidden = inputs
          x = self.embedding(input)
          output, state = self.gru(x, initial_state=hidden)
          return output, state
  
      def initialize_hidden_state(self):
          return tf.zeros((self.batch_sz, self.enc_units))
  
  
  # Attention layer
  class BahdanauAttention(tf.keras.layers.Layer):
      def __init__(self, units):
          super(BahdanauAttention, self).__init__()
          self.W1 = tf.keras.layers.Dense(units)
          self.W2 = tf.keras.layers.Dense(units)
          self.V = tf.keras.layers.Dense(1)
  
      def call(self, inputs):
          query, values = inputs
          hidden_with_time_axis = tf.expand_dims(query, 1)
  
          score = self.V(tf.nn.tanh(self.W1(values) + self.W2(hidden_with_time_axis)))
  
          attention_weights = tf.nn.softmax(score, axis=1)
  
          contex_vector = attention_weights * values
          contex_vector = tf.reduce_sum(contex_vector, axis=1)
  
          return contex_vector, attention_weights
  
  
  # Decoder
  class Decoder(tf.keras.Model):
      def __init__(self, vocab_size, embedding_dim, dec_units, batch_sz):
          super(Decoder, self).__init__()
          self.batch_sz = batch_sz
          self.dec_units = dec_units
          self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
          self.gru = tf.keras.layers.GRU(self.dec_units,
                                         return_sequences=True,
                                         return_state=True,
                                         recurrent_initializer='glorot_uniform')
          self.fc = tf.keras.layers.Dense(vocab_size)
  
          self.attention = BahdanauAttention(self.dec_units)
  
      def call(self, inputs):
          x, hidden, enc_output = inputs
          context_vector, attention_weights = self.attention((hidden, enc_output))
          x = self.embedding(x)
          x = tf.concat([tf.expand_dims(context_vector, 1), x], axis=-1)
  
          output, state = self.gru(x)
          output = tf.reshape(output, (-1, output.shape[2]))
          x = self.fc(output)
  
          return x, state, attention_weights
  