Update api.py

api.py

@@ -171,47 +171,49 @@ def is_greedy_response_acceptable(text):
 
     return True
 
-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 _ 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)  
+def generate_text_with_temp_and_rep_penalty(model, prompt, max_len=100, max_gen=98,
+                                            repetition_penalty=1.2, temperature=0.7,
+                                            min_len=20):
+    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()
+
+        # Repetition penalty
+        for t in set(generated):
+            count = generated.count(t)
             next_logits[t] /= (repetition_penalty ** count)
 
-        # Stop token filtering  
-        stop_tokens = ["음", "어", "그", "뭐지", "..."]
-        for tok in stop_tokens:
-            tok_id = sp.piece_to_id(tok)
-            next_logits[tok_id] -= 5.0
+        # Temperature scaling
+        next_logits = next_logits / temperature
 
-        next_logits[pad_id] -= 10.0
-        next_token = np.argmax(next_logits)
-        generated.append(int(next_token))  
+        # Softmax로 확률 계산
+        exp_logits = np.exp(next_logits - np.max(next_logits))
+        probs = exp_logits / exp_logits.sum()
 
-        decoded = sp.decode(generated)  
-        for t in ["<start>", "<sep>", "<end>"]:  
-            decoded = decoded.replace(t, "")  
+        # 다음 토큰 샘플링
+        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, "")
         decoded = decoded.strip()
 
-        if len(generated) >= min_len and (next_token == end_id or decoded.endswith(('요', '다', '.', '!', '?'))):  
+        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  
@@ -333,9 +335,9 @@ def respond(input_text):
         return f"{summary}\n다른 궁금한 점 있으신가요?"
 
     # 일상 대화: 샘플링 + fallback
-    response = generate_text_greedy_strong(model, input_text)
+    response = generate_text_with_temp_and_rep_penalty(model, input_text)
     if not is_valid_response(response) or mismatch_tone(input_text, response):
-        response = generate_text_greedy_strong(model, input_text)
+        response = generate_text_with_temp_and_rep_penalty(model, input_text)
     return response
 
 @app.get("/generate", response_class=PlainTextResponse)