Update app.py

app.py

@@ -32,9 +32,9 @@ if not os.path.exists(TOKENIZER_PATH):
         TOKENIZER_PATH
     )
 
-MAX_LEN = 128
-EMBED_DIM = 384
-LATENT_DIM = 384
+MAX_LEN = 384
+EMBED_DIM = 512
+LATENT_DIM = 512
 DROP_RATE = 0.1 
 
 # ===============================
@@ -50,51 +50,45 @@ def encode_sentence(sentence, max_len=MAX_LEN):
 def pad_sentence(tokens):
     return tokens + [pad_id]*(MAX_LEN - len(tokens))
 
+
 class DynamicConv(layers.Layer):
-    def __init__(self, k=7):
+    def __init__(self, d_model, k=7):
         super().__init__()
-        assert k % 2 == 1, "kernel size should be odd for symmetric padding"
+        assert k % 2 == 1
         self.k = k
-        # generator는 각 토큰에 대해 k개의 로짓을 뱉음 -> softmax로 가중치화
-        self.generator = layers.Dense(k)
-
+        self.dense = layers.Dense(d_model, activation='silu')
+        self.proj = layers.Dense(d_model)
+        self.generator = layers.Dense(k, dtype='float32')
     def call(self, x):
-        # x: (B, L, D)
+        x_in = x
+        x = tf.cast(x, tf.float32)
+
         B = tf.shape(x)[0]
         L = tf.shape(x)[1]
         D = tf.shape(x)[2]
 
-        # (B, L, k) logits -> softmax -> (B, L, k)
-        kernels = self.generator(x)
+        kernels = self.generator(self.dense(x))
         kernels = tf.nn.softmax(kernels, axis=-1)
 
-        # padding (same)
         pad = (self.k - 1) // 2
-        x_pad = tf.pad(x, [[0, 0], [pad, pad], [0, 0]])  # (B, L+2pad, D)
+        x_pad = tf.pad(x, [[0,0],[pad,pad],[0,0]])
 
-        # extract patches using tf.image.extract_patches:
-        # make 4D: (B, H=1, W=L+2pad, C=D)
         x_pad_4d = tf.expand_dims(x_pad, axis=1)
         patches = tf.image.extract_patches(
             images=x_pad_4d,
-            sizes=[1, 1, self.k, 1],
-            strides=[1, 1, 1, 1],
-            rates=[1, 1, 1, 1],
+            sizes=[1,1,self.k,1],
+            strides=[1,1,1,1],
+            rates=[1,1,1,1],
             padding='VALID'
-        )  # (B, 1, L, k*D)
-
-        # reshape -> (B, L, k, D)
-        patches = tf.reshape(patches, [B, 1, L, self.k * D])
-        patches = tf.squeeze(patches, axis=1)
+        )
         patches = tf.reshape(patches, [B, L, self.k, D])
 
-        # kernels: (B, L, k) -> (B, L, k, 1)
         kernels_exp = tf.expand_dims(kernels, axis=-1)
-
-        # weighted sum over kernel dim -> (B, L, D)
         out = tf.reduce_sum(patches * kernels_exp, axis=2)
+        out = self.proj(out)
 
-        return out
+        # 🔥 원래 dtype으로 돌려줌
+        return tf.cast(out, x_in.dtype)
 
 class EncoderBlock(tf.keras.layers.Layer):
     def __init__(self, embed_dim=EMBED_DIM, ff_dim=1152, seq_len=MAX_LEN, num_conv_layers=2):
@@ -105,10 +99,7 @@ class EncoderBlock(tf.keras.layers.Layer):
         # MLP / FFN
         self.fc1 = layers.Dense(ff_dim)
         self.fc2 = layers.Dense(embed_dim)
-
-        # DynamicConv 블록 여러 개 쌓기
-        self.blocks = [DynamicConv(k=7) for _ in range(num_conv_layers)]
-
+        self.blocks = [DynamicConv(d_model=embed_dim, k=7) for _ in range(num_conv_layers)]
         # LayerNorm
         self.ln = layers.LayerNormalization(epsilon=1e-5)   # 입력 정규화
         self.ln1 = layers.LayerNormalization(epsilon=1e-5)  # Conv residual
@@ -120,9 +111,7 @@ class EncoderBlock(tf.keras.layers.Layer):
 
         # DynamicConv 여러 층 통과
         out = x_norm
-        for block in self.blocks:
-            out = block(out)
-
+        for block in self.blocks: out = block(out)
         # Conv residual 연결
         x = x_norm + self.ln1(out)
 
@@ -138,6 +127,7 @@ class EncoderBlock(tf.keras.layers.Layer):
 
         return x
 
+
 class L2NormLayer(layers.Layer):
     def __init__(self, axis=1, epsilon=1e-10, **kwargs):
         super().__init__(**kwargs)
@@ -145,37 +135,46 @@ class L2NormLayer(layers.Layer):
         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=pad_id):
+class SentenceEncoder(Model):
+    def __init__(self, vocab_size, embed_dim=EMBED_DIM, latent_dim=LATENT_DIM, max_len=MAX_LEN, pad_id=pad_id, dropout_rate=EMBED_DROPOUT):
         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.dropout = layers.Dropout(dropout_rate)
         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 제거
-        self.l2norm = L2NormLayer()  # 추가
+        self.latent = layers.Dense(latent_dim, activation=None)
+        self.l2norm = L2NormLayer(axis=1)
 
-    def call(self, x):
+    def call(self, x, training=None):
         positions = tf.range(tf.shape(x)[1])[tf.newaxis, :]
         x_embed = self.embed(x) + self.pos_embed(positions)
+        x_embed = self.dropout(x_embed, training=training)
+
         mask = tf.cast(tf.not_equal(x, self.pad_id), tf.float32)
-        x = x_embed
+
+        h = x_embed
         for block in self.blocks:
-            x = block(x, mask)
-        x = self.ln_f(x)
+            h = block(h, training=training)
+
+        h = self.ln_f(h)
 
-        scores = self.attn_pool(x)
-        scores = tf.where(tf.equal(mask[..., tf.newaxis], 0), -1e9, scores)
+        # 🔥 scores를 float32 강제
+        scores = self.attn_pool(h)
+        scores = tf.cast(scores, tf.float32)
+
+        scores = tf.where(mask[..., tf.newaxis] == 0, tf.constant(-1e9, tf.float32), scores)
         scores = tf.nn.softmax(scores, axis=1)
-        pooled = tf.reduce_sum(x * scores, axis=1)
 
+        pooled = tf.reduce_sum(h * scores, axis=1)
         latent = self.latent(pooled)
-        return self.l2norm(latent)  # L2 정규화 후 반환
+        latent = self.l2norm(latent)
+
+        # 🔥 출력만 float32
+        return tf.cast(latent, tf.float32)
 
 # 3️⃣ 모델 로드
 # ===============================