Update app.py

app.py

@@ -1,6 +1,10 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 import os
+import openai  # OpenAI API를 사용하기 위해 추가
+
+# OpenAI API 클라이언트 설정
+openai.api_key = os.getenv("OPENAI_API_KEY")
 
 MODELS = {
     "Zephyr 7B Beta": "HuggingFaceH4/zephyr-7b-beta",
@@ -11,16 +15,32 @@ MODELS = {
     "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",
-    "Cohere Aya-23-35B": "CohereForAI/aya-23-35B"
+    "Cohere Aya-23-35B": "CohereForAI/aya-23-35B",
+    "GPT-4o Mini": "gpt-4o-mini"  # GPT-4o Mini 모델 추가
 }
 
 def get_client(model_name):
+    if model_name == "GPT-4o Mini":
+        return None  # OpenAI 모델을 위해 HuggingFace 클라이언트를 사용하지 않음
     model_id = MODELS[model_name]
     hf_token = os.getenv("HF_TOKEN")
     if not hf_token:
         raise ValueError("HF_TOKEN environment variable is required")
     return InferenceClient(model_id, token=hf_token)
 
+def call_openai_api(content, system_message, max_tokens, temperature, top_p):
+    response = openai.ChatCompletion.create(
+        model="gpt-4o-mini",  # OpenAI 모델 사용
+        messages=[
+            {"role": "system", "content": system_message},
+            {"role": "user", "content": content},
+        ],
+        max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+    )
+    return response.choices[0].message['content']
+
 def respond(
     message,
     chat_history,
@@ -30,52 +50,64 @@ def respond(
     top_p,
     system_message,
 ):
-    try:
-        client = get_client(model_name)
-    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:
-        messages.append({"role": "user", "content": human})
-        messages.append({"role": "assistant", "content": assistant})
-    messages.append({"role": "user", "content": message})
-
-    try:
-        if "Cohere" in model_name:
-            # Cohere 모델을 위한 비스트리밍 처리
-            response = client.chat_completion(
-                messages,
-                max_tokens=max_tokens,
-                temperature=temperature,
-                top_p=top_p,
+    if model_name == "GPT-4o Mini":
+        try:
+            assistant_message = call_openai_api(
+                message, system_message, max_tokens, temperature, top_p
             )
-            assistant_message = response.choices[0].message.content
             chat_history.append((message, assistant_message))
             yield 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
-    except Exception as e:
-        error_message = f"An error occurred: {str(e)}"
-        chat_history.append((message, error_message))
-        yield chat_history
+        except Exception as e:
+            error_message = f"An error occurred with GPT-4o Mini: {str(e)}"
+            chat_history.append((message, error_message))
+            yield chat_history
+    else:
+        try:
+            client = get_client(model_name)
+        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:
+            messages.append({"role": "user", "content": human})
+            messages.append({"role": "assistant", "content": assistant})
+        messages.append({"role": "user", "content": message})
+
+        try:
+            if "Cohere" in model_name:
+                # Cohere 모델을 위한 비스트리밍 처리
+                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))
+                yield 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
+        except Exception as e:
+            error_message = f"An error occurred: {str(e)}"
+            chat_history.append((message, error_message))
+            yield chat_history
 
 def clear_conversation():
     return []