Update api.py

api.py

@@ -167,7 +167,7 @@ def is_greedy_response_acceptable(text):
     return True
 
 def generate_text_sample(model, prompt, max_len=100, max_gen=98,
-                         temperature=0.7, top_k=40, min_len=12):
+                         temperature=0.7, top_k=40, top_p=0.9, min_len=12):
     model_input = text_to_ids(f"<start> {prompt} <sep>")
     model_input = model_input[:max_len]
     generated = list(model_input)
@@ -190,6 +190,20 @@ def generate_text_sample(model, prompt, max_len=100, max_gen=98,
             probs[indices_to_remove] = 0
             probs /= probs.sum()
 
+        # Top-P (누적 확률) 필터링
+        if top_p is not None and 0 < top_p < 1:
+            sorted_indices = np.argsort(probs)[::-1]
+            sorted_probs = probs[sorted_indices]
+            cumulative_probs = np.cumsum(sorted_probs)
+            # 누적 확률이 top_p 초과하는 토큰들은 제거
+            cutoff_index = np.searchsorted(cumulative_probs, top_p, side='right')
+            probs_to_keep = sorted_indices[:cutoff_index+1]
+            
+            mask = np.ones_like(probs, dtype=bool)
+            mask[probs_to_keep] = False
+            probs[mask] = 0
+            probs /= probs.sum()
+
         # 샘플링
         next_token = np.random.choice(len(probs), p=probs)
         generated.append(int(next_token))
@@ -210,7 +224,6 @@ def generate_text_sample(model, prompt, max_len=100, max_gen=98,
     for t in ["<start>", "<sep>", "<end>"]:
         decoded = decoded.replace(t, "")
     return decoded.strip()
-
     
 def mismatch_tone(input_text, output_text):  
     if "ㅋㅋ" in input_text and not re.search(r'ㅋㅋ|ㅎ|재밌|놀|만나|맛집|여행', output_text):