Update AlphaS2S.py

AlphaS2S.py

@@ -267,6 +267,26 @@ class AlphaS2S(tf.keras.Model):
         for layer in self.dec_layers: y = layer(y, enc_out, training=training)
         return self.final_layer(y)
 
+def masked_loss(y_true, y_pred):
+    loss = loss_fn(y_true, y_pred)
+    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)
+    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
+    )
 
 chat_model = AlphaS2S(num_layers=4, d_model=160, num_heads=8,
                          input_vocab_size=chat_vocab_size, target_vocab_size=chat_vocab_size)
@@ -274,4 +294,61 @@ dummy_input = {
     "enc_inputs": tf.zeros((1, max_len), dtype=tf.int32),
     "dec_inputs": tf.zeros((1, max_len), dtype=tf.int32)
 }
-_ = chat_model(dummy_input)
+_ = chat_model(dummy_input)
+loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction='none')
+
+
+# 옵티마이저 설정
+optimizer = tf.keras.optimizers.Adam(
+    learning_rate=create_lr_schedule(),
+    beta_1=0.9,
+    beta_2=0.95,
+    epsilon=1e-8,
+    clipnorm=1.0
+)
+
+# 모델 컴파일
+chat_model.compile(
+    optimizer=optimizer,
+    loss=masked_loss,
+    metrics=[
+        masked_perplexity
+    ]
+)
+
+history = chat_model.fit(dataset, epochs=1, verbose=1)
+# 가중치 저장
+chat_model.save_weights("chat_model.weights.h5")
+print("모델 가중치 저장 완료!")
+
+def generate_text_topp(model, prompt, max_len=150, max_gen=150, p=0.9, temperature=0.8, min_len=20):
+    model_input = text_to_ids(f"<start> {prompt}")
+    model_input = model_input[:max_len]
+    generated = list(model_input)
+    for step in range(max_gen):
+        if len(generated) > max_len:
+            input_seq = generated[-max_len:]
+        else:
+            input_seq = generated
+        input_padded = np.pad(input_seq, (0, max_len - len(input_seq)), constant_values=pad_id)
+        input_tensor = tf.convert_to_tensor([input_padded])
+        logits = model(input_tensor, training=False)
+        next_token_logits = logits[0, len(input_seq) - 1].numpy()
+        next_token_logits[end_id] -= 5.0
+        next_token_logits[pad_id] -= 10.0
+        probs = tf.nn.softmax(next_token_logits / temperature).numpy()
+        sorted_indices = np.argsort(probs)[::-1]
+        sorted_probs = probs[sorted_indices]
+        cumulative_probs = np.cumsum(sorted_probs)
+        cutoff = np.searchsorted(cumulative_probs, p)
+        top_indices = sorted_indices[:cutoff + 1]
+        top_probs = sorted_probs[:cutoff + 1]
+        top_probs /= np.sum(top_probs)
+        next_token_id = np.random.choice(top_indices, p=top_probs)
+        if next_token_id == end_id and len(generated) >= min_len:
+            break
+        generated.append(int(next_token_id))
+    return ids_to_text(generated)
+
+print("\n\n===== 생성 결과 =====")  
+print(generate_text_topp(chat_model, "지난 2년 동안 출연연이 국가가 필요한 연구를", p=0.9))