Update app.py

app.py

@@ -22,13 +22,13 @@ TOKENIZER_PATH = "bpe.model"
 
 if not os.path.exists(MODEL_PATH):
     download_file(
-        "https://huggingface.co/OpenLab-NLP/openlem1/resolve/main/encoder.weights.h5?download=true",
+        "https://huggingface.co/OpenLab-NLP/openlem-prototype/resolve/main/encoder_simcse.weights.h5?download=true",
         MODEL_PATH
     )
 
 if not os.path.exists(TOKENIZER_PATH):
     download_file(
-        "https://huggingface.co/OpenLab-NLP/openlem1/resolve/main/bpe.model?download=true",
+        "https://huggingface.co/OpenLab-NLP/openlem-prototype/resolve/main/bpe.model?download=true",
         TOKENIZER_PATH
     )
 
@@ -50,45 +50,72 @@ def encode_sentence(sentence, max_len=MAX_LEN):
 def pad_sentence(tokens):
     return tokens + [pad_id]*(MAX_LEN - len(tokens))
 
-class EncoderBlock(layers.Layer):
-    def __init__(self, embed_dim=EMBED_DIM, latent_dim=LATENT_DIM):
-        super().__init__()  # ✅ 반드시 맨 위에 추가
-        self.mha = layers.MultiHeadAttention(num_heads=8, key_dim=embed_dim//8)
-        self.WB = layers.Dense(1152)
-        self.W = layers.Dense(embed_dim)
-        self.ln1 = tf.keras.layers.LayerNormalization(epsilon=1e-5, dtype=tf.float32)
-        self.ln2 = tf.keras.layers.LayerNormalization(epsilon=1e-5, dtype=tf.float32)
-        self.ln3 = tf.keras.layers.LayerNormalization(epsilon=1e-5, dtype=tf.float32)
-    def call(self, x):
-        x = self.ln1(x)
-        attn = self.mha(x, x, x)
-        x = self.ln2(attn) + x
-        re = x
-        w = self.WB(x)
-        a, b = tf.split(w, 2, axis=-1)
-        g = tf.nn.silu(a) * b
-        o = self.W(g)
-        return self.ln3(o) + re
+class EncoderBlock(tf.keras.layers.Layer):
+    def __init__(self, embed_dim=EMBED_DIM, ff_dim=1152, seq_len=MAX_LEN):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.seq_len = seq_len
+
+        self.fc1 = layers.Dense(ff_dim)
+        self.fc2 = layers.Dense(embed_dim)
+        self.fc3 = layers.Dense(ff_dim//2)
+        self.fc4 = layers.Dense(embed_dim)
+
+        self.attn = layers.Dense(1)
+        self.token_mixer = layers.Dense(seq_len)
+        self.token_gate = layers.Dense(seq_len, activation='sigmoid')
+
+        self.ln = layers.LayerNormalization(epsilon=1e-5)
+        self.ln1 = layers.LayerNormalization(epsilon=1e-5)
+        self.ln2 = layers.LayerNormalization(epsilon=1e-5)
+        self.ln3 = layers.LayerNormalization(epsilon=1e-5)
+        self.ln4 = layers.LayerNormalization(epsilon=1e-5)
+
+    def call(self, x, mask):
+        mask = mask
+        # x: (B, L, D)
+        x_norm = self.ln(x)
+
+        h = self.fc1(x_norm)
+        g, v = tf.split(h, 2, axis=-1) 
+        h = tf.nn.silu(g) * v
+        h = self.fc2(h)  
+
+        h = x + self.ln1(h)
+
+        scores = self.attn(h)
+        scores = tf.where(tf.equal(mask[..., tf.newaxis], 0), -1e9, scores)
+        scores = tf.nn.softmax(scores, axis=1)
+        attn = h + self.ln2(h * scores)
+
+        v = tf.transpose(attn, [0, 2, 1])
+        v = self.token_mixer(v) * self.token_gate(v)
+        v = tf.transpose(v, [0, 2, 1])
+
+        x_norm = attn + self.ln3(v)
+        x = self.fc3(x_norm)
+        x = tf.nn.silu(x)
+        x = self.fc4(x)
+
+        return x_norm + self.ln4(x)
 
 class L2NormLayer(layers.Layer):
     def __init__(self, axis=1, epsilon=1e-10, **kwargs):
         super().__init__(**kwargs)
         self.axis = axis
         self.epsilon = epsilon
-
     def call(self, inputs):
         return tf.math.l2_normalize(inputs, axis=self.axis, epsilon=self.epsilon)
-
     def get_config(self):
         return {"axis": self.axis, "epsilon": self.epsilon, **super().get_config()}
 
 class SentenceEncoder(tf.keras.Model):
-    def __init__(self, vocab_size, embed_dim=384, latent_dim=384, max_len=128, pad_id=3):
+    def __init__(self, vocab_size, embed_dim=384, latent_dim=384, max_len=128, pad_id=pad_id):
         super().__init__()
         self.pad_id = pad_id
         self.embed = layers.Embedding(vocab_size, embed_dim)
         self.pos_embed = layers.Embedding(input_dim=max_len, output_dim=embed_dim)
-        self.blocks = [EncoderBlock(embed_dim=embed_dim, latent_dim=latent_dim) for _ in range(2)]
+        self.blocks = [EncoderBlock() for _ in range(2)]
         self.attn_pool = layers.Dense(1)
         self.ln_f = layers.LayerNormalization(epsilon=1e-5, dtype=tf.float32)
         self.latent = layers.Dense(latent_dim, activation=None)  # tanh 제거
@@ -98,10 +125,9 @@ class SentenceEncoder(tf.keras.Model):
         positions = tf.range(tf.shape(x)[1])[tf.newaxis, :]
         x_embed = self.embed(x) + self.pos_embed(positions)
         mask = tf.cast(tf.not_equal(x, self.pad_id), tf.float32)
-
         x = x_embed
         for block in self.blocks:
-            x = block(x)
+            x = block(x, mask)
         x = self.ln_f(x)
 
         scores = self.attn_pool(x)
@@ -110,7 +136,8 @@ class SentenceEncoder(tf.keras.Model):
         pooled = tf.reduce_sum(x * scores, axis=1)
 
         latent = self.latent(pooled)
-        return self.l2norm(latent)  # L2 정규��
+        return self.l2norm(latent)  # L2 정규화 후 반환
+
 # 3️⃣ 모델 로드
 # ===============================
 encoder = SentenceEncoder(vocab_size=vocab_size)