Update Model.py

Model.py

@@ -120,13 +120,22 @@ dataset = dataset.shuffle(1000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
 
 print("✅ TF Dataset 생성 완료!")
 
-class Adapter(layers.Layer):
+class Lo(layers.Layer):
     def __init__(self, d_model):
         super().__init__()
-        self.Wr = layers.Dense(64)
-        self.W = layers.Dense(d_model)
+        # 내부 계산은 float32로 유지
+        self.proj = layers.Dense(d_model, use_bias=True, dtype='float32')
+        self.p = layers.Dense(96, use_bias=True, dtype='float32')
+        self._out_dtype = 'float32'
+
     def call(self, x):
-        return self.W(tf.nn.gelu(self.Wr(x)))
+        # x may be bfloat16; cast to float32 for stable intermediate computation
+        x_f32 = tf.cast(x, tf.float32)
+        x = self.proj(x_f32)
+        x = tf.nn.gelu(x)
+        x = self.p(x)
+        # cast back to model dtype for consistency
+        return tf.cast(x, self._out_dtype)
 
 class LoSoU(layers.Layer):
     """
@@ -145,12 +154,12 @@ class LoSoU(layers.Layer):
         self.eps = float(eps)
 
         # projection / gating layers in float32
-        self.Q = layers.Dense(d_model, dtype='float32')
-        self.K = layers.Dense(d_model, dtype='float32')
-        self.adapter = Adapter(d_model)
+        self.Q = layers.Dense(96, dtype='float32')
+        self.K = layers.Dense(96, dtype='float32')
+        self.V = Lo(d_model)  # Lo already handles casting to model dtype; we'll cast back to float32
+        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.norm2 = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
 
         # 동적 alpha 계산을 위한 레이어
         # alpha는 [0, 1] 범위여야 하므로 sigmoid 사용
@@ -199,22 +208,22 @@ class LoSoU(layers.Layer):
         # x: (B, L, d_model) maybe bfloat16 or float32
         # cast to float32 for all internal computations
         x_f32 = tf.cast(x, tf.float32)
-        x_f32 = self.norm2(x_f32)
         residual = x_f32
 
         # Q, K, V
-        q = self.Q(x_f32)  # (B, L, 96)
-        k = self.K(x_f32)  # (B, L, 96)
+        q = self.Q(x_f32)   # (B, L, 96)
+        k = self.K(x_f32)   # (B, L, 96)
+        V = tf.cast(self.V(x), tf.float32)  # ensure V's output is float32
 
         # gating signals in (0,1)
         g_q = tf.nn.sigmoid(q)
-        g_k = tf.nn.tanh(k)
+        g_k = tf.nn.sigmoid(k)
 
         # elementwise product -> bounded roughly [0,1]
         score = g_q * g_k
 
         # 동적 alpha 계산: (B, L, d_model) -> (B, L, 1)
-        alpha_dynamic = self.alpha_linear(x_f32) # (B, L, 1)
+        alpha_dynamic = self.alpha_linear(x_f32) * 0.8 + 0.1 # (B, L, 1)
         # 필요시 alpha_dynamic에 대한 후처리 (예: min/max 등) 가능
         # ex: alpha_dynamic = tf.clip_by_value(alpha_dynamic, 0.01, 0.99)
 
@@ -230,9 +239,20 @@ class LoSoU(layers.Layer):
         score_clipped = tf.clip_by_value(score_norm, -self.clip_value, self.clip_value)
 
         # combine with V
-        x_comb = tf.nn.silu(score_clipped)  # (B, L, d_model)
-        out = self.norm(x_comb) + residual
-        out = self.norm1(self.adapter(out)) + out
+        x_comb = score_clipped * V  # (B, L, d_model)
+
+        out = self.proj(x_comb)  # (B, L, d_model)
+
+        # ensure out dim even for split
+        d = out.shape[-1]  # this is an int (static shape)
+        if d is not None and d % 2 == 1:
+            out = tf.pad(out, [[0,0],[0,0],[0,1]])
+
+        a, b = tf.split(out, 2, axis=-1)
+        gated = tf.nn.silu(a) * b
+        out = self.O(gated)
+
+        out = self.norm(out + residual)
 
         # cast back to original dtype for downstream layers
         return tf.cast(out, x.dtype)
@@ -251,8 +271,10 @@ 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.pos_embedding = layers.Embedding(max_seq_len, 128)
+        self.pos_embedding = layers.Embedding(max_seq_len, d_model)
         self.blocks = [Block(d_model, hyper_n=1) for _ in range(n_layers)]
+
+        # LayerNormalization은 float32로 해서 정밀도 문제 방지
         self.ln_f = layers.LayerNormalization(epsilon=1e-5, dtype="float32")
 
     def call(self, x, training=False):
@@ -262,6 +284,7 @@ class ReLaM(tf.keras.Model):
         x = self.token_embedding(x) + self.pos_embedding(positions)
         for block in self.blocks:
             x = block(x)
+
         x = self.ln_f(x)
 
         embedding_matrix = tf.cast(self.token_embedding.embeddings, x.dtype)
@@ -294,7 +317,7 @@ 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=128,
+    d_model=256,
     n_layers=1
 )
 
@@ -363,4 +386,4 @@ def generate_text_topp(model, prompt, max_len=100, max_gen=98, p=0.9, temperatur
     return ids_to_text(generated)
 
 print("\n\n===== 생성 결과 =====")  
-print(generate_text_topp(model, "제가 이따가 버스를 타야 해서 준비 좀 해야겠어요. 재미있는 대화였습니다!", p=0.9))
+print(generate_text_topp(model, "안녕", p=0.9))