Update api.py

api.py

@@ -138,9 +138,9 @@ _ = model(dummy_input)  # 모델이 빌드됨
 model.load_weights("InteractGPT.weights.h5")      
 print("모델 가중치 로드 완료!")      
 
-def generate_text_top_p(model, prompt, max_len=100, max_gen=98,
-                        temperature=1.0, min_len=20,
-                        repetition_penalty=1.1, top_p=0.9):
+def generate_text_top_kp(model, prompt, max_len=100, max_gen=98,
+                         temperature=1.0, min_len=20,
+                         repetition_penalty=1.1, top_k=40, top_p=0.9):
     model_input = text_to_ids(f"<start> {prompt} <sep>")
     model_input = model_input[:max_len]
     generated = list(model_input)
@@ -153,12 +153,12 @@ def generate_text_top_p(model, prompt, max_len=100, max_gen=98,
         logits = model(input_tensor, training=False)
         next_logits = logits[0, len(generated) - 1].numpy()
 
-        # 반복 억제 penalty
+        # 반복 억제
         for t in set(generated):
             count = generated.count(t)
             next_logits[t] /= (repetition_penalty ** count)
 
-        # 종료 조건 방지
+        # 조기 종료 방지
         if len(generated) < min_len:
             next_logits[end_id] -= 5.0
         next_logits[pad_id] -= 10.0
@@ -168,17 +168,23 @@ def generate_text_top_p(model, prompt, max_len=100, max_gen=98,
         probs = np.exp(next_logits - np.max(next_logits))
         probs /= probs.sum()
 
-        # Top-p 필터링
-        sorted_idx = np.argsort(-probs)
-        sorted_probs = probs[sorted_idx]
+        # Top-K 적용
+        top_k = min(top_k, len(probs))
+        top_k_idx = np.argsort(-probs)[:top_k]
+        top_k_probs = probs[top_k_idx]
+        top_k_probs /= top_k_probs.sum()
+
+        # Top-P 필터링
+        sorted_idx = np.argsort(-top_k_probs)
+        sorted_probs = top_k_probs[sorted_idx]
         cum_probs = np.cumsum(sorted_probs)
         cutoff = np.searchsorted(cum_probs, top_p) + 1
 
-        filtered_idx = sorted_idx[:cutoff]
-        filtered_probs = sorted_probs[:cutoff]
-        filtered_probs /= filtered_probs.sum()
+        final_idx = top_k_idx[sorted_idx[:cutoff]]
+        final_probs = sorted_probs[:cutoff]
+        final_probs /= final_probs.sum()
 
-        sampled = np.random.choice(filtered_idx, p=filtered_probs)
+        sampled = np.random.choice(final_idx, p=final_probs)
         generated.append(int(sampled))
 
         decoded = sp.decode(generated)
@@ -189,9 +195,9 @@ def generate_text_top_p(model, prompt, max_len=100, max_gen=98,
         if len(generated) >= min_len and (sampled == end_id or decoded.endswith(('.', '!', '?'))):
             yield decoded
             break
-              
+
 async def async_generator_wrapper(prompt: str):
-    gen = generate_text_top_p(model, prompt)
+    gen = generate_text_top_kp(model, prompt)
     for text_piece in gen:
         yield text_piece
         await asyncio.sleep(0.1)