Update app.py

app.py

@@ -3,9 +3,11 @@ from huggingface_hub import InferenceClient
 import os
 from threading import Event
 
+# Hugging Face API Token을 환경 변수로부터 가져옴
 hf_token = os.getenv("HF_TOKEN")
 stop_event = Event()
 
+# 모델 목록 정의
 models = {
     "deepseek-ai/DeepSeek-Coder-V2-Instruct": "(한국회사)DeepSeek-Coder-V2-Instruct",
     "meta-llama/Meta-Llama-3.1-8B-Instruct": "Meta-Llama-3.1-8B-Instruct",
@@ -13,13 +15,16 @@ models = {
     "CohereForAI/c4ai-command-r-plus": "Cohere Command-R Plus"
 }
 
+# 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": "user", "content": message}
@@ -27,6 +32,9 @@ def respond(message, system_message, max_tokens, temperature, top_p, selected_mo
     
     try:
         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,
@@ -38,97 +46,77 @@ def respond(message, system_message, max_tokens, temperature, top_p, selected_mo
                 break
             if chunk:
                 response += chunk
-                yield [(message, response)]
+                total_tokens_used += len(chunk.split())  # 청크당 사용된 토큰 수 추산
+                yield [(message, response, f"사용된 토큰 수: {total_tokens_used}/{max_tokens}")]
         
     except Exception as e:
-        yield [(message, f"오류 발생: {str(e)}")]
-        
+        yield [(message, f"오류 발생: {str(e)}", "에러 처리 필요")]
+
+# 이전 응답을 확인하는 함수
 def get_last_response(chatbot):
     if chatbot and len(chatbot) > 0:
         return chatbot[-1][1]
-    return ""
+    return None
 
-def continue_writing(chatbot, system_message, max_tokens, temperature, top_p, selected_model):
-    last_response = get_last_response(chatbot)
-    stop_event.clear()
-    client = get_client(selected_model)
+# 프롬프트 비교 및 최적화를 위한 히스토리 기록 추가
+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 인터페이스 함수 정의
+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 = f"이전 응답을 이어서 작성해주세요. 이전 응답: {last_response}"
-    messages = [
-        {"role": "system", "content": system_message},
-        {"role": "user", "content": prompt}
-    ]
+    # 프롬프트와 결과를 히스토리에 추가
+    prompt_history.add_entry(
+        message,
+        result[0][1],  # 모델 응답
+        selected_model,
+        {"max_tokens": max_tokens, "temperature": temperature, "top_p": top_p}
+    )
     
-    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 chatbot + [("계속 작성", response)]
-        
-    except Exception as e:
-        yield chatbot + [("계속 작성", f"오류 발생: {str(e)}")]
+    return result
 
-def stop_generation():
-    stop_event.set()
-    return "생성이 중단되었습니다."
+# 히스토리 확인용 함수
+def view_history():
+    return prompt_history.get_history()
 
+# Gradio UI 구성
 with gr.Blocks() as demo:
-    gr.Markdown("# 프롬프트 최적화 Playground")
+    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")
     
-    gr.Markdown("""
-    **주의사항:**
-    - '전송' 버튼을 클릭하거나 입력 필드에서 Shift+Enter를 눌러 메시지를 전송할 수 있습니다.
-    - Enter 키는 줄바꿈으로 작동합니다.
-    - 입력한 내용에 대해서만 응답하도록 설정되어 있지만, 모델이 때때로 예상치 못한 방식으로 응답할 수 있습니다.
-    """)
+    response_output = gr.Textbox(label="모델 응답")
+    token_usage = gr.Textbox(label="토큰 사용량")
+    history_button = gr.Button("히스토리 보기")
     
-    with gr.Row():
-        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 (핵 샘플링)")
-            
-            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를 눌러 전송합니다."
-            )
-            
-            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, outputs=[prompt])
-    clear.click(lambda: None, outputs=[chatbot])
+    # 버튼을 눌러 응답을 받는 함수 연결
+    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)
 
-if __name__ == "__main__":
-    if not hf_token:
-        print("경고: HF_TOKEN 환경 변수가 설정되지 않았습니다. 일부 모델에 접근할 수 없을 수 있습니다.")
-    demo.launch()
+# UI 실행
+demo.launch()