Update Model.py

Model.py

@@ -68,7 +68,7 @@ unk_id = sp.piece_to_id("<unk>")
 vocab_size = sp.get_piece_size()
 print(f"✅ Vocabulary size: {vocab_size}")
 
-max_len = 200
+max_len = 512
 batch_size = 128
 
 def text_to_ids(text):
@@ -156,7 +156,7 @@ class Lo(layers.Layer):
         # cast back to model dtype for consistency
         return tf.cast(x, self._out_dtype)
 
-class LoSoU(layers.Layer):
+class LoU(layers.Layer):
     """
     안정화된 LoSoU 레이어 (동적 alpha 사용)
     - alpha 값을 입력에 따라 동적으로 계산: alpha = sigmoid(Linear(x))
@@ -182,6 +182,8 @@ class LoSoU(layers.Layer):
         self.proj = layers.Dense(d_model, use_bias=True, dtype='float32')
         self.O = layers.Dense(d_model, dtype='float32')
         self.norm = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
+        self.norm1 = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
+        
         self.alpha_linear = layers.Dense(1, activation='sigmoid', dtype='float32')
 
     def _ema_over_time(self, score, alpha_dynamic):
@@ -222,6 +224,7 @@ class LoSoU(layers.Layer):
         # cast to float32 for all internal computations
         x_f32 = tf.cast(x, tf.float32)
         residual = x_f32
+        x_f32 = self.norm1(x)
 
         # Q, K, V
         q = self.Q(x_f32)
@@ -274,22 +277,12 @@ class LoSoU(layers.Layer):
         # cast back to original dtype for downstream layers
         return tf.cast(out, x.dtype)
 
-class Block(layers.Layer):
-    def __init__(self, d_model, hyper_n):
-        super().__init__()
-        self.losou = [LoSoU(d_model) for _ in range(hyper_n)]
-
-    def call(self, x):
-        for losou in self.losou:
-            x = losou(x)
-        return x
-
 class ReLaM(tf.keras.Model):
     def __init__(self, vocab_size, max_seq_len, d_model, n_layers, dropout_rate=0.1):
         super().__init__()
-        self.token_embedding = layers.Embedding(vocab_size, 128)
+        self.token_embedding = layers.Embedding(vocab_size, d_model)
         self.pos_embedding = layers.Embedding(max_seq_len, d_model)
-        self.blocks = [Block(d_model, hyper_n=1) for _ in range(n_layers)]
+        self.blocks = [LoU(d_model) for _ in range(n_layers)]
 
         # LayerNormalization은 float32로 해서 정밀도 문제 방지
         self.ln_f = layers.LayerNormalization(epsilon=1e-5, dtype="float32")
@@ -334,8 +327,8 @@ def create_lr_schedule(initial_lr=5e-5, decay_steps=10000, decay_rate=0.9):
 model = ReLaM(
     vocab_size=vocab_size,
     max_seq_len=max_len,
-    d_model=256,
-    n_layers=1
+    d_model=512,
+    n_layers=16
 )
 
 # 옵티마이저 설정