Update AlphaS2S.py

AlphaS2S.py

@@ -298,64 +298,50 @@ with strategy.scope():
 chat_model.save_weights("chat_model.weights.h5")
 print("\n✅ 모델 가중치 저장 완료!")
 
-# =======================
-# 6) 추론 함수 (기존 코드 유지)
-# =======================
+def generate_text_topp(model, context, prompt, max_len=256, max_gen=100, p=0.9, temperature=0.8, min_len=20):
+    # Encoder input: ID 레벨로 특수 토큰 삽입
+    enc_ids = [context_s_id] + text_to_ids(context) + [context_e_id] + \
+              [user_s_id] + text_to_ids(prompt) + [user_e_id]
+    enc_ids = enc_ids[-max_len:]  # 길이 제한
+    enc_tensor = tf.convert_to_tensor([np.pad(enc_ids, (0, max_len - len(enc_ids)), constant_values=pad_id)], dtype=tf.int32)
+
+    # Decoder input: <sos>로 시작
+    generated = [start_id]
 
-def generate_text_topp(model, prompt, max_len=150, max_gen=100, p=0.9, temperature=0.8, min_len=20):
-    # 인코더 입력은 <start> Prompt <sep> 만 사용
-    model_input = text_to_ids(f"<start> {prompt} <sep>")
-    model_input = model_input[:max_len]
-    generated = list(model_input)  
     for step in range(max_gen):
-        current_len = len(generated)
-        # 현재까지 생성된 시퀀스를 입력으로 사용
-        if current_len > 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])
-        # 모델 추론 (enc_inputs, dec_inputs 모두 동일한 시퀀스를 사용)
-        dummy_input = {
-            "enc_inputs": input_tensor,
-            "dec_inputs": input_tensor
-        }
-        logits = model(dummy_input, training=False)
-        # 다음 토큰의 로짓은 시퀀스의 마지막 토큰 위치에서 가져옴 (0-based index: current_len - 1)
-        # 하지만 패딩 후 input_tensor의 실제 시퀀스 길이는 len(input_seq)
-        next_token_logits = logits[0, len(input_seq) - 1].numpy()
-        # 특수 토큰 생성 억제
-        next_token_logits[end_id] -= 5.0
+        dec_input = generated[-max_len:]  # max_len 유지
+        dec_tensor = tf.convert_to_tensor([np.pad(dec_input, (0, max_len - len(dec_input)), constant_values=pad_id)], dtype=tf.int32)
+
+        # 모델 추론
+        logits = model({"enc_inputs": enc_tensor, "dec_inputs": dec_tensor}, training=False)
+        # 마지막 토큰 위치 logits 사용
+        next_token_logits = logits[0, len(dec_input) - 1].numpy()
+        # 특수 토큰 억제
         next_token_logits[pad_id] -= 10.0
+        next_token_logits[context_s_id] -= 5.0
+        next_token_logits[context_e_id] -= 5.0
+        next_token_logits[user_s_id] -= 5.0
+        next_token_logits[user_e_id] -= 5.0
+
+        # Softmax + Top-p
         probs = tf.nn.softmax(next_token_logits / temperature).numpy()
         sorted_indices = np.argsort(probs)[::-1]
         sorted_probs = probs[sorted_indices]
-        # Top-p (Nucleus) Sampling
         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    
+            break
         generated.append(int(next_token_id))
-    # <start> 토큰 제거 및 <sep> 이전 부분 제거
-    try:
-        sep_index = generated.index(sep_id)
-        # <sep> 이후부터 <end> 이전까지의 응답만 반환
-        result_ids = generated[sep_index + 1:]
-        try:
-            end_index = result_ids.index(end_id)
-            result_ids = result_ids[:end_index]
-        except ValueError:
-            pass
-        return ids_to_text(result_ids)
-    except ValueError:
-        return ids_to_text(generated) # <sep>이 없으면 전체 반환
 
+    # <sos> 제거 후 텍스트로 변환
+    result_ids = generated[1:]  # 첫 토큰 <sos> 제거
+    return ids_to_text(result_ids)
+
+# 예시 사용
 print("\n\n===== 생성 결과 =====")  
-# 모델이 1 epoch만 학습되었으므로 의미 있는 결과가 아닐 수 있습니다.
-print(generate_text_topp(chat_model, "제가 이따가 버스를 타야 해서 준비 좀 해야겠어요. 재미있는 대화였습니다!", p=0.9))
+print(generate_text_topp(chat_model, "대화 시작", "안녕하세요! 어떻게 지내셨나요?", p=0.9))