Update api.py

api.py

@@ -140,14 +140,19 @@ _ = model(dummy_input)  # 모델이 빌드됨
 model.load_weights("InteractGPT.weights.h5")      
 print("모델 가중치 로드 완료!")      
 
-def generate_text_topp(model, prompt, max_len=100, max_gen=98,  
-                       temperature=0.50, min_len=20,  
-                       repetition_penalty=1.2, top_p=0.90):    
-    def top_p_filtering(logits, top_p):  
+def generate_text_topkp(model, prompt, max_len=100, max_gen=98,  
+                        temperature=0.50, min_len=20,  
+                        repetition_penalty=1.2, top_p=0.90, top_k=50):    
+    def top_kp_filtering(logits, top_k, top_p):  
         probs = np.exp(logits - np.max(logits))  
         probs /= probs.sum()  
         sorted_idx = np.argsort(-probs)  
         sorted_probs = probs[sorted_idx]  
+        # Top-K 필터링  
+        if top_k > 0:  
+            sorted_idx = sorted_idx[:top_k]  
+            sorted_probs = sorted_probs[:top_k]  
+        # Top-P 필터링  
         cum_probs = np.cumsum(sorted_probs)  
         cutoff = np.searchsorted(cum_probs, top_p) + 1  
         final_idx = sorted_idx[:cutoff]  
@@ -174,8 +179,8 @@ def generate_text_topp(model, prompt, max_len=100, max_gen=98,
         next_logits[pad_id] -= 10.0  
         # 온도 적용  
         next_logits = next_logits / temperature  
-        # Top-P Sampling 적용  
-        final_idx, final_probs = top_p_filtering(next_logits, top_p=top_p)  
+        # Top-KP Sampling 적용  
+        final_idx, final_probs = top_kp_filtering(next_logits, top_k=top_k, top_p=top_p)  
         sampled = np.random.choice(final_idx, p=final_probs)  
         generated.append(int(sampled))  
         decoded = sp.decode(generated)  
@@ -266,10 +271,10 @@ def respond(input_text):
         summary = get_wikipedia_summary(keyword)
         return f"{summary}\n다른 궁금한 점 있으신가요?"
 
-    return generate_text_topp(model, input_text)
+    return generate_text_topkp(model, input_text)
 
 async def async_generator_wrapper(prompt: str):
-    gen = generate_text_topp(model, prompt)
+    gen = generate_text_topkp(model, prompt)
     for text_piece in gen:
         yield text_piece
         await asyncio.sleep(0.1)