Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
+import openai  # OpenAI 라이브러리 추가
 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": "openai/gpt-4o-mini"  # GPT-4o Mini 모델 추가
 }
 
 # HuggingFace 토큰 설정
@@ -26,22 +28,50 @@ hf_token = os.getenv("HF_TOKEN")
 if not hf_token:
     raise ValueError("HF_TOKEN 환경 변수가 설정되지 않았습니다.")
 
+# OpenAI API 키 설정
+openai.api_key = os.getenv("OPENAI_API_KEY")
+if not openai.api_key:
+    raise ValueError("OPENAI_API_KEY 환경 변수가 설정되지 않았습니다.")
+
 def call_hf_api(prompt, reference_text, max_tokens, temperature, top_p, model):
-    client = InferenceClient(model=model, token=hf_token)
+    if model == "openai/gpt-4o-mini":
+        return call_openai_api(prompt, reference_text, max_tokens, temperature, top_p)
+    else:
+        client = InferenceClient(model=model, token=hf_token)
+        combined_prompt = f"{prompt}\n\n참고 텍스트:\n{reference_text}"
+        random_seed = random.randint(0, 1000000)
+        
+        try:
+            response = client.text_generation(
+                combined_prompt,
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                seed=random_seed
+            )
+            return response
+        except Exception as e:
+            logging.error(f"HuggingFace API 호출 중 오류 발생: {str(e)}")
+            return f"응답 생성 중 오류 발생: {str(e)}. 나중에 다시 시도해 주세요."
+
+def call_openai_api(prompt, reference_text, max_tokens, temperature, top_p):
+    system_message = "You are a helpful assistant."
     combined_prompt = f"{prompt}\n\n참고 텍스트:\n{reference_text}"
-    random_seed = random.randint(0, 1000000)
     
     try:
-        response = client.text_generation(
-            combined_prompt,
-            max_new_tokens=max_tokens,
+        response = openai.ChatCompletion.create(
+            model="gpt-4o-mini",
+            messages=[
+                {"role": "system", "content": system_message},
+                {"role": "user", "content": combined_prompt},
+            ],
+            max_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
-            seed=random_seed
         )
-        return response
+        return response.choices[0].message['content']
     except Exception as e:
-        logging.error(f"HuggingFace API 호출 중 오류 발생: {str(e)}")
+        logging.error(f"OpenAI API 호출 중 오류 발생: {str(e)}")
         return f"응답 생성 중 오류 발생: {str(e)}. 나중에 다시 시도해 주세요."
 
 def generate_response(prompt, reference_text, max_tokens, temperature, top_p, model):
@@ -79,4 +109,4 @@ with gr.Blocks() as demo:
     )
 
 # 인터페이스 실행
-demo.launch(share=True)
+demo.launch(share=True)