Update api.py

api.py

@@ -150,75 +150,47 @@ from sklearn.metrics.pairwise import cosine_similarity
 from fastapi import Request
 from fastapi.responses import PlainTextResponse
 
-# 1. Top-KP 기반 생성기
-def generate_text_topkp(model, prompt, max_len=100, max_gen=98,  
-                        temperature=0.90, min_len=20,  
-                        repetition_penalty=1.2, top_p=0.90, top_k=50):    
-    def top_kp_filtering(logits, top_k, top_p):  
-        probs = np.exp(logits - np.max(logits))  
-        probs /= probs.sum()  
-        sorted_idx = np.argsort(-probs)  
-        sorted_probs = probs[sorted_idx]  
-        if top_k > 0:  
-            sorted_idx = sorted_idx[:top_k]  
-            sorted_probs = sorted_probs[:top_k]  
-        cum_probs = np.cumsum(sorted_probs)  
-        cutoff = np.searchsorted(cum_probs, top_p) + 1  
-        final_idx = sorted_idx[:cutoff]  
-        final_probs = probs[final_idx]  
-        final_probs /= final_probs.sum()  
-        return final_idx, final_probs  
-
+def generate_text_greedy_strong(model, prompt, max_len=100, max_gen=98,  
+                                repetition_penalty=1.2, min_len=20):  
     model_input = text_to_ids(f"<start> {prompt} <sep>")  
     model_input = model_input[:max_len]  
     generated = list(model_input)  
-    for step in range(max_gen):  
+
+    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()  
+
+        # Repetition Penalty
         for t in set(generated):  
             count = generated.count(t)  
-            next_logits[t] /= (repetition_penalty ** count)  
-        if len(generated) < min_len:  
-            next_logits[end_id] -= 5.0  
-        next_logits[pad_id] -= 10.0  
-        next_logits = next_logits / temperature  
-        final_idx, final_probs = top_kp_filtering(next_logits, top_k=top_k, top_p=top_p)  
-        sampled = np.random.choice(final_idx, p=final_probs)  
-        generated.append(int(sampled))  
-        decoded = sp.decode(generated)  
-        for t in ["<start>", "<sep>", "<end>"]:  
-            decoded = decoded.replace(t, "")  
-        decoded = decoded.strip()  
-        if len(generated) >= min_len and (sampled == end_id or decoded.endswith(('.', '!', '?'))):  
-            return decoded
-    return sp.decode(generated)
+            next_logits[t] /= (repetition_penalty ** count)
 
-# Greedy 버전 생성기
-def generate_text_greedy(model, prompt, max_len=100, max_gen=98):  
-    model_input = text_to_ids(f"<start> {prompt} <sep>")  
-    model_input = model_input[:max_len]  
-    generated = list(model_input)  
-    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()  
-        next_logits[pad_id] -= 10.0  
-        next_token = np.argmax(next_logits)  
+        # Stop token filtering  
+        stop_tokens = ["음", "어", "그", "뭐지", "..."]
+        for tok in stop_tokens:
+            tok_id = sp.piece_to_id(tok)
+            next_logits[tok_id] -= 5.0
+
+        next_logits[pad_id] -= 10.0
+        next_token = np.argmax(next_logits)
         generated.append(int(next_token))  
+
         decoded = sp.decode(generated)  
         for t in ["<start>", "<sep>", "<end>"]:  
             decoded = decoded.replace(t, "")  
-        decoded = decoded.strip()  
-        if next_token == end_id or decoded.endswith(('.', '!', '?')):  
-            return decoded
+        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
+
     return sp.decode(generated)
 
-# 톤 불일치 체크
 def mismatch_tone(input_text, output_text):  
     if "ㅋㅋ" in input_text and not re.search(r'ㅋㅋ|ㅎ|재밌|놀|만나|맛집|여행', output_text):  
         return True  
@@ -340,9 +312,9 @@ def respond(input_text):
         return f"{summary}\n다른 궁금한 점 있으신가요?"
 
     # 일상 대화: 샘플링 + fallback
-    response = generate_text_topkp(model, input_text)
+    response = generate_text_greedy_strong(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_greedy_strong(model, input_text)
     return response
 
 @app.get("/generate", response_class=PlainTextResponse)