Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
+import openai  # OpenAI API를 사용하기 위해 추가
 import os
 import random
 import logging
@@ -18,7 +19,8 @@ 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",
-    "Aya-23-35B": "CohereForAI/aya-23-35B"
+    "Aya-23-35B": "CohereForAI/aya-23-35B",
+    "GPT-4o Mini": "gpt-4o-mini"  # 새로운 모델 추가
 }
 
 # HuggingFace 토큰 설정
@@ -26,7 +28,13 @@ hf_token = os.getenv("HF_TOKEN")
 if not hf_token:
     raise ValueError("HF_TOKEN 환경 변수가 설정되지 않았습니다.")
 
+# OpenAI API 키 설정
+openai.api_key = os.getenv("OPENAI_API_KEY")
+
 def call_hf_api(prompt, reference_text, max_tokens, temperature, top_p, model):
+    if model == "gpt-4o-mini":
+        return call_openai_api(prompt, reference_text, max_tokens, temperature, top_p)
+    
     client = InferenceClient(model=model, token=hf_token)
     combined_prompt = f"{prompt}\n\n참고 텍스트:\n{reference_text}"
     random_seed = random.randint(0, 1000000)
@@ -44,8 +52,26 @@ def call_hf_api(prompt, reference_text, max_tokens, temperature, top_p, model):
         logging.error(f"HuggingFace API 호출 중 오류 발생: {str(e)}")
         return f"응답 생성 중 오류 발생: {str(e)}. 나중에 다시 시도해 주세요."
 
+def call_openai_api(content, system_message, max_tokens, temperature, top_p):
+    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 generate_response(prompt, reference_text, max_tokens, temperature, top_p, model):
-    response = call_hf_api(prompt, reference_text, max_tokens, temperature, top_p, MODELS[model])
+    if model == "GPT-4o Mini":
+        system_message = "이것은 사용자 요청에 대한 참고 텍스트를 활용하여 응답을 생성하는 작업입니다."
+        response = call_openai_api(prompt, reference_text, max_tokens, temperature, top_p)
+    else:
+        response = call_hf_api(prompt, reference_text, max_tokens, temperature, top_p, MODELS[model])
+    
     response_html = f"""
     <h3>생성된 응답:</h3>
     <div style='max-height: 500px; overflow-y: auto; white-space: pre-wrap; word-wrap: break-word;'>
@@ -79,4 +105,4 @@ with gr.Blocks() as demo:
     )
 
 # 인터페이스 실행
-demo.launch(share=True)
+demo.launch(share=True)