Update AlphaS2S.py

AlphaS2S.py

@@ -166,11 +166,13 @@ class EncoderBlock(layers.Layer):
         super().__init__()
         self.mha = layers.MultiHeadAttention(num_heads=num_heads, key_dim=d_model)
         self.ffn = SwiGLU(d_model, dff)
+        self.proj = layers.Dense(d_model)
         self.norm1 = layers.LayerNormalization(epsilon=1e-6)
         self.norm2 = layers.LayerNormalization(epsilon=1e-6)
         self.dropout1 = layers.Dropout(dropout)
         self.dropout2 = layers.Dropout(dropout)
     def call(self, x, mask=None, training=False):
+        x = self.proj(x)
         attn_out = self.dropout1(self.mha(x, x, x, attention_mask=mask), training=training)
         out1 = self.norm1(attn_out + x)
         ffn_out = self.dropout2(self.ffn(out1), training=training)
@@ -182,6 +184,7 @@ class DecoderBlock(layers.Layer):
         self.self_mha = layers.MultiHeadAttention(num_heads=num_heads, key_dim=d_model)
         self.cross_mha = layers.MultiHeadAttention(num_heads=num_heads, key_dim=d_model)
         self.ffn = SwiGLU(d_model, dff)
+        self.proj = layers.Dense(d_model)
         self.norm1 = layers.LayerNormalization(epsilon=1e-6)
         self.norm2 = layers.LayerNormalization(epsilon=1e-6)
         self.norm3 = layers.LayerNormalization(epsilon=1e-6)
@@ -189,6 +192,7 @@ class DecoderBlock(layers.Layer):
         self.dropout2 = layers.Dropout(dropout)
         self.dropout3 = layers.Dropout(dropout)
     def call(self, x, enc_out, training=False):
+        x = self.proj(x)
         attn1 = self.dropout1(self.self_mha(x, x, x, use_causal_mask=True), training=training)
         out1 = self.norm1(attn1 + x)
         attn2 = self.dropout2(self.cross_mha(out1, enc_out, enc_out), training=training)
@@ -201,9 +205,9 @@ class Transformer(tf.keras.Model):
         super().__init__()
         self.max_len = max_len
         self.d_model = d_model
-        self.enc_embedding = layers.Embedding(input_vocab_size, d_model)
-        self.enc_pos_embedding = layers.Embedding(max_len, d_model)
-        self.dec_embedding = layers.Embedding(target_vocab_size, d_model)
+        self.enc_embedding = layers.Embedding(input_vocab_size, 256)
+        self.enc_pos_embedding = layers.Embedding(max_len, 256)
+        self.dec_embedding = layers.Embedding(target_vocab_size, 256)
         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.dec_layers = [DecoderBlock(d_model, num_heads, dff, dropout) for _ in range(num_layers)]