Update Mo.py

Mo.py

@@ -215,60 +215,51 @@ class ReLM(tf.keras.Model):
         logits = tf.matmul(x, embedding_matrix, transpose_b=True)
         return tf.cast(logits, tf.float32)
 
-loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction='none')
-
-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)
-    masked_loss = tf.reduce_sum(loss * mask) / tf.reduce_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_accuracy(y_true, y_pred):
+    y_true = tf.cast(y_true, tf.int32)
     mask = tf.cast(tf.not_equal(y_true, pad_id), tf.float32)
-    avg_loss = tf.reduce_sum(loss * mask) / tf.reduce_sum(mask)
-    return tf.exp(tf.minimum(avg_loss, 10.0))  # 수치 안정성 확보
-
-def create_lr_schedule(initial_lr=5e-5, decay_steps=10000, decay_rate=0.9):
-    return tf.keras.optimizers.schedules.ExponentialDecay(
-        initial_learning_rate=initial_lr,
-        decay_steps=decay_steps,
-        decay_rate=decay_rate,
-        staircase=False
-    )
+    pred_id = tf.argmax(y_pred, axis=-1, output_type=tf.int32)
+    acc = tf.cast(tf.equal(y_true, pred_id), tf.float32) * mask
+    return tf.reduce_sum(acc) / (tf.reduce_sum(mask) + 1e-8)
 
-# 모델 생성
-model = ReLM(
-    vocab_size=vocab_size,
-    max_seq_len=max_len,
-    d_model=700,
-    n_layers=16
-)
+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)
+    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)
 
-# 옵티마이저 설정
-optimizer = tf.keras.optimizers.Adam(
-    learning_rate=create_lr_schedule(),
-    beta_1=0.9,
-    beta_2=0.95,
-    epsilon=1e-8,
-    clipnorm=1.0
-)
 
-# 모델 컴파일
-model.compile(
-    optimizer=optimizer,
-    loss=masked_loss,
-    metrics=[
-        masked_perplexity
-    ]
-)
+# =======================
+# 모델 생성 & 컴파일
+# =======================
+with strategy.scope():
+    model = ReLM(vocab_size=vocab_size, max_seq_len=max_len, d_ff=768, n_layers=12)
+    dummy_input = tf.zeros((batch_size, max_len), dtype=tf.int32)
+    _ = model(dummy_input, training=False)
+    model.summary()
 
-# 더미 인풋으로 모델 초기화
-dummy_input = np.zeros((1, max_len), dtype=np.int32)
-model(dummy_input)
-model.summary()
+    optimizer = tf.keras.optimizers.Adam(1e-4, beta_1=0.9, beta_2=0.95, epsilon=1e-8, clipnorm=1.0)
+    model.compile(optimizer=optimizer, loss=smoothed_loss_keras, metrics=[masked_accuracy, masked_perplexity])
 
-history = model.fit(dataset, epochs=1, verbose=1)
+    # 학습
+    history = model.fit(dist_dataset, epochs=1, verbose=1)
 
 
 # 가중치 저장