Update app.py

app.py

@@ -4,13 +4,11 @@ import os
 
 MODELS = {
     "Zephyr 7B Beta": "HuggingFaceH4/zephyr-7b-beta",
-    "DeepSeek Coder V2": "deepseek-ai/DeepSeek-Coder-V2-Instruct",
     "Meta Llama 3.1 8B": "meta-llama/Meta-Llama-3.1-8B-Instruct",
     "Meta-Llama 3.1 70B-Instruct": "meta-llama/Meta-Llama-3.1-70B-Instruct",
     "Microsoft": "microsoft/Phi-3-mini-4k-instruct",
     "Mixtral 8x7B": "mistralai/Mistral-7B-Instruct-v0.3",
     "Mixtral Nous-Hermes": "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
-    "Cohere Command R+": "CohereForAI/c4ai-command-r-plus",
     "Aya-23-35B": "CohereForAI/aya-23-35B"
 }
 
@@ -18,7 +16,7 @@ def get_client(model_name):
     model_id = MODELS[model_name]
     hf_token = os.getenv("HF_TOKEN")
     if not hf_token:
-        raise ValueError("HF_TOKEN environment variable is required")
+        raise ValueError("HF_TOKEN 환경 변수가 필요합니다.")
     return InferenceClient(model_id, token=hf_token)
 
 def respond(
@@ -73,28 +71,102 @@ def respond(
                         chat_history.append((message, partial_message))
                     yield chat_history
     except Exception as e:
-        error_message = f"An error occurred: {str(e)}"
+        error_message = f"오류가 발생했습니다: {str(e)}"
         chat_history.append((message, error_message))
         yield chat_history
 
 def clear_conversation():
     return []
 
+# Cohere Command R+ 전용 응답 함수
+from openai import OpenAI
+
+ACCESS_TOKEN = os.getenv("HF_TOKEN")
+
+cohere_client = OpenAI(
+    base_url="https://api-inference.huggingface.co/v1/",
+    api_key=ACCESS_TOKEN,
+)
+
+def cohere_respond(
+    message,
+    chat_history,
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+):
+    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})
+
+    messages.append({"role": "user", "content": message})
+
+    response = ""
+    
+    try:
+        for msg in cohere_client.chat.completions.create(
+            model="CohereForAI/c4ai-command-r-plus-08-2024",
+            max_tokens=max_tokens,
+            stream=True,
+            temperature=temperature,
+            top_p=top_p,
+            messages=messages,
+        ):
+            token = msg.choices[0].delta.content
+            response += token
+            if len(chat_history) > 0 and chat_history[-1][0] == message:
+                chat_history[-1] = (message, response)
+            else:
+                chat_history.append((message, response))
+            yield chat_history
+    except Exception as e:
+        error_message = f"오류가 발생했습니다: {str(e)}"
+        chat_history.append((message, error_message))
+        yield chat_history
+
 with gr.Blocks() as demo:
     gr.Markdown("# Prompting AI Chatbot")
     gr.Markdown("언어모델별 프롬프트 테스트 챗봇입니다.")
-
-    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(
+    
+    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="""반드시 한글로 답변할 것.
+너는 최고의 비서이다.
+내가 요구하는것들을 최대한 자세하고 정확하게 답변하라.
+""",
+                    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("대화 내역 지우기")
+    
+        msg.submit(respond, [msg, chatbot, model_name, max_tokens, temperature, top_p, system_message], chatbot)
+        submit_button.click(respond, [msg, chatbot, model_name, max_tokens, temperature, top_p, system_message], chatbot)
+        clear_button.click(clear_conversation, outputs=chatbot, queue=False)
+    
+    with gr.Tab("Cohere Command R+"):
+        with gr.Row():
+            cohere_system_message = gr.Textbox(
                 value="""반드시 한글로 답변할 것.
 너는 최고의 비서이다.
 내가 요구하는것들을 최대한 자세하고 정확하게 답변하라.
@@ -102,17 +174,33 @@ with gr.Blocks() as demo:
                 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("대화 내역 지우기")
-
-    msg.submit(respond, [msg, chatbot, model_name, max_tokens, temperature, top_p, system_message], chatbot)
-    submit_button.click(respond, [msg, chatbot, model_name, max_tokens, temperature, top_p, system_message], chatbot)
-    clear_button.click(clear_conversation, outputs=chatbot, queue=False)
+            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_msg.submit(
+            cohere_respond,
+            [cohere_msg, cohere_chatbot, cohere_system_message, cohere_max_tokens, cohere_temperature, cohere_top_p],
+            cohere_chatbot
+        )
+        cohere_submit_button.click(
+            cohere_respond,
+            [cohere_msg, cohere_chatbot, cohere_system_message, cohere_max_tokens, cohere_temperature, cohere_top_p],
+            cohere_chatbot
+        )
+        cohere_clear_button.click(clear_conversation, outputs=cohere_chatbot, queue=False)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()