Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
+import openai
 import os
 
 MODELS = {
@@ -11,16 +12,33 @@ 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": "openai/gpt-4o-mini"  # 추가된 모델
 }
 
 def get_client(model_name):
+    if model_name == "GPT-4o Mini":
+        return None  # OpenAI API는 별도로 처리
     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):
+    openai.api_key = os.getenv("OPENAI_API_KEY")
+    response = openai.ChatCompletion.create(
+        model="gpt-4o-mini",  # 또는 다른 모델 ID 사용
+        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,
@@ -36,46 +54,58 @@ def respond(
         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
+    if model_name == "GPT-4o Mini":
+        # GPT-4o Mini 모델에 대한 비스트리밍 처리
+        try:
+            assistant_message = call_openai_api(message, system_message, max_tokens, temperature, top_p)
             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: {str(e)}"
+            chat_history.append((message, error_message))
+            yield chat_history
+    else:
+        # 다른 모델들에 대한 처리
+        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 []
@@ -115,4 +145,4 @@ with gr.Blocks() as demo:
     clear_button.click(clear_conversation, outputs=chatbot, queue=False)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()