Update app.py

app.py

@@ -21,15 +21,14 @@ MAX_HISTORY_LENGTH = 5  # 히스토리에 유지할 최대 대화 수
 def truncate_history(history):
     return history[-MAX_HISTORY_LENGTH:] if len(history) > MAX_HISTORY_LENGTH else history
 
-def respond(message, history, system_message, max_tokens, temperature, top_p, selected_model):
+def respond(message, system_message, max_tokens, temperature, top_p, selected_model):
     stop_event.clear()
     client = InferenceClient(model=selected_model, token=hf_token)
     
-    truncated_history = truncate_history(history)
-    
-    messages = [{"role": "system", "content": system_message + "\n사용자의 입력에만 직접적으로 답변하세요. 추가 질문을 생성하거나 사용자의 입력을 확장하지 마세요."}]
-    messages.extend([{"role": "user" if i % 2 == 0 else "assistant", "content": m} for h in truncated_history for i, m in enumerate(h) if m])
-    messages.append({"role": "user", "content": message})
+    messages = [
+        {"role": "system", "content": system_message},
+        {"role": "user", "content": message}
+    ]
     
     try:
         response = ""
@@ -44,26 +43,15 @@ def respond(message, history, system_message, max_tokens, temperature, top_p, se
                 break
             if chunk:
                 response += chunk
-                if response.startswith(message):
-                    response = response[len(message):].lstrip()
-                yield truncated_history + [(message, response)]
+                yield [(message, response)]
         
     except Exception as e:
-        yield truncated_history + [(message, f"오류 발생: {str(e)}")]
-
-def continue_writing(message, history, system_message, max_tokens, temperature, top_p, selected_model):
-    if not history:
-        yield [("시스템", "대화 내역이 없습니다.")]
-        return
+        yield [(message, f"오류 발생: {str(e)}")]
 
-    truncated_history = truncate_history(history)
-    last_assistant_message = truncated_history[-1][1]
-    
-    prompt = f"이전 대화를 간단히 요약하고 이어서 작성해주세요. 마지막 응답: {last_assistant_message[:100]}..."
+def stop_generation():
+    stop_event.set()
+    return "생성이 중단되었습니다."
     
-    for response in respond(prompt, truncated_history[:-1], system_message, max_tokens, temperature, top_p, selected_model):
-        yield response
-
 def stop_generation():
     stop_event.set()
     return "생성이 중단되었습니다."
@@ -74,14 +62,42 @@ def regenerate(chat_history, system_message, max_tokens, temperature, top_p, sel
     last_user_message = chat_history[-1][0]
     return respond(last_user_message, chat_history[:-1], system_message, max_tokens, temperature, top_p, selected_model)
 
+def continue_writing(last_response, system_message, max_tokens, temperature, top_p, selected_model):
+    stop_event.clear()
+    client = InferenceClient(model=selected_model, token=hf_token)
+    
+    prompt = f"이전 응답을 이어서 작성해주세요. 이전 응답: {last_response}"
+    messages = [
+        {"role": "system", "content": system_message},
+        {"role": "user", "content": prompt}
+    ]
+    
+    try:
+        response = last_response
+        for chunk in client.text_generation(
+            prompt="\n".join([f"{m['role']}: {m['content']}" for m in messages]),
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            stream=True
+        ):
+            if stop_event.is_set():
+                break
+            if chunk:
+                response += chunk
+                yield [("계속 작성", response)]
+        
+    except Exception as e:
+        yield [("계속 작성", f"오류 발생: {str(e)}")]
+
+# Gradio 인터페이스 수정
 with gr.Blocks() as demo:
     chatbot = gr.Chatbot()
-    msg = gr.Textbox(label="메시지 입력", placeholder="메시지를 입력하세요. Enter로 전송, Shift+Enter로 줄바꿈")
+    msg = gr.Textbox(label="메시지 입력")
 
     with gr.Row():
         send = gr.Button("전송")
         continue_btn = gr.Button("계속 작성")
-        regenerate_btn = gr.Button("🔄 재생성")
         stop = gr.Button("🛑 생성 중단")
         clear = gr.Button("🗑️ 대화 내역 지우기")
 
@@ -97,12 +113,11 @@ with gr.Blocks() as demo:
         model = gr.Radio(list(models.keys()), value=list(models.keys())[0], label="언어 모델 선택", info="사용할 언어 모델을 선택하세요")
 
     # Event handlers
-    msg.submit(respond, [msg, chatbot, system_message, max_tokens, temperature, top_p, model], [chatbot])
-    send.click(respond, [msg, chatbot, system_message, max_tokens, temperature, top_p, model], [chatbot])
+    send.click(respond, inputs=[msg, system_message, max_tokens, temperature, top_p, model], outputs=[chatbot])
+    msg.submit(respond, inputs=[msg, system_message, max_tokens, temperature, top_p, model], outputs=[chatbot])
     continue_btn.click(continue_writing, 
-        inputs=[msg, chatbot, system_message, max_tokens, temperature, top_p, model],
-        outputs=[chatbot])
-    regenerate_btn.click(regenerate, [chatbot, system_message, max_tokens, temperature, top_p, model], [chatbot])
+                       inputs=[lambda: chatbot[-1][1] if chatbot else "", system_message, max_tokens, temperature, top_p, model], 
+                       outputs=[chatbot])
     stop.click(stop_generation, outputs=[msg])
     clear.click(lambda: None, outputs=[chatbot])