Update app.py

app.py

@@ -1,6 +1,10 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 import os
+import openai # OpenAI API 클라이언트 추가
+
+# OpenAI API 키 설정
+openai.api_key = os.getenv("OPENAI_API_KEY")
 
 MODELS = {
  "Zephyr 7B Beta": "HuggingFaceH4/zephyr-7b-beta",
@@ -11,16 +15,32 @@ 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",
+ "ChatGPT-4o-mini": "gpt-4o-mini" # ChatGPT-4o-mini 모델 추가
 }
 
 def get_client(model_name):
+ if model_name == "ChatGPT-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):
+ 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 respond(
  message,
  chat_history,
@@ -31,7 +51,14 @@ def respond(
  system_message,
 ):
  try:
- client = get_client(model_name)
+ if model_name == "ChatGPT-4o-mini":
+ assistant_message = call_openai_api(message, system_message, max_tokens, temperature, top_p)
+ chat_history.append((message, assistant_message))
+ yield chat_history
+ else:
+ client = get_client(model_name)
+ if client is None:
+ raise ValueError(f"No client available for model: {model_name}")
  except ValueError as e:
  chat_history.append((message, str(e)))
  return chat_history
@@ -44,7 +71,6 @@ def respond(
 
  try:
  if "Cohere" in model_name:
- # Cohere 모델을 위한 비스트리밍 처리
  response = client.chat_completion(
  messages,
  max_tokens=max_tokens,
@@ -55,7 +81,6 @@ def respond(
  chat_history.append((message, assistant_message))
  yield chat_history
  else:
- # 다른 모델들을 위한 스트리밍 처리
  stream = client.chat_completion(
  messages,
  max_tokens=max_tokens,
@@ -115,4 +140,4 @@ with gr.Blocks() as demo:
  clear_button.click(clear_conversation, outputs=chatbot, queue=False)
 
 if __name__ == "__main__":
- demo.launch()
+ demo.launch()