Update AlphaS2S.py

AlphaS2S.py

@@ -13,7 +13,7 @@ tf.get_logger().setLevel("ERROR")
 SEED = 42
 tf.random.set_seed(SEED)
 np.random.seed(SEED)
-max_len = 224 # 기존 코드에서 200으로 설정됨
+max_len = 220 # 기존 코드에서 200으로 설정됨
 batch_size = 48
 
 # TPU 초기화 (기존 코드와 동일)
@@ -166,13 +166,11 @@ 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)
@@ -184,7 +182,6 @@ 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)
@@ -192,7 +189,6 @@ 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)
@@ -205,10 +201,10 @@ 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, 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, 256)
+        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.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)]
         self.final_layer = layers.Dense(target_vocab_size, use_bias=False)
@@ -259,7 +255,7 @@ def create_lr_schedule(initial_lr=5e-5, decay_steps=10000, decay_rate=0.9):
 
 with strategy.scope():
     # ⚠️ 수정: chat_vocab_size 대신 정의된 vocab_size 사용
-    chat_model = Transformer(num_layers=2, d_model=320, num_heads=4, dff=960, input_vocab_size=vocab_size, target_vocab_size=vocab_size, max_len=256, dropout=0.1)
+    chat_model = Transformer(num_layers=2, d_model=304, num_heads=4, dff=912, input_vocab_size=vocab_size, target_vocab_size=vocab_size, max_len=256, dropout=0.1)
     
     dummy_input = {
         "enc_inputs": tf.zeros((1, max_len), dtype=tf.int32),