Update AlphaS2S.py

AlphaS2S.py

@@ -233,23 +233,28 @@ class Transformer(tf.keras.Model):
 
 # 5) 학습 설정 및 실행
 # =======================
-
-def masked_loss(y_true, y_pred):
-    loss = loss_fn(y_true, y_pred)
+def smoothed_loss_keras(y_true, y_pred, eps=0.1):
+    y_true = tf.cast(y_true, tf.int32)
     mask = tf.cast(tf.not_equal(y_true, pad_id), tf.float32)
-    # mixed_bfloat16 사용 시 나눗셈 시 NaN 방지
-    sum_mask = tf.reduce_sum(mask)
-    safe_sum_mask = tf.where(sum_mask == 0.0, 1.0, sum_mask)
-    masked_loss = tf.reduce_sum(loss * mask) / safe_sum_mask
-    return masked_loss
-
-def masked_perplexity(y_true, y_pred):
-    loss = loss_fn(y_true, y_pred)
+    vocab = tf.shape(y_pred)[-1]
+    y_true_oh = tf.one_hot(y_true, depth=vocab, dtype=tf.float32)
+    y_true_ls = (1.0 - eps) * y_true_oh + eps / tf.cast(vocab, tf.float32)
+    log_probs = tf.nn.log_softmax(y_pred, axis=-1)
+    per_tok = -tf.reduce_sum(y_true_ls * log_probs, axis=-1)
+    per_tok = per_tok * mask
+    return tf.reduce_sum(per_tok) / (tf.reduce_sum(mask) + 1e-8)
+
+def masked_perplexity(y_true, y_pred, eps=0.1):
+    y_true = tf.cast(y_true, tf.int32)
     mask = tf.cast(tf.not_equal(y_true, pad_id), tf.float32)
-    sum_mask = tf.reduce_sum(mask)
-    safe_sum_mask = tf.where(sum_mask == 0.0, 1.0, sum_mask)
-    avg_loss = tf.reduce_sum(loss * mask) / safe_sum_mask
-    return tf.exp(tf.minimum(avg_loss, 10.0))  
+    vocab = tf.shape(y_pred)[-1]
+    y_true_oh = tf.one_hot(y_true, depth=vocab, dtype=tf.float32)
+    y_true_ls = (1.0 - eps) * y_true_oh + eps / tf.cast(vocab, tf.float32)
+    log_probs = tf.nn.log_softmax(y_pred, axis=-1)
+    per_tok = -tf.reduce_sum(y_true_ls * log_probs, axis=-1)
+    per_tok = per_tok * mask
+    mean_loss = tf.reduce_sum(per_tok) / (tf.reduce_sum(mask) + 1e-8)
+    return tf.exp(mean_loss) 
 
 def create_lr_schedule(initial_lr=5e-5, decay_steps=10000, decay_rate=0.9):
     return tf.keras.optimizers.schedules.ExponentialDecay(
@@ -271,7 +276,6 @@ with strategy.scope():
     
     loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction='none')
 
-
     # 옵티마이저 설정
     optimizer = tf.keras.optimizers.Adam(
         learning_rate=create_lr_schedule(),
@@ -280,15 +284,7 @@ with strategy.scope():
         epsilon=1e-8,
         clipnorm=1.0
     )
-
-    # 모델 컴파일
-    chat_model.compile(
-        optimizer=optimizer,
-        loss=masked_loss,
-        metrics=[
-            masked_perplexity
-        ]
-    )
+    chat_model.compile(optimizer=optimizer, loss=smoothed_loss_keras, metrics=[masked_perplexity])
     chat_model.summary()
     print("✅ 모델 컴파일 완료, 학습 시작...")
     # ⚠️ 학습 실행