Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 import openai
-import os
 
 # 제거할 모델들을 MODELS 사전에서 제외
 MODELS = {
@@ -17,7 +16,6 @@ MODELS = {
 # Cohere Command R+ 모델 ID 정의
 COHERE_MODEL = "CohereForAI/c4ai-command-r-plus-08-2024"
 
-
 def get_client(model_name, hf_token):
     """
     모델 이름에 맞춰 InferenceClient 생성.
@@ -35,91 +33,68 @@ def get_client(model_name, hf_token):
     return InferenceClient(model_id, token=hf_token)
 
 
-def respond(
-    message,
-    chat_history,
-    model_name,
-    max_tokens,
-    temperature,
-    top_p,
-    system_message,
-    hf_token,       # HF 토큰을 추가로 받음
+def respond_hf_qna(
+    question: str,
+    model_name: str,
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
+    system_message: str,
+    hf_token: str
 ):
+    """
+    HuggingFace 모델(Zephyr 등)에 대해 한 번의 질문(question)에 대한 답변을 반환하는 함수.
+    """
     try:
         client = get_client(model_name, hf_token)
     except ValueError as e:
-        chat_history.append((message, str(e)))
-        return chat_history
+        return f"오류: {str(e)}"
 
-    messages = [{"role": "system", "content": system_message}]
-    for human, assistant in chat_history:
-        messages.append({"role": "user", "content": human})
-        messages.append({"role": "assistant", "content": assistant})
-    messages.append({"role": "user", "content": message})
+    # 시스템 메시지 + 유저 질문을 한 번만 전달
+    messages = [
+        {"role": "system", "content": system_message},
+        {"role": "user", "content": question}
+    ]
 
     try:
-        if model_name == "Cohere Command R+":
-            # Cohere Command R+ 모델을 위한 비스트리밍 처리
-            response = client.chat_completion(
-                messages,
-                max_tokens=max_tokens,
-                temperature=temperature,
-                top_p=top_p,
-            )
-            assistant_message = response.choices[0].message.content
-            chat_history.append((message, assistant_message))
-            return chat_history
-        else:
-            # 다른 모델들을 위한 스트리밍 처리
-            stream = client.chat_completion(
-                messages,
-                max_tokens=max_tokens,
-                temperature=temperature,
-                top_p=top_p,
-                stream=True,
-            )
-            partial_message = ""
-            for response in stream:
-                if response.choices[0].delta.content is not None:
-                    partial_message += response.choices[0].delta.content
-                    if len(chat_history) > 0 and chat_history[-1][0] == message:
-                        chat_history[-1] = (message, partial_message)
-                    else:
-                        chat_history.append((message, partial_message))
-                    yield chat_history
+        # 스트리밍 대신 전체 답변(비스트리밍) 호출
+        response = client.chat_completion(
+            messages,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            stream=False,  # 스트리밍 비활성화
+        )
+        assistant_message = response.choices[0].message.content
+        return assistant_message
+
     except Exception as e:
-        error_message = f"오류가 발생했습니다: {str(e)}"
-        chat_history.append((message, error_message))
-        yield chat_history
-
-
-def cohere_respond(
-    message,
-    chat_history,
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    hf_token,   # HF 토큰 추가
+        return f"오류가 발생했습니다: {str(e)}"
+
+
+def respond_cohere_qna(
+    question: str,
+    system_message: str,
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
+    hf_token: str
 ):
+    """
+    Cohere Command R+ 모델을 이용해 한 번의 질문(question)에 대한 답변을 반환하는 함수.
+    """
     model_name = "Cohere Command R+"
     try:
         client = get_client(model_name, hf_token)
     except ValueError as e:
-        chat_history.append((message, str(e)))
-        return chat_history
-
-    messages = [{"role": "system", "content": system_message}]
-    for human, assistant in chat_history:
-        if human:
-            messages.append({"role": "user", "content": human})
-        if assistant:
-            messages.append({"role": "assistant", "content": assistant})
+        return f"오류: {str(e)}"
 
-    messages.append({"role": "user", "content": message})
+    messages = [
+        {"role": "system", "content": system_message},
+        {"role": "user", "content": question}
+    ]
 
     try:
-        # Cohere Command R+ 모델을 위한 비스트리밍 처리
         response_full = client.chat_completion(
             messages,
             max_tokens=max_tokens,
@@ -127,65 +102,50 @@ def cohere_respond(
             top_p=top_p,
         )
         assistant_message = response_full.choices[0].message.content
-        chat_history.append((message, assistant_message))
-        return chat_history
+        return assistant_message
     except Exception as e:
-        error_message = f"오류가 발생했습니다: {str(e)}"
-        chat_history.append((message, error_message))
-        return chat_history
-
-
-def chatgpt_respond(
-    message,
-    chat_history,
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    openai_token,  # openai 토큰 추가
+        return f"오류가 발생했습니다: {str(e)}"
+
+
+def respond_chatgpt_qna(
+    question: str,
+    system_message: str,
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
+    openai_token: str
 ):
     """
-    chatgpt용 응답. openai_token을 UI에서 입력받아 사용하도록 변경.
+    ChatGPT(OpenAI) 모델을 이용해 한 번의 질문(question)에 대한 답변을 반환하는 함수.
     """
     if not openai_token:
-        chat_history.append((message, "OpenAI API 토큰이 필요합니다."))
-        return chat_history
+        return "OpenAI API 토큰이 필요합니다."
 
-    # openai.api_key = os.getenv("OPENAI_API_KEY")  # 기존 코드 주석
-    openai.api_key = openai_token  # UI에서 받은 토큰 사용
+    openai.api_key = openai_token
 
-    messages = [{"role": "system", "content": system_message}]
-    for human, assistant in chat_history:
-        messages.append({"role": "user", "content": human})
-        messages.append({"role": "assistant", "content": assistant})
-    messages.append({"role": "user", "content": message})
+    messages = [
+        {"role": "system", "content": system_message},
+        {"role": "user", "content": question}
+    ]
 
     try:
         response = openai.ChatCompletion.create(
-            model="gpt-4o-mini",  # 또는 다른 모델 ID 사용
+            model="gpt-4o-mini",  # 필요한 경우 변경
             messages=messages,
             max_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
         )
         assistant_message = response.choices[0].message['content']
-        chat_history.append((message, assistant_message))
-        return chat_history
+        return assistant_message
     except Exception as e:
-        error_message = f"오류가 발생했습니다: {str(e)}"
-        chat_history.append((message, error_message))
-        return chat_history
-
-
-def clear_conversation():
-    return []
+        return f"오류가 발생했습니다: {str(e)}"
 
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Prompting AI Chatbot")
-    gr.Markdown("언어모델별 프롬프트 테스트 챗봇입니다.")
+    gr.Markdown("# Prompting AI - 일반 문답형 데모")
+    gr.Markdown("언어모델별 문답형 테스트 데모입니다. 한 번에 한 질문씩만 주고받습니다.")
 
-    # --- 토큰 입력 UI 추가 ---
     with gr.Row():
         hf_token_box = gr.Textbox(
             label="HuggingFace 토큰 (비공개)",
@@ -198,127 +158,108 @@ with gr.Blocks() as demo:
             placeholder="OpenAI API 토큰을 입력하세요..."
         )
 
+    # --- Tab: 일반 모델 ---
     with gr.Tab("일반 모델"):
-        with gr.Row():
-            with gr.Column(scale=1):
-                model_name = gr.Radio(
-                    choices=list(MODELS.keys()),
-                    label="Language Model",
-                    value="Zephyr 7B Beta"
-                )
-                max_tokens = gr.Slider(minimum=0, maximum=2000, value=500, step=100, label="Max Tokens")
-                temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.05, label="Temperature")
-                top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p")
-                system_message = gr.Textbox(
-                    value="""반드시 한글로 답변할 것.
+        model_name = gr.Radio(
+            choices=list(MODELS.keys()),
+            label="Language Model (HuggingFace)",
+            value="Zephyr 7B Beta"
+        )
+        max_tokens = gr.Slider(minimum=0, maximum=2000, value=500, step=100, label="Max Tokens")
+        temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.05, label="Temperature")
+        top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p")
+        system_message = gr.Textbox(
+            value="""반드시 한글로 답변할 것.
 너는 최고의 비서이다.
 내가 요구하는것들을 최대한 자세하고 정확하게 답변하라.
 """,
-                    label="System Message",
-                    lines=3
-                )
-    
-            with gr.Column(scale=2):
-                chatbot = gr.Chatbot()
-                msg = gr.Textbox(label="메세지를 입력하세요")
-                with gr.Row():
-                    submit_button = gr.Button("전송")
-                    clear_button = gr.Button("대화 내역 지우기")
-    
-        # respond 함수에 hf_token 인자를 전달하도록 수정
-        inputs_for_normal = [
-            msg, 
-            chatbot, 
-            model_name, 
-            max_tokens, 
-            temperature, 
-            top_p, 
-            system_message,
-            hf_token_box
-        ]
-        msg.submit(respond, inputs_for_normal, chatbot)
-        submit_button.click(respond, inputs_for_normal, chatbot)
-        clear_button.click(clear_conversation, outputs=chatbot, queue=False)
-    
+            label="System Message",
+            lines=3
+        )
+
+        question_input = gr.Textbox(label="질문을 입력하세요")
+        answer_output = gr.Textbox(label="답변", interactive=False)
+
+        # '전송' 버튼 클릭 시 HF 모델 QnA 함수 호출
+        submit_button = gr.Button("전송")
+
+        submit_button.click(
+            fn=respond_hf_qna,
+            inputs=[
+                question_input,
+                model_name,
+                max_tokens,
+                temperature,
+                top_p,
+                system_message,
+                hf_token_box
+            ],
+            outputs=answer_output
+        )
+
+    # --- Tab: Cohere Command R+ ---
     with gr.Tab("Cohere Command R+"):
-        with gr.Row():
-            cohere_system_message = gr.Textbox(
-                value="""반드시 한글로 답변할 것.
+        cohere_system_message = gr.Textbox(
+            value="""반드시 한글로 답변할 것.
 너는 최고의 비서이다.
 내가 요구하는것들을 최대한 자세하고 정확하게 답변하라.
 """,
-                label="System Message",
-                lines=3
-            )
-            cohere_max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens")
-            cohere_temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
-            cohere_top_p = gr.Slider(
-                minimum=0.1,
-                maximum=1.0,
-                value=0.95,
-                step=0.05,
-                label="Top-P",
-            )
-        
-        cohere_chatbot = gr.Chatbot(height=600)
-        cohere_msg = gr.Textbox(label="메세지를 입력하세요")
-        with gr.Row():
-            cohere_submit_button = gr.Button("전송")
-            cohere_clear_button = gr.Button("대화 내역 지우기")
-        
-        # cohere_respond 함수에 hf_token 인자를 전달하도록 수정
-        inputs_for_cohere = [
-            cohere_msg, 
-            cohere_chatbot, 
-            cohere_system_message, 
-            cohere_max_tokens, 
-            cohere_temperature, 
-            cohere_top_p,
-            hf_token_box
-        ]
-        cohere_msg.submit(cohere_respond, inputs_for_cohere, cohere_chatbot)
-        cohere_submit_button.click(cohere_respond, inputs_for_cohere, cohere_chatbot)
-        cohere_clear_button.click(clear_conversation, outputs=cohere_chatbot, queue=False)
-    
+            label="System Message",
+            lines=3
+        )
+        cohere_max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max Tokens")
+        cohere_temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
+        cohere_top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-P")
+
+        cohere_question_input = gr.Textbox(label="질문을 입력하세요")
+        cohere_answer_output = gr.Textbox(label="답변", interactive=False)
+
+        cohere_submit_button = gr.Button("전송")
+
+        cohere_submit_button.click(
+            fn=respond_cohere_qna,
+            inputs=[
+                cohere_question_input,
+                cohere_system_message,
+                cohere_max_tokens,
+                cohere_temperature,
+                cohere_top_p,
+                hf_token_box
+            ],
+            outputs=cohere_answer_output
+        )
+
+    # --- Tab: ChatGPT ---
     with gr.Tab("ChatGPT"):
-        with gr.Row():
-            chatgpt_system_message = gr.Textbox(
-                value="""반드시 한글로 답변할 것.
+        chatgpt_system_message = gr.Textbox(
+            value="""반드시 한글로 답변할 것.
 너는 ChatGPT, OpenAI에서 개발한 언어 모델이다.
 내가 요구하는 것을 최대한 자세하고 정확하게 답변하라.
 """,
-                label="System Message",
-                lines=3
-            )
-            chatgpt_max_tokens = gr.Slider(minimum=1, maximum=4096, value=1024, step=1, label="Max Tokens")
-            chatgpt_temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.05, label="Temperature")
-            chatgpt_top_p = gr.Slider(
-                minimum=0.1,
-                maximum=1.0,
-                value=0.95,
-                step=0.05,
-                label="Top-P",
-            )
-        
-        chatgpt_chatbot = gr.Chatbot(height=600)
-        chatgpt_msg = gr.Textbox(label="메세지를 입력하세요")
-        with gr.Row():
-            chatgpt_submit_button = gr.Button("전송")
-            chatgpt_clear_button = gr.Button("대화 내역 지우기")
-        
-        # chatgpt_respond 함수에 openai_token 인자를 전달하도록 수정
-        inputs_for_chatgpt = [
-            chatgpt_msg, 
-            chatgpt_chatbot, 
-            chatgpt_system_message, 
-            chatgpt_max_tokens, 
-            chatgpt_temperature, 
-            chatgpt_top_p,
-            openai_token_box
-        ]
-        chatgpt_msg.submit(chatgpt_respond, inputs_for_chatgpt, chatgpt_chatbot)
-        chatgpt_submit_button.click(chatgpt_respond, inputs_for_chatgpt, chatgpt_chatbot)
-        chatgpt_clear_button.click(clear_conversation, outputs=chatgpt_chatbot, queue=False)
+            label="System Message",
+            lines=3
+        )
+        chatgpt_max_tokens = gr.Slider(minimum=1, maximum=4096, value=1024, step=1, label="Max Tokens")
+        chatgpt_temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.05, label="Temperature")
+        chatgpt_top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-P")
+
+        chatgpt_question_input = gr.Textbox(label="질문을 입력하세요")
+        chatgpt_answer_output = gr.Textbox(label="답변", interactive=False)
+
+        chatgpt_submit_button = gr.Button("전송")
+
+        chatgpt_submit_button.click(
+            fn=respond_chatgpt_qna,
+            inputs=[
+                chatgpt_question_input,
+                chatgpt_system_message,
+                chatgpt_max_tokens,
+                chatgpt_temperature,
+                chatgpt_top_p,
+                openai_token_box
+            ],
+            outputs=chatgpt_answer_output
+        )
 
 if __name__ == "__main__":
     demo.launch()