Update api.py

api.py

@@ -124,14 +124,16 @@ _ = model(dummy_input)
 model.load_weights("InteractGPT.weights.h5")
 print("모델 가중치 로드 완료!")
 
-async def generate_text_mirostat_top_p(model, prompt, max_len=100, max_gen=98,    
+
+async def generate_text_mirostat_top_p_with_buffer(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):    
+                                 repetition_penalty=1.2, eta=0.1, m=100, p=0.9, buffer_size=3):    
     model_input = text_to_ids(f"<start> {prompt} <sep>")    
     model_input = model_input[:max_len]    
     generated = list(model_input)    
     
     tau = 5.0  # 초기 목표 surprise    
+    buffer_tokens = []    
     
     for step in range(max_gen):    
         pad_length = max(0, max_len - len(generated))    
@@ -177,18 +179,30 @@ async def generate_text_mirostat_top_p(model, prompt, max_len=100, max_gen=98,
     
         final_token = np.random.choice(filtered_indices, p=filtered_probs)    
     
-        # 특수 토큰 처리    
         if final_token == end_id:    
+            # 버퍼에 남은 거 다 출력    
+            if buffer_tokens:    
+                decoded = sp.decode(buffer_tokens)    
+                for token in ["<start>", "<sep>", "<end>"]:    
+                    decoded = decoded.replace(token, "")    
+                yield decoded.strip()    
             break    
+    
         if final_token in [start_id, pad_id] or sp.id_to_piece(final_token) == "<sep>":    
             continue    
     
-        # 정상 토큰만 추가 및 yield    
         generated.append(int(final_token))    
-        decoded = sp.decode([final_token])    
-        yield decoded
+        buffer_tokens.append(final_token)    
+    
+        if len(buffer_tokens) >= buffer_size or sp.id_to_piece(final_token).endswith("▁"):    
+            # 띄어쓰기 있는 토큰 나오거나 버퍼 꽉 찼으면 출력    
+            decoded = sp.decode(buffer_tokens)    
+            for token in ["<start>", "<sep>", "<end>"]:    
+                decoded = decoded.replace(token, "")    
+            yield decoded.strip()    
+            buffer_tokens = []
         
 @app.get("/generate")
 async def generate(request: Request):
     prompt = request.query_params.get("prompt", "안녕하세요")
-    return StreamingResponse(generate_text_stream(prompt), media_type="text/plain")
+    return StreamingResponse(generate_text_mirostat_top_p_with_buffer(prompt), media_type="text/plain")