Update api.py

api.py

@@ -138,81 +138,60 @@ _ = model(dummy_input)  # 모델이 빌드됨
 model.load_weights("InteractGPT.weights.h5")      
 print("모델 가중치 로드 완료!")      
 
+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)
 
-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)      
-      
-    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)      
-        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)      
-        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()      
-      
-        # 4. 최종 토큰 샘플링      
-        final_token = np.random.choice(filtered_indices, p=filtered_probs)      
-        generated.append(int(final_token))      
-      
-        decoded_text = sp.decode(generated)      
-        # 특수 토큰 제거      
-        for token in ["<start>", "<sep>", "<end>"]:      
-            decoded_text = decoded_text.replace(token, "")      
-      
-        decoded_text = decoded_text.strip()      
-      
-        if len(generated) >= min_len and (final_token == end_id or decoded_text.endswith(('.', '!', '?', '<end>'))):      
-            yield decoded_text      
-            break      
+    for step 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()
+
+        # 반복 억제 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()
+
+        # 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
+
+        filtered_idx = sorted_idx[:cutoff]
+        filtered_probs = sorted_probs[:cutoff]
+        filtered_probs /= filtered_probs.sum()
+
+        sampled = np.random.choice(filtered_idx, p=filtered_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(('.', '!', '?'))):
+            yield decoded
+            break
               
 async def async_generator_wrapper(prompt: str):
-    gen = generate_text_mirostat_top_p(model, prompt)
+    gen = generate_text_top_p(model, prompt)
     for text_piece in gen:
         yield text_piece
         await asyncio.sleep(0.1)