Update api.py

api.py

@@ -138,9 +138,30 @@ _ = model(dummy_input)  # 모델이 빌드됨
 model.load_weights("InteractGPT.weights.h5")      
 print("모델 가중치 로드 완료!")      
 
-def generate_text_top_kp(model, prompt, max_len=100, max_gen=98,
-                         temperature=0.7, min_len=20,
-                         repetition_penalty=1.1, top_k=50, top_p=0.9):
+def generate_text_typical(model, prompt, max_len=100, max_gen=98,
+                          temperature=0.7, min_len=20,
+                          repetition_penalty=1.1, typical_p=0.95):
+    
+    def typical_filtering(logits, typical_p):
+        probs = np.exp(logits - np.max(logits))
+        probs /= probs.sum()
+
+        log_probs = np.log(probs + 1e-9)
+        entropy = -np.sum(probs * log_probs)
+
+        shifted = np.abs(-log_probs - entropy)
+        sorted_idx = np.argsort(shifted)
+        sorted_probs = probs[sorted_idx]
+
+        cum_probs = np.cumsum(sorted_probs)
+        cutoff = np.searchsorted(cum_probs, typical_p) + 1
+
+        final_idx = sorted_idx[:cutoff]
+        final_probs = probs[final_idx]
+        final_probs /= final_probs.sum()
+
+        return final_idx, final_probs
+
     model_input = text_to_ids(f"<start> {prompt} <sep>")
     model_input = model_input[:max_len]
     generated = list(model_input)
@@ -165,26 +186,11 @@ def generate_text_top_kp(model, prompt, max_len=100, max_gen=98,
 
         # 온도 적용
         next_logits = next_logits / temperature
-        probs = np.exp(next_logits - np.max(next_logits))
-        probs /= probs.sum()
-
-        # 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
-
-        final_idx = top_k_idx[sorted_idx[:cutoff]]
-        final_probs = sorted_probs[:cutoff]
-        final_probs /= final_probs.sum()
 
+        # Typical Sampling 적용
+        final_idx, final_probs = typical_filtering(next_logits, typical_p=typical_p)
         sampled = np.random.choice(final_idx, p=final_probs)
+
         generated.append(int(sampled))
 
         decoded = sp.decode(generated)
@@ -197,7 +203,7 @@ def generate_text_top_kp(model, prompt, max_len=100, max_gen=98,
             break
 
 async def async_generator_wrapper(prompt: str):
-    gen = generate_text_top_kp(model, prompt)
+    gen = generate_text_typical(model, prompt)
     for text_piece in gen:
         yield text_piece
         await asyncio.sleep(0.1)