Update Model.py

Model.py

@@ -112,21 +112,6 @@ dataset = dataset.shuffle(2000, seed=SEED).batch(batch_size, drop_remainder=True
 with strategy.scope():
     dist_dataset = strategy.experimental_distribute_dataset(dataset)
 
-
-class Lo(layers.Layer):
-    def __init__(self, d_model):
-        super().__init__()
-        self.proj = layers.Dense(d_model, use_bias=True, dtype='float32')
-        self.p = layers.Dense(128, use_bias=True, dtype='float32')
-        self._out_dtype = 'float32'
-
-    def call(self, x):
-        x_f32 = tf.cast(x, tf.float32)
-        x = self.proj(x_f32)
-        x = tf.nn.gelu(x)
-        x = self.p(x)
-        return tf.cast(x, self._out_dtype)
-
 class SwiGLU(layers.Layer):
     def __init__(self, d_model):
         super().__init__()
@@ -148,9 +133,6 @@ class LoU(layers.Layer):
         self.Q = layers.Dense(d_model, dtype='float32')
         self.K = layers.Dense(d_model, dtype='float32')
         self.V = layers.Dense(d_model, dtype='float32')
-        self.Qr = Lo(d_model)
-        self.Kr = Lo(d_model)
-        self.Vr = Lo(d_model)
         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')
@@ -185,10 +167,6 @@ class LoU(layers.Layer):
         q = self.Q(x_f32)
         k = self.K(x_f32)
         V = self.V(x_f32)
-        q = self.Qr(q)
-        k = self.Kr(k)
-        V = self.Vr(V)
-
         # 기존 코드:
         # g_q = tf.nn.sigmoid(q)
         # g_k = tf.nn.sigmoid(k)
@@ -208,16 +186,28 @@ class LoU(layers.Layer):
         out = self.norm(out + residual)
         return tf.cast(out, x.dtype)
 
+class Lo(layers.Layer):
+    def __init__(self, d_model):
+        super().__init__()
+        self.d = layers.Dense(256, activation='silu')
+        self.w = layers.Dense(d_model)
+    def call(self, x):
+        p = self.d(x)
+        p = self.w(p)
+        return p + x
+
 class Block(layers.Layer):
     def __init__(self, d_model):
         super().__init__()
         self.lou = LoU(d_model)
         self.glu = SwiGLU(d_model)
         self.norm = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
+        self.lo = Lo(d_model)
 
     def call(self, x):
         x = self.lou(x)
         x = self.norm(self.glu(x)) + x
+        x = self.lo(x)
         return x
 
 class ReLM(tf.keras.Model):