Update Inference.py

Inference.py

@@ -13,7 +13,7 @@ tf.get_logger().setLevel("ERROR")
 SEED = 42
 tf.random.set_seed(SEED)
 np.random.seed(SEED)
-max_len = 150 # 기존 코드에서 200으로 설정됨
+max_len = 512 # 기존 코드에서 200으로 설정됨
 batch_size = 128
 
 # TPU 초기화 (기존 코드와 동일)
@@ -186,83 +186,44 @@ model = ReLM(
     d_model=256,
     n_layers=1
 )
-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.load_weights('/kaggle/working/chat_model.weights.h5')
+chat_model.load_weights('/kaggle/working/model.weights.h5')
 print("모델 가중치 로드 완료!")
 # =======================
 # 6) 추론 함수 (기존 코드 유지)
-# =======================
+# 더미 인풋으로 모델 초기화
+dummy_input = np.zeros((1, max_len), dtype=np.int32)
+_ = model(dummy_input)
+model.summary()
 
-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>")
+def generate_text_topp(model, prompt, max_len=512, max_gen=512, 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):
-        current_len = len(generated)
-        
-        # 현재까지 생성된 시퀀스를 입력으로 사용
-        if current_len > max_len:
+        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])
-        
-        # 모델 추론 (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)
+        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]
-        
-        # 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
-            
         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>이 없으면 전체 반환
+    return ids_to_text(generated)
 
 print("\n\n===== 생성 결과 =====")  
-# 모델이 1 epoch만 학습되었으므로 의미 있는 결과가 아닐 수 있습니다.
-print(generate_text_topp(chat_model, "제가 이따가 버스를 타야 해서 준비 좀 해야겠어요. 재미있는 대화였습니다!", p=0.9))
+print(generate_text_topp(model, "지난 2년 동안", p=0.9))