Update Model.py

Model.py

@@ -217,7 +217,7 @@ class LoSoU(layers.Layer):
 
         # gating signals in (0,1)
         g_q = tf.nn.sigmoid(q)
-        g_k = tf.nn.sigmoid(k)
+        g_k = tf.nn.tanh(k)
 
         # elementwise product -> bounded roughly [0,1]
         score = g_q * g_k
@@ -273,7 +273,7 @@ class ReLaM(tf.keras.Model):
         self.token_embedding = layers.Embedding(vocab_size, 128)
         self.pos_embedding = layers.Embedding(max_seq_len, d_model)
         self.blocks = [Block(d_model, hyper_n=1) for _ in range(n_layers)]
-
+        self.proj = layers.Dense(128)
         # LayerNormalization은 float32로 해서 정밀도 문제 방지
         self.ln_f = layers.LayerNormalization(epsilon=1e-5, dtype="float32")
 
@@ -284,7 +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.proj(x)
         x = self.ln_f(x)
 
         embedding_matrix = tf.cast(self.token_embedding.embeddings, x.dtype)