Update api.py

api.py

@@ -172,13 +172,11 @@ def is_greedy_response_acceptable(text):
   
     return True
 
-def generate_text_with_mirostat(model, prompt, max_len=100, max_gen=98,
-                                repetition_penalty=1.2, tau_init=5.0, mu=5.0,
-                                eta=0.1, min_len=20):
+def generate_text_typical(model, prompt, max_len=100, max_gen=98,
+                          typical_p=0.9, min_len=20):
     model_input = text_to_ids(f"<start> {prompt} <sep>")
     model_input = model_input[:max_len]
     generated = list(model_input)
-    tau = tau_init
 
     for _ in range(max_gen):
         pad_len = max(0, max_len - len(generated))
@@ -187,29 +185,33 @@ def generate_text_with_mirostat(model, prompt, max_len=100, max_gen=98,
         logits = model(input_tensor, training=False)
         next_logits = logits[0, len(generated) - 1].numpy()
 
-        # Repetition penalty
-        for t in set(generated):
-            count = generated.count(t)
-            next_logits[t] /= (repetition_penalty ** count)
+        # 🔥 Typical Sampling
+        probs = tf.nn.softmax(next_logits).numpy()
+        log_probs = -np.log(probs + 1e-10)
+        info_content = log_probs
+        mean_info = np.mean(info_content)
+        deviation = np.abs(info_content - mean_info)
+        sorted_indices = np.argsort(deviation)
 
-        # Mirostat: softmax with current τ
-        scaled_logits = next_logits / tau
-        exp_logits = np.exp(scaled_logits - np.max(scaled_logits))
-        probs = exp_logits / exp_logits.sum()
-
-        # 샘플링
-        next_token = np.random.choice(len(probs), p=probs)
-        generated.append(int(next_token))
+        filtered_indices = []
+        cumulative_prob = 0.0
+        for idx in sorted_indices:
+            cumulative_prob += probs[idx]
+            filtered_indices.append(idx)
+            if cumulative_prob >= typical_p:
+                break
 
-        # 정보량 계산 (r = -log(p(x)))
-        prob = probs[next_token]
-        r = -np.log(prob + 1e-10)
+        filtered_probs = np.zeros_like(probs)
+        filtered_probs[filtered_indices] = probs[filtered_indices]
+        filtered_probs /= filtered_probs.sum()
+        next_token = np.random.choice(len(filtered_probs), p=filtered_probs)
 
-        # τ 업데이트
-        tau = tau - eta * (r - mu)
+        generated.append(int(next_token))
+        decoded = sp.decode(generated)
+        for t in ["<start>", "<sep>", "<end>"]:
+            decoded = decoded.replace(t, "")
+        decoded = decoded.strip()
 
-        # 디코딩 및 조건 검증
-        decoded = sp.decode(generated).replace("<start>", "").replace("<sep>", "").replace("<end>", "").strip()
         if len(generated) >= min_len and (next_token == end_id or decoded.endswith(('요', '다', '.', '!', '?'))):
             if is_greedy_response_acceptable(decoded):
                 return decoded
@@ -341,9 +343,9 @@ def respond(input_text):
         return f"{summary}\n다른 궁금한 점 있으신가요?"
 
     # 일상 대화: 샘플링 + fallback
-    response = generate_text_with_mirostat(model, input_text)
+    response = generate_text_typical(model, input_text)
     if not is_valid_response(response) or mismatch_tone(input_text, response):
-        response = generate_text_with_mirostat(model, input_text)
+        response = generate_text_typical(model, input_text)
     return response
 
 @app.get("/generate", response_class=PlainTextResponse)