Update app.py

app.py

@@ -126,86 +126,64 @@ def decode_sp_tokens(tokens):
     text = ''.join(tokens).replace('▁', ' ').strip()  
     return text  
 
-def generate_text_mirostat_top_p(model, prompt, max_len=100, max_gen=98,
-                                 temperature=1.0, min_len=20,
-                                 repetition_penalty=1.2, eta=0.1, m=100, p=0.9):
+def generate_text_top_p(model, prompt, max_len=100, max_gen=98,
+                        temperature=1.0, min_len=20,
+                        repetition_penalty=1.1, top_p=0.9):
     model_input = text_to_ids(f"<start> {prompt} <sep>")
     model_input = model_input[:max_len]
     generated = list(model_input)
-    text_so_far = []
-
-    tau = 5.0  # 초기 목표 surprise
 
     for step in range(max_gen):
-        pad_length = max(0, max_len - len(generated))
-        input_padded = np.pad(generated, (0, pad_length), constant_values=pad_id)
+        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_token_logits = logits[0, len(generated) - 1].numpy()
-
-        # 반복 페널티 적용
-        token_counts = {}
-        for t in generated:
-            token_counts[t] = token_counts.get(t, 0) + 1
-        for token_id, count in token_counts.items():
-            next_token_logits[token_id] /= (repetition_penalty ** count)
-
-        # 최소 길이 넘으면 종료 토큰 확률 낮추기
-        if len(generated) >= min_len:
-            next_token_logits[end_id] -= 5.0
-        next_token_logits[pad_id] -= 10.0
-
-        # 온도 조절
-        next_token_logits = next_token_logits / temperature
-
-        # --- 미로스타트 + Top-p 샘플링 시작 ---
-        logits_stable = next_token_logits - np.max(next_token_logits)
-        probs = np.exp(logits_stable)
+        next_logits = logits[0, len(generated) - 1].numpy()
+
+        # 반복 억제 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
+        probs = np.exp(next_logits - np.max(next_logits))
         probs /= probs.sum()
 
-        # 1. mirostat top-m 후보 추리기
-        sorted_indices = np.argsort(-probs)
-        top_indices = sorted_indices[:m]
-        top_probs = probs[top_indices]
-        top_probs /= top_probs.sum()
-
-        # 2. mirostat 샘플링
-        sampled_index = np.random.choice(top_indices, p=top_probs)
-        sampled_prob = probs[sampled_index]
-        observed_surprise = -np.log(sampled_prob + 1e-9)
-        tau += eta * (observed_surprise - tau)
-
-        # 3. top-p 필터링
-        sorted_top_indices = top_indices[np.argsort(-top_probs)]
-        sorted_top_probs = np.sort(top_probs)[::-1]
-        cumulative_probs = np.cumsum(sorted_top_probs)
-        cutoff = np.searchsorted(cumulative_probs, p, side='left') + 1
-        filtered_indices = sorted_top_indices[:cutoff]
-        filtered_probs = sorted_top_probs[:cutoff]
-        filtered_probs /= filtered_probs.sum()
+        # Top-p 필터링
+        sorted_idx = np.argsort(-probs)
+        sorted_probs = probs[sorted_idx]
+        cum_probs = np.cumsum(sorted_probs)
+        cutoff = np.searchsorted(cum_probs, top_p) + 1
 
-        # 4. 최종 토큰은 filtered 집합에서 다시 샘플링
-        final_token = np.random.choice(filtered_indices, p=filtered_probs)
+        filtered_idx = sorted_idx[:cutoff]
+        filtered_probs = sorted_probs[:cutoff]
+        filtered_probs /= filtered_probs.sum()
 
-        generated.append(int(final_token))
+        sampled = np.random.choice(filtered_idx, p=filtered_probs)
+        generated.append(int(sampled))
 
-        next_word = sp.id_to_piece(int(final_token))
-        text_so_far.append(next_word)
-        decoded_text = decode_sp_tokens(text_so_far)
+        decoded = sp.decode(generated)
+        for t in ["<start>", "<sep>", "<end>"]:
+            decoded = decoded.replace(t, "")
+        decoded = decoded.strip()
 
-        if len(generated) >= min_len and final_token == end_id:
+        if len(generated) >= min_len and (sampled == end_id or decoded.endswith(('.', '!', '?'))):
+            yield decoded
             break
-        if len(generated) >= min_len and decoded_text.endswith(('.', '!', '?', '<end>')):
-            break
-
-        yield decoded_text
 
 nickname = "사용자"
 
 def respond(message, chat_history):
     message = message.replace("@사용자1@", nickname)
     response = ""
-    for partial in generate_text_mirostat_top_p(model, message):
+    for partial in generate_text_top_p(model, message):
         response = partial
         yield response