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
  )