Update app.py

app.py

@@ -19,14 +19,14 @@ def get_client(model):
     return InferenceClient(model=model, token=hf_token)
 
 # 응답 생성 함수
-def respond(message, system_message, max_tokens, temperature, top_p, selected_model):
+def respond(prompt, system_message, max_tokens, temperature, top_p, selected_model):
     stop_event.clear()
     client = get_client(selected_model)
     
     # 프롬프트 설정 - 시스템 메시지를 자유롭게 설정 가능
     messages = [
         {"role": "system", "content": system_message},
-        {"role": "user", "content": message}
+        {"role": "user", "content": prompt}
     ]
     
     try:
@@ -46,42 +46,68 @@ def respond(message, system_message, max_tokens, temperature, top_p, selected_mo
             if chunk:
                 response += chunk
                 total_tokens_used += len(chunk.split())  # 청크당 사용된 토큰 수 추산
-                yield [(message, response, f"사용된 토큰 수: {total_tokens_used}/{max_tokens}")]
         
+        return response
+    
     except Exception as e:
-        yield [(message, f"오류 발생: {str(e)}", "에러 처리 필요")]
+        return f"오류 발생: {str(e)}"
 
-# 이전 응답을 확인하는 함수
-def get_last_response(chatbot):
-    if chatbot and len(chatbot) > 0:
-        return chatbot[-1][1]
-    return None
+# 대화 계속 함수
+def continue_writing(chatbot, system_message, max_tokens, temperature, top_p, selected_model):
+    last_message = chatbot[-1][1] if chatbot else ""
+    return respond(last_message, system_message, max_tokens, temperature, top_p, selected_model)
 
-# Gradio UI 구성
-def gradio_interface(message, system_message, max_tokens, temperature, top_p, selected_model):
-    result = None
-    for output in respond(message, system_message, max_tokens, temperature, top_p, selected_model):
-        result = output
-    return result
+# 응답 중단 함수
+def stop_generation():
+    stop_event.set()
 
+# Gradio UI 구성
 with gr.Blocks() as demo:
+    gr.Markdown("# 프롬프트 최적화 Playground")
+    
+    gr.Markdown("""
+    **주의사항:**
+    - '전송' 버튼을 클릭하거나 입력 필드에서 Shift+Enter를 눌러 메시지를 전송할 수 있습니다.
+    - Enter 키는 줄바꿈으로 작동합니다.
+    - 입력한 내용에 대해서만 응답하도록 설정되어 있지만, 모델이 때때로 예상치 못한 방식으로 응답할 수 있습니다.
+    """)
+    
     with gr.Row():
-        with gr.Column():
-            selected_model = gr.Dropdown(choices=list(models.keys()), value="deepseek-ai/DeepSeek-Coder-V2-Instruct", label="모델 선택")
-            system_message = gr.Textbox(label="시스템 메시지", value="이 메시지를 기준으로 대화 흐름을 설정합니다.")
-            message = gr.Textbox(label="사용자 메시지")
+        with gr.Column(scale=1):
+            with gr.Accordion("모델 설정", open=True):
+                model = gr.Radio(list(models.keys()), value=list(models.keys())[0], label="언어 모델 선택", info="사용할 언어 모델을 선택하세요")
+                max_tokens = gr.Slider(minimum=1, maximum=2000, value=500, step=100, label="최대 새 토큰 수")
+                temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.05, label="온도")
+                top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.90, step=0.05, label="Top-p (핵 샘플링)")
             
-            max_tokens = gr.Slider(minimum=10, maximum=512, value=128, label="최대 토큰 수")
-            temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, label="Temperature")
-            top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.9, label="Top-p")
+            system_message = gr.Textbox(
+                value="너는 나의 최고의 비서이다.\n내가 요구하는것들을 최대한 자세하고 정확하게 답변하라.\n반드시 한글로 답변할것.\n사용자의 입력 내용에만 직접적으로 답변하고, 추가 질문을 만들거나 입력을 확장하지 마라.",
+                label="시스템 메시지",
+                lines=5
+            )
+        
+        with gr.Column(scale=2):
+            chatbot = gr.Chatbot(height=400, label="대화 결과")
+            prompt = gr.Textbox(
+                label="내용 입력", 
+                lines=3,
+                placeholder="메시지를 입력하세요. 전송 버튼을 클릭하거나 Shift+Enter를 눌러 전송합니다."
+            )
             
-            submit_button = gr.Button("응답 생성")
-        with gr.Column():
-            chatbot = gr.Chatbot()
-            token_usage = gr.Textbox(label="토큰 사용량", interactive=False)
-    
-    # 버튼을 눌러 응답을 받는 함수 연결
-    submit_button.click(gradio_interface, inputs=[message, system_message, max_tokens, temperature, top_p, selected_model], outputs=[chatbot, token_usage])
+            with gr.Row():
+                send = gr.Button("전송")
+                continue_btn = gr.Button("계속 작성")
+                stop = gr.Button("🛑 생성 중단")
+                clear = gr.Button("🗑️ 대화 내역 지우기")
+
+    # Event handlers
+    send.click(respond, inputs=[prompt, system_message, max_tokens, temperature, top_p, model], outputs=[chatbot])
+    prompt.submit(respond, inputs=[prompt, system_message, max_tokens, temperature, top_p, model], outputs=[chatbot])
+    continue_btn.click(continue_writing, 
+                       inputs=[chatbot, system_message, max_tokens, temperature, top_p, model], 
+                       outputs=[chatbot])
+    stop.click(stop_generation)
+    clear.click(lambda: None, outputs=[chatbot])
 
 # UI 실행
 demo.launch()