Update app.py

app.py

@@ -126,13 +126,15 @@ def decode_sp_tokens(tokens):
     text = ''.join(tokens).replace('▁', ' ').strip()  
     return text  
 
-def generate_text_mirostat(model, prompt, max_len=100, max_gen=98, temperature=1.0, min_len=20, repetition_penalty=1.2, eta=0.1, m=100):
+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):
     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 (정보량)
+    tau = 5.0  # 초기 목표 surprise
 
     for step in range(max_gen):
         pad_length = max(0, max_len - len(generated))
@@ -141,7 +143,7 @@ def generate_text_mirostat(model, prompt, max_len=100, max_gen=98, temperature=1
         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
@@ -156,45 +158,54 @@ def generate_text_mirostat(model, prompt, max_len=100, max_gen=98, temperature=1
         # 온도 조절
         next_token_logits = next_token_logits / temperature
 
-        # --- 미로스타트 샘플링 ---
-        logits_stable = next_token_logits - np.max(next_token_logits)  # 안정화
+        # --- 미로스타트 + Top-p 샘플링 시작 ---
+        logits_stable = next_token_logits - np.max(next_token_logits)
         probs = np.exp(logits_stable)
         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 업데이트
         tau += eta * (observed_surprise - tau)
 
-        generated.append(int(sampled_index))
+        # 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()
+
+        # 4. 최종 토큰은 filtered 집합에서 다시 샘플링
+        final_token = np.random.choice(filtered_indices, p=filtered_probs)
+
+        generated.append(int(final_token))
 
-        next_word = sp.id_to_piece(int(sampled_index))
+        next_word = sp.id_to_piece(int(final_token))
         text_so_far.append(next_word)
         decoded_text = decode_sp_tokens(text_so_far)
 
-        if len(generated) >= min_len and sampled_index == end_id:
+        if len(generated) >= min_len and final_token == end_id:
             break
         if len(generated) >= min_len and decoded_text.endswith(('.', '!', '?', '<end>')):
             break
 
         yield decoded_text
 
-
-import gradio as gr
-
 nickname = "사용자"
 
 def respond(message, chat_history):
     message = message.replace("@사용자1@", nickname)
     response = ""
-    for partial in generate_text_mirostat(model, message):
+    for partial in generate_text_mirostat_top_p(model, message):
         response = partial
         yield response