Update api.py

api.py

@@ -166,23 +166,35 @@ def is_greedy_response_acceptable(text):
   
     return True
 
-def generate_text_greedy(model, prompt, max_len=100, min_len=12):
+def generate_text_sample(model, prompt, max_len=100, max_gen=98,
+                         temperature=0.7, top_k=40, min_len=12):
     model_input = text_to_ids(f"<start> {prompt} <sep>")
     model_input = model_input[:max_len]
     generated = list(model_input)
 
-    for _ in range(max_len):
+    for _ in range(max_gen):
         pad_len = max(0, max_len - len(generated))
         input_padded = np.pad(generated, (0, pad_len), constant_values=pad_id)
         input_tensor = tf.convert_to_tensor([input_padded])
         logits = model(input_tensor, training=False)
         next_logits = logits[0, len(generated) - 1].numpy()
 
-        # Greedy: 확률 가장 높은 토큰만 선택
-        next_token = int(np.argmax(next_logits))
-        generated.append(next_token)
+        # Temperature 적용
+        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()
+
+        # 샘플링
+        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, "")
@@ -194,7 +206,6 @@ def generate_text_greedy(model, prompt, max_len=100, min_len=12):
             else:
                 continue
 
-    # 끝까지 만족하는 문장 없으면 그냥 최종 출력
     decoded = sp.decode(generated)
     for t in ["<start>", "<sep>", "<end>"]:
         decoded = decoded.replace(t, "")
@@ -324,9 +335,9 @@ def respond(input_text):
         return f"{summary}\n다른 궁금한 점 있으신가요?"
 
     # 일상 대화: 샘플링 + fallback
-    response = generate_text_greedy(model, input_text)
+    response = generate_text_sample(model, input_text)
     if not is_valid_response(response) or mismatch_tone(input_text, response):
-        response = generate_text_greedy(model, input_text)
+        response = generate_text_sample(model, input_text)
     return response
 
 @app.get("/generate", response_class=PlainTextResponse)