Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from huggingface_hub import InferenceClient
 import os
 
+# 제거할 모델들을 MODELS 사전에서 제외
 MODELS = {
     "Zephyr 7B Beta": "HuggingFaceH4/zephyr-7b-beta",
     "Meta Llama 3.1 8B": "meta-llama/Meta-Llama-3.1-8B-Instruct",
@@ -12,11 +13,19 @@ MODELS = {
     "Aya-23-35B": "CohereForAI/aya-23-35B"
 }
 
+# Cohere Command R+ 모델 ID 정의
+COHERE_MODEL = "CohereForAI/c4ai-command-r-plus-08-2024"
+
 def get_client(model_name):
-    model_id = MODELS[model_name]
     hf_token = os.getenv("HF_TOKEN")
     if not hf_token:
         raise ValueError("HF_TOKEN 환경 변수가 필요합니다.")
+    if model_name in MODELS:
+        model_id = MODELS[model_name]
+    elif model_name == "Cohere Command R+":
+        model_id = COHERE_MODEL
+    else:
+        raise ValueError("유효하지 않은 모델 이름입니다.")
     return InferenceClient(model_id, token=hf_token)
 
 def respond(
@@ -41,8 +50,8 @@ def respond(
     messages.append({"role": "user", "content": message})
 
     try:
-        if "Cohere" in model_name:
-            # Cohere 모델을 위한 비스트리밍 처리
+        if model_name == "Cohere Command R+":
+            # Cohere Command R+ 모델을 위한 비스트리밍 처리
             response = client.chat_completion(
                 messages,
                 max_tokens=max_tokens,
@@ -51,7 +60,7 @@ def respond(
             )
             assistant_message = response.choices[0].message.content
             chat_history.append((message, assistant_message))
-            yield chat_history
+            return chat_history
         else:
             # 다른 모델들을 위한 스트리밍 처리
             stream = client.chat_completion(
@@ -78,57 +87,6 @@ def respond(
 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("언어모델별 프롬프트 테스트 챗봇입니다.")
@@ -191,13 +149,13 @@ with gr.Blocks() as demo:
             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],
+            respond,
+            [cohere_msg, cohere_chatbot, "Cohere Command R+", cohere_max_tokens, cohere_temperature, cohere_top_p, cohere_system_message],
             cohere_chatbot
         )
         cohere_submit_button.click(
-            cohere_respond,
-            [cohere_msg, cohere_chatbot, cohere_system_message, cohere_max_tokens, cohere_temperature, cohere_top_p],
+            respond,
+            [cohere_msg, cohere_chatbot, "Cohere Command R+", cohere_max_tokens, cohere_temperature, cohere_top_p, cohere_system_message],
             cohere_chatbot
         )
         cohere_clear_button.click(clear_conversation, outputs=cohere_chatbot, queue=False)