Update api.py

api.py

@@ -172,8 +172,10 @@ def is_greedy_response_acceptable(text):
   
     return True
 
-def generate_text_typical(model, prompt, max_len=100, max_gen=98,
-                          typical_p=0.5, min_len=20):
+def generate_text_flex(model, prompt, max_len=100, max_gen=98,
+                       repetition_penalty=1.2, temperature=1.0,
+                       top_k=50, top_p=0.85, typical_p=0.72,
+                       min_len=20):
     model_input = text_to_ids(f"<start> {prompt} <sep>")
     model_input = model_input[:max_len]
     generated = list(model_input)
@@ -185,28 +187,61 @@ def generate_text_typical(model, prompt, max_len=100, max_gen=98,
         logits = model(input_tensor, training=False)
         next_logits = logits[0, len(generated) - 1].numpy()
 
-        # 🔥 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)
-
-        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
+        # Repetition penalty 적용
+        for t in set(generated):
+            count = generated.count(t)
+            next_logits[t] /= (repetition_penalty ** count)
+
+        # Temperature scaling
+        next_logits = next_logits / temperature
+
+        # 확률 계산
+        probs = np.exp(next_logits - np.max(next_logits))
+        probs = probs / probs.sum()
+
+        # Top-K 필터링
+        if top_k is not None and top_k > 0:
+            indices_to_remove = probs < np.sort(probs)[-top_k]
+            probs[indices_to_remove] = 0
+            probs /= probs.sum()
+
+        # Top-P (Nucleus) 필터링
+        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)
+            cutoff_index = np.searchsorted(cumulative_probs, top_p, side='right')
+            keep_indices = sorted_indices[:cutoff_index + 1]
 
-        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)
+            filtered_probs = np.zeros_like(probs)
+            filtered_probs[keep_indices] = probs[keep_indices]
+            filtered_probs /= filtered_probs.sum()
+            probs = filtered_probs
 
+        # Typical-p 필터링
+        if typical_p is not None and 0 < typical_p < 1:
+            log_probs = -np.log(probs + 1e-10)
+            mean_info = np.mean(log_probs)
+            deviation = np.abs(log_probs - mean_info)
+            sorted_indices = np.argsort(deviation)
+
+            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
+
+            filtered_probs = np.zeros_like(probs)
+            filtered_probs[filtered_indices] = probs[filtered_indices]
+            filtered_probs /= filtered_probs.sum()
+            probs = filtered_probs
+
+        # 다음 토큰 샘플링
+        next_token = np.random.choice(len(probs), p=probs)
         generated.append(int(next_token))
+
         decoded = sp.decode(generated)
         for t in ["<start>", "<sep>", "<end>"]:
             decoded = decoded.replace(t, "")
@@ -343,9 +378,9 @@ def respond(input_text):
         return f"{summary}\n다른 궁금한 점 있으신가요?"
 
     # 일상 대화: 샘플링 + fallback
-    response = generate_text_typical(model, input_text)
+    response = generate_text_flex(model, input_text)
     if not is_valid_response(response) or mismatch_tone(input_text, response):
-        response = generate_text_typical(model, input_text)
+        response = generate_text_flex(model, input_text)
     return response
 
 @app.get("/generate", response_class=PlainTextResponse)