Update AlphaS2S.py

AlphaS2S.py

@@ -243,8 +243,8 @@ class LoU(layers.Layer):
         out = self.glu(out)
         return tf.cast(out, x.dtype)
         
-class Transformer(tf.keras.Model):
-    def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, target_vocab_size, max_len=100, dropout=0.1):
+class AlphaS2S(tf.keras.Model):
+    def __init__(self, num_layers, d_model, num_heads, input_vocab_size, target_vocab_size, max_len=100, dropout=0.1):
         super().__init__()
         self.max_len = max_len
         self.d_model = d_model
@@ -252,9 +252,9 @@ class Transformer(tf.keras.Model):
         self.enc_pos_embedding = layers.Embedding(max_len, d_model)
         self.dec_embedding = layers.Embedding(target_vocab_size, d_model)
         self.dec_pos_embedding = layers.Embedding(max_len, d_model)
-        self.enc_layers = [EncoderBlock(d_model, num_heads, dff, dropout) for _ in range(num_layers)]
+        self.enc_layers = [EncoderBlock(d_model, num_heads, dropout) for _ in range(num_layers)]
         self.dec_layers = [LoU(d_model) for _ in range(num_layers)]
-        self.final_layer = layers.Dense(target_vocab_size)
+        self.final_layer = layers.Dense(target_vocab_size, use_bias=False)
     def call(self, inputs, training=False):
         enc_inputs = inputs["enc_inputs"]
         dec_inputs = inputs["dec_inputs"]
@@ -266,3 +266,12 @@ class Transformer(tf.keras.Model):
         y = self.dec_embedding(dec_inputs) + self.dec_pos_embedding(dec_pos)
         for layer in self.dec_layers: y = layer(y, enc_out, training=training)
         return self.final_layer(y)
+
+
+chat_model = AlphaS2S(num_layers=4, d_model=160, num_heads=8,
+                         input_vocab_size=chat_vocab_size, target_vocab_size=chat_vocab_size)
+dummy_input = {
+    "enc_inputs": tf.zeros((1, max_len), dtype=tf.int32),
+    "dec_inputs": tf.zeros((1, max_len), dtype=tf.int32)
+}
+_ = chat_model(dummy_input)