Update app.py

app.py

@@ -11,11 +11,14 @@ import torch
 from PIL import Image
 from transformers import AutoProcessor, AutoModelForCausalLM
 import random
+import openai  # OpenAI API 라이브러리 추가
 
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
 
+# OpenAI API 클라이언트 설정
+openai.api_key = os.getenv("OPENAI_API_KEY")
 
 # Initialize Florence model
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -177,6 +180,32 @@ class HuggingFaceInferenceNode:
         # 생략된 기능들...
         pass
 
+# gpt-4o-mini와 Cohere Command R+를 사용한 프롬프트 생성 함수
+def call_gpt4o_mini(content, system_message, max_tokens=1000, temperature=0.7, top_p=1):
+    response = openai.ChatCompletion.create(
+        model="gpt-4o-mini",
+        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 call_cohere(content, temperature=0.7, max_tokens=1000):
+    response = openai.ChatCompletion.create(
+        model="Cohere-Command-R+",
+        messages=[
+            {"role": "user", "content": content},
+        ],
+        max_tokens=max_tokens,
+        temperature=temperature,
+    )
+    return response.choices[0].message['content']
+
+
 # Gradio 인터페이스 생성 함수
 def create_interface():
     prompt_generator = PromptGenerator()  # PromptGenerator 클래스가 정의되었으므로 사용 가능
@@ -238,7 +267,7 @@ def create_interface():
             
             with gr.Column(scale=2):
                 with gr.Accordion("LLM을 사용한 프롬프트 생성", open=False):
-                    model = gr.Dropdown(["Mixtral", "Mistral", "Llama 3", "Mistral-Nemo"], label="모델", value="Llama 3")
+                    model = gr.Dropdown(["Mixtral", "Mistral", "Llama 3", "Mistral-Nemo", "gpt-4o-mini", "Cohere-Command-R+"], label="모델", value="Llama 3")
                     happy_talk = gr.Checkbox(label="행복한 대화", value=True)
                     compress = gr.Checkbox(label="압축", value=True)
                     compression_level = gr.Radio(["부드럽게", "중간", "강하게"], label="압축 레벨", value="강하게")
@@ -273,7 +302,7 @@ def create_interface():
         )
 
         generate_text_button.click(
-            huggingface_node.generate,
+            lambda model, input_text, happy_talk, compress, compression_level, poster, custom_base_prompt: call_gpt4o_mini(input_text, custom_base_prompt) if model == "gpt-4o-mini" else call_cohere(input_text),
             inputs=[model, output, happy_talk, compress, compression_level, poster, custom_base_prompt],
             outputs=text_output
         )