Update app.py

app.py

@@ -3,7 +3,6 @@ from huggingface_hub import InferenceClient
 import os
 from threading import Event
 
-# Hugging Face API Token을 환경 변수로부터 가져옴
 hf_token = os.getenv("HF_TOKEN")
 stop_event = Event()
 
@@ -15,18 +14,18 @@ models = {
     "CohereForAI/c4ai-command-r-plus": "Cohere Command-R Plus"
 }
 
-# Inference 클라이언트를 반환하는 함수
+# Inference 클라이언트 반환
 def get_client(model):
     return InferenceClient(model=model, token=hf_token)
 
-# 메시지 응답 생성 함수
+# 응답 생성 함수
 def respond(message, system_message, max_tokens, temperature, top_p, selected_model):
     stop_event.clear()
     client = get_client(selected_model)
     
-    # 프롬프트 설정
+    # 프롬프트 설정 - 시스템 메시지를 자유롭게 설정 가능
     messages = [
-        {"role": "system", "content": system_message + "\n주어진 입력에만 정확히 답변하세요. 추가 질문을 만들거나 입력을 확장하지 마세요."},
+        {"role": "system", "content": system_message},
         {"role": "user", "content": message}
     ]
     
@@ -34,7 +33,7 @@ def respond(message, system_message, max_tokens, temperature, top_p, selected_mo
         response = ""
         total_tokens_used = 0  # 사용된 토큰 수 추적
         
-        # 모델에서 응답을 청크 단위로 스트리밍
+        # 모델에서 응답을 스트리밍
         for chunk in client.text_generation(
             prompt="\n".join([f"{m['role']}: {m['content']}" for m in messages]),
             max_new_tokens=max_tokens,
@@ -58,65 +57,31 @@ def get_last_response(chatbot):
         return chatbot[-1][1]
     return None
 
-# 프롬프트 비교 및 최적화를 위한 히스토리 기록 추가
-class PromptHistory:
-    def __init__(self):
-        self.history = []
-
-    def add_entry(self, prompt, response, model, settings):
-        self.history.append({
-            "prompt": prompt,
-            "response": response,
-            "model": model,
-            "settings": settings
-        })
-
-    def get_history(self):
-        return self.history
-
-# 히스토리 객체 생성
-prompt_history = PromptHistory()
-
-# Gradio 인터페이스 함수 정의
+# 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
-    
-    # 프롬프트와 결과를 히스토리에 추가
-    prompt_history.add_entry(
-        message,
-        result[0][1],  # 모델 응답
-        selected_model,
-        {"max_tokens": max_tokens, "temperature": temperature, "top_p": top_p}
-    )
-    
     return result
 
-# 히스토리 확인용 함수
-def view_history():
-    return prompt_history.get_history()
-
-# Gradio UI 구성
 with gr.Blocks() as demo:
-    selected_model = gr.Dropdown(choices=list(models.keys()), label="모델 선택")
-    message = gr.Textbox(label="사용자 메시지")
-    system_message = gr.Textbox(label="시스템 메시지", value="이 메시지를 기준으로 대화 흐름을 설정합니다.")
-    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")
-    
-    response_output = gr.Textbox(label="모델 응답")
-    token_usage = gr.Textbox(label="토큰 사용량")
-    history_button = gr.Button("히스토리 보기")
+    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="사용자 메시지")
+            
+            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")
+            
+            submit_button = gr.Button("응답 생성")
+        with gr.Column():
+            chatbot = gr.Chatbot()
+            token_usage = gr.Textbox(label="토큰 사용량", interactive=False)
     
     # 버튼을 눌러 응답을 받는 함수 연결
-    submit_button = gr.Button("응답 생성")
-    submit_button.click(gradio_interface, inputs=[message, system_message, max_tokens, temperature, top_p, selected_model], outputs=[response_output, token_usage])
-    
-    # 히스토리 보기 기능 연결
-    history_output = gr.Textbox(label="히스토리", interactive=False)
-    history_button.click(view_history, outputs=history_output)
+    submit_button.click(gradio_interface, inputs=[message, system_message, max_tokens, temperature, top_p, selected_model], outputs=[chatbot, token_usage])
 
 # UI 실행
 demo.launch()