Update app.py

app.py

@@ -1,3 +1,4 @@
+# app.py
 import gradio as gr
 from huggingface_hub import InferenceClient
 import os
@@ -5,63 +6,91 @@ import os
 # Cohere Command R+ 모델 ID 정의
 COHERE_MODEL = "CohereForAI/c4ai-command-r-plus-08-2024"
 
-# Hugging Face API 토큰을 환경 변수에서 가져옴
-def get_client():
+def get_client(model_name):
+    """
+    모델 이름에 맞춰 InferenceClient 생성.
+    토큰은 환경 변수에서 가져옴.
+    """
     hf_token = os.getenv("HF_TOKEN")
     if not hf_token:
         raise ValueError("HuggingFace API 토큰이 필요합니다.")
-    return InferenceClient(COHERE_MODEL, token=hf_token)
 
-def generate_blog(tone: str, ref1: str, ref2: str, ref3: str):
-    """
-    Cohere Command R+ 모델을 사용하여 블로그 글을 생성하는 함수
+    if model_name == "Cohere Command R+":
+        model_id = COHERE_MODEL
+    else:
+        raise ValueError("유효하지 않은 모델 이름입니다.")
+    return InferenceClient(model_id, token=hf_token)
+
+def respond_cohere_qna(
+    tone: str,
+    reference1: str,
+    reference2: str,
+    reference3: str,
+    system_message: str,
+    max_tokens: int,
+    temperature: float,
+    top_p: float
+):
     """
-    system_message = f"""반드시 한글로 답변할 것.
-    너는 최고의 비서이며 요청에 따라 주어진 말투를 사용하여 블로그를 작성한다.
-    말투: {tone}.
+    Cohere Command R+ 모델을 이용해 블로그 생성 함수.
     """
-    question = f"참조글 1: {ref1}\n참조글 2: {ref2}\n참조글 3: {ref3}\n블로그 글을 생성하라."
+    model_name = "Cohere Command R+"
+    try:
+        client = get_client(model_name)
+    except ValueError as e:
+        return f"오류: {str(e)}"
+
+    question = f"말투: {tone} \n\n 참조글1: {reference1} \n\n 참조글2: {reference2} \n\n 참조글3: {reference3}"
 
     try:
-        client = get_client()
-        response = client.chat_completion(
-            messages=[
+        response_full = client.chat_completion(
+            [
                 {"role": "system", "content": system_message},
                 {"role": "user", "content": question}
             ],
-            max_tokens=4000,
-            temperature=0.7,
-            top_p=0.95
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
         )
-        return response.choices[0].message.content
+        assistant_message = response_full.choices[0].message.content
+        return assistant_message
     except Exception as e:
         return f"오류가 발생했습니다: {str(e)}"
 
-# Gradio 인터페이스 구축
+# Gradio UI 설정
 with gr.Blocks() as demo:
     gr.Markdown("# 블로그 생성기")
 
     with gr.Row():
         tone = gr.Radio(
-            label="말투 바꾸기",
             choices=["친근하게", "일반적인", "전문적인"],
-            value="일반적인"
+            value="일반적인",
+            label="말투바꾸기"
         )
 
-    ref1 = gr.Textbox(label="참조글 1", lines=3)
-    ref2 = gr.Textbox(label="참조글 2", lines=3)
-    ref3 = gr.Textbox(label="참조글 3", lines=3)
+    reference1 = gr.Textbox(label="참조글 1", lines=3, placeholder="블로그 글에 포함할 주요 참조 내용 1")
+    reference2 = gr.Textbox(label="참조글 2", lines=3, placeholder="블로그 글에 포함할 주요 참조 내용 2")
+    reference3 = gr.Textbox(label="참조글 3", lines=3, placeholder="블로그 글에 포함할 주요 참조 내용 3")
 
     output = gr.Textbox(label="생성된 블로그 글", lines=10, interactive=False)
 
-    generate_button = gr.Button("생성하기")
+    with gr.Accordion("고급 설정 (Cohere)", open=False):
+        system_message = gr.Textbox(
+            value="""반드시 한글로 답변할 것.\n너는 블로그 작성을 도와주는 비서이다.\n사용자의 요구사항을 정확히 반영하여 작성하라.""",
+            label="System Message",
+            lines=3
+        )
+        max_tokens = gr.Slider(minimum=100, maximum=5000, value=2000, step=100, label="Max Tokens")
+        temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature")
+        top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-P")
+
+    generate_button = gr.Button("생성")
 
     generate_button.click(
-        fn=generate_blog,
-        inputs=[tone, ref1, ref2, ref3],
+        fn=respond_cohere_qna,
+        inputs=[tone, reference1, reference2, reference3, system_message, max_tokens, temperature, top_p],
         outputs=output
     )
 
 if __name__ == "__main__":
     demo.launch()
-