Update api.py

api.py

@@ -166,54 +166,39 @@ def is_greedy_response_acceptable(text):
   
     return True
 
-def generate_text_beam(model, prompt, max_len=100, beam_width=4, length_penalty=0.7):
+def generate_text_greedy(model, prompt, max_len=100, min_len=12):
     model_input = text_to_ids(f"<start> {prompt} <sep>")
     model_input = model_input[:max_len]
-
-    beams = [{
-        "sequence": list(model_input),
-        "score": 0.0
-    }]
+    generated = list(model_input)
 
     for _ in range(max_len):
-        all_candidates = []
-
-        for beam in beams:
-            seq = beam["sequence"]
-            pad_len = max(0, max_len - len(seq))
-            input_padded = np.pad(seq, (0, pad_len), constant_values=pad_id)
-            input_tensor = tf.convert_to_tensor([input_padded])
-            logits = model(input_tensor, training=False)[0, len(seq) - 1].numpy()
-
-            probs = np.exp(logits - np.max(logits))
-            probs = probs / probs.sum()
-
-            top_indices = probs.argsort()[-beam_width:][::-1]
-
-            for idx in top_indices:
-                new_seq = seq + [int(idx)]
-                new_score = beam["score"] + np.log(probs[idx])
-                all_candidates.append({
-                    "sequence": new_seq,
-                    "score": new_score
-                })
-
-        # 길이 보정
-        for cand in all_candidates:
-            cand["score"] /= (len(cand["sequence"]) ** length_penalty)
-
-        # 상위 beam_width개만 유지
-        beams = sorted(all_candidates, key=lambda x: x["score"], reverse=True)[:beam_width]
-
-        # 조기 종료 (EOS 토큰 또는 끝나는 말투)
-        for b in beams:
-            decoded = sp.decode(b["sequence"]).strip()
-            if end_id in b["sequence"] and is_greedy_response_acceptable(decoded):
+        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)
+
+        # 디코딩 및 필터링
+        decoded = sp.decode(generated)
+        for t in ["<start>", "<sep>", "<end>"]:
+            decoded = decoded.replace(t, "")
+        decoded = decoded.strip()
+
+        if len(generated) >= min_len and (next_token == end_id or decoded.endswith(('요', '다', '.', '!', '?'))):
+            if is_greedy_response_acceptable(decoded):
                 return decoded
+            else:
+                continue
 
-    # 최종 후보 중 가장 점수 높은 거 반환
-    final = beams[0]["sequence"]
-    return sp.decode(final)
+    # 끝까지 만족하는 문장 없으면 그냥 최종 출력
+    decoded = sp.decode(generated)
+    for t in ["<start>", "<sep>", "<end>"]:
+        decoded = decoded.replace(t, "")
+    return decoded.strip()
 
     
 def mismatch_tone(input_text, output_text):  
@@ -339,9 +324,9 @@ def respond(input_text):
         return f"{summary}\n다른 궁금한 점 있으신가요?"
 
     # 일상 대화: 샘플링 + fallback
-    response = generate_text_beam(model, input_text)
+    response = generate_text_greedy(model, input_text)
     if not is_valid_response(response) or mismatch_tone(input_text, response):
-        response = generate_text_beam(model, input_text)
+        response = generate_text_greedy(model, input_text)
     return response
 
 @app.get("/generate", response_class=PlainTextResponse)