Update Model.py

Model.py

@@ -124,7 +124,7 @@ class Lo(layers.Layer):
     def __init__(self):
         super().__init__()
         # 내부 계산은 float32로 유지
-        self.p = layers.Dense(48, use_bias=True, dtype='float32')
+        self.p = layers.Dense(64, use_bias=True, dtype='float32')
         self._out_dtype = 'float32'
 
     def call(self, x):
@@ -137,26 +137,22 @@ class Lo(layers.Layer):
 class rGLU(layers.Layer):
     def __init__(self, d_model, hyper_n):
         super().__init__()
-        self.Wr = layers.Dense(48)
-        self.WB = layers.Dense(768)
-        self.Wr1 = layers.Dense(48)
+        self.Wr = Lo()
+        self.W2 = layers.Dense(256)
+        self.W1 = layers.Dense(256)
+        self.Wr1 = Lo()
         self.W = layers.Dense(d_model)
     def call(self, x):
-        x = self.Wr(x)
-        x = self.WB(x)
-        a, b = tf.split(x, 2, axis=-1)
-        o = tf.nn.silu(a) * b
-        o = self.Wr1(o)
-        o = self.W(o)
+        x = tf.nn.silu(self.W1(Wr(x)) + x) * (self.W2(self.Wr1(x)) + x)
         return o
 
 class Adapter(layers.Layer):
     def __init__(self, d_model, hyper_n):
         super().__init__()
-        self.Wr = layers.Dense(48, activation='gelu')
+        self.Wr = Lo()
         self.W = layers.Dense(d_model)
     def call(self, x):
-        return self.W(self.Wr(x))
+        return self.W(tf.nn.gelu(self.Wr(x)))
 
 class LoSoU(layers.Layer):
     """
@@ -237,9 +233,9 @@ class LoSoU(layers.Layer):
         residual = x_f32
 
         # Q, K, V
-        q = self.Q(self.Qr(x_f32))  # (B, L, 96)
-        k = self.K(self.Kr(x_f32))   # (B, L, 96)
-        V = tf.cast(self.V(self.Vr(x)), tf.float32)  # ensure V's output is float32
+        q = self.Q(self.Qr(x_f32)) + x_f32  # (B, L, 96)
+        k = self.K(self.Kr(x_f32)) + x_f32  # (B, L, 96)
+        V = self.V(self.Vr(x)) + x # ensure V's output is float32
 
         # gating signals in (0,1)
         g_q = tf.nn.sigmoid(q)
@@ -333,7 +329,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=256,
+    d_model=192,
     n_layers=1
 )