Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from huggingface_hub import InferenceClient
 import openai
 import anthropic
+import os
 from typing import Optional
 
 #############################
@@ -21,11 +22,12 @@ MODELS = {
 # Cohere Command R+ 모델 ID 정의
 COHERE_MODEL = "CohereForAI/c4ai-command-r-plus-08-2024"
 
-def get_client(model_name, hf_token):
+def get_client(model_name):
     """
     모델 이름에 맞춰 InferenceClient 생성.
-    hf_token을 UI에서 입력받은 값으로 사용하도록 변경.
+    토큰은 환경 변수에서 가져옴.
     """
+    hf_token = os.getenv("HF_TOKEN")
     if not hf_token:
         raise ValueError("HuggingFace API 토큰이 필요합니다.")
 
@@ -44,14 +46,13 @@ def respond_hf_qna(
     max_tokens: int,
     temperature: float,
     top_p: float,
-    system_message: str,
-    hf_token: str
+    system_message: str
 ):
     """
     HuggingFace 모델(Zephyr 등)에 대해 한 번의 질문(question)에 대한 답변을 반환하는 함수.
     """
     try:
-        client = get_client(model_name, hf_token)
+        client = get_client(model_name)
     except ValueError as e:
         return f"오류: {str(e)}"
 
@@ -80,15 +81,14 @@ def respond_cohere_qna(
     system_message: str,
     max_tokens: int,
     temperature: float,
-    top_p: float,
-    hf_token: str
+    top_p: float
 ):
     """
     Cohere Command R+ 모델을 이용해 한 번의 질문(question)에 대한 답변을 반환하는 함수.
     """
     model_name = "Cohere Command R+"
     try:
-        client = get_client(model_name, hf_token)
+        client = get_client(model_name)
     except ValueError as e:
         return f"오류: {str(e)}"
 
@@ -115,12 +115,12 @@ def respond_chatgpt_qna(
     system_message: str,
     max_tokens: int,
     temperature: float,
-    top_p: float,
-    openai_token: str
+    top_p: float
 ):
     """
     ChatGPT(OpenAI) 모델을 이용해 한 번의 질문(question)에 대한 답변을 반환하는 함수.
     """
+    openai_token = os.getenv("OPENAI_TOKEN")
     if not openai_token:
         return "OpenAI API 토큰이 필요합니다."
 
@@ -150,12 +150,12 @@ def respond_deepseek_qna(
     system_message: str,
     max_tokens: int,
     temperature: float,
-    top_p: float,
-    deepseek_token: str
+    top_p: float
 ):
     """
     DeepSeek 모델을 이용해 한 번의 질문(question)에 대한 답변을 반환하는 함수.
     """
+    deepseek_token = os.getenv("DEEPSEEK_TOKEN")
     if not deepseek_token:
         return "DeepSeek API 토큰이 필요합니다."
 
@@ -186,12 +186,12 @@ def respond_claude_qna(
     system_message: str,
     max_tokens: int,
     temperature: float,
-    top_p: float,
-    claude_api_key: str
+    top_p: float
 ) -> str:
     """
     Claude API를 사용한 개선된 응답 생성 함수
     """
+    claude_api_key = os.getenv("CLAUDE_TOKEN")
     if not claude_api_key:
         return "Claude API 토큰이 필요합니다."
 
@@ -229,30 +229,6 @@ def respond_claude_qna(
 with gr.Blocks() as demo:
     gr.Markdown("# LLM 플레이그라운드")
 
-    # 한 줄에 세 토큰 텍스트박스 배치
-    with gr.Row():
-        hf_token_box = gr.Textbox(
-            label="HuggingFace 토큰",
-            type="password",
-            placeholder="HuggingFace API 토큰을 입력하세요..."
-        )
-        openai_token_box = gr.Textbox(
-            label="OpenAI 토큰",
-            type="password",
-            placeholder="OpenAI API 토큰을 입력하세요..."
-        )
-        claude_token_box = gr.Textbox(
-            label="Claude 토큰",
-            type="password",
-            placeholder="Claude API 토큰을 입력하세요...",
-            show_copy_button=False
-        )
-        deepseek_token_box = gr.Textbox(
-            label="DeepSeek 토큰",
-            type="password",
-            placeholder="DeepSeek API 토큰을 입력하세요..."
-        )
-
     #################
     # 일반 모델 탭
     #################
@@ -290,7 +266,7 @@ with gr.Blocks() as demo:
 
         submit_button = gr.Button("전송")
 
-        def merge_and_call_hf(i1, i2, i3, i4, i5, m_name, mt, temp, top_p_, sys_msg, hf_token):
+        def merge_and_call_hf(i1, i2, i3, i4, i5, m_name, mt, temp, top_p_, sys_msg):
             # 입력1~5를 공백 기준으로 합쳐서 question 구성
             question = " ".join([i1, i2, i3, i4, i5])
             return respond_hf_qna(
@@ -299,8 +275,7 @@ with gr.Blocks() as demo:
                 max_tokens=mt,
                 temperature=temp,
                 top_p=top_p_,
-                system_message=sys_msg,
-                hf_token=hf_token
+                system_message=sys_msg
             )
 
         submit_button.click(
@@ -311,8 +286,7 @@ with gr.Blocks() as demo:
                 max_tokens,
                 temperature,
                 top_p,
-                system_message,
-                hf_token_box
+                system_message
             ],
             outputs=answer_output
         )
@@ -344,15 +318,14 @@ with gr.Blocks() as demo:
 
         cohere_submit_button = gr.Button("전송")
 
-        def merge_and_call_cohere(i1, i2, i3, i4, i5, sys_msg, mt, temp, top_p_, hf_token):
+        def merge_and_call_cohere(i1, i2, i3, i4, i5, sys_msg, mt, temp, top_p_):
             question = " ".join([i1, i2, i3, i4, i5])
             return respond_cohere_qna(
                 question=question,
                 system_message=sys_msg,
                 max_tokens=mt,
                 temperature=temp,
-                top_p=top_p_,
-                hf_token=hf_token
+                top_p=top_p_
             )
 
         cohere_submit_button.click(
@@ -362,8 +335,7 @@ with gr.Blocks() as demo:
                 cohere_system_message,
                 cohere_max_tokens,
                 cohere_temperature,
-                cohere_top_p,
-                hf_token_box
+                cohere_top_p
             ],
             outputs=cohere_answer_output
         )
@@ -395,15 +367,14 @@ with gr.Blocks() as demo:
 
         chatgpt_submit_button = gr.Button("전송")
 
-        def merge_and_call_chatgpt(i1, i2, i3, i4, i5, sys_msg, mt, temp, top_p_, openai_token):
+        def merge_and_call_chatgpt(i1, i2, i3, i4, i5, sys_msg, mt, temp, top_p_):
             question = " ".join([i1, i2, i3, i4, i5])
             return respond_chatgpt_qna(
                 question=question,
                 system_message=sys_msg,
                 max_tokens=mt,
                 temperature=temp,
-                top_p=top_p_,
-                openai_token=openai_token
+                top_p=top_p_
             )
 
         chatgpt_submit_button.click(
@@ -413,8 +384,7 @@ with gr.Blocks() as demo:
                 chatgpt_system_message,
                 chatgpt_max_tokens,
                 chatgpt_temperature,
-                chatgpt_top_p,
-                openai_token_box
+                chatgpt_top_p
             ],
             outputs=chatgpt_answer_output
         )
@@ -463,15 +433,14 @@ with gr.Blocks() as demo:
 
         claude_submit_button = gr.Button("전송")
 
-        def merge_and_call_claude(i1, i2, i3, i4, i5, sys_msg, mt, temp, top_p_, claude_key):
+        def merge_and_call_claude(i1, i2, i3, i4, i5, sys_msg, mt, temp, top_p_):
             question = " ".join([i1, i2, i3, i4, i5])
             return respond_claude_qna(
                 question=question,
                 system_message=sys_msg,
                 max_tokens=mt,
                 temperature=temp,
-                top_p=top_p_,
-                claude_api_key=claude_key
+                top_p=top_p_
             )
 
         claude_submit_button.click(
@@ -481,8 +450,7 @@ with gr.Blocks() as demo:
                 claude_system_message,
                 claude_max_tokens,
                 claude_temperature,
-                claude_top_p,
-                claude_token_box
+                claude_top_p
             ],
             outputs=claude_answer_output
         )
@@ -514,15 +482,14 @@ with gr.Blocks() as demo:
 
         deepseek_submit_button = gr.Button("전송")
 
-        def merge_and_call_deepseek(i1, i2, i3, i4, i5, sys_msg, mt, temp, top_p_, deepseek_token):
+        def merge_and_call_deepseek(i1, i2, i3, i4, i5, sys_msg, mt, temp, top_p_):
             question = " ".join([i1, i2, i3, i4, i5])
             return respond_deepseek_qna(
                 question=question,
                 system_message=sys_msg,
                 max_tokens=mt,
                 temperature=temp,
-                top_p=top_p_,
-                deepseek_token=deepseek_token
+                top_p=top_p_
             )
 
         deepseek_submit_button.click(
@@ -532,8 +499,7 @@ with gr.Blocks() as demo:
                 deepseek_system_message,
                 deepseek_max_tokens,
                 deepseek_temperature,
-                deepseek_top_p,
-                deepseek_token_box
+                deepseek_top_p
             ],
             outputs=deepseek_answer_output
         )