app.py

import gradio as gr
import os
from huggingface_hub import InferenceClient

#############################
# [블로그 생성기 - Cohere Command R+]
#############################

# Cohere Command R+ 모델 ID 정의
COHERE_MODEL = "CohereForAI/c4ai-command-r-plus-08-2024"
SYSTEM_MESSAGE = "블로그 글을 작성해주세요."
MAX_TOKENS = 4000
TEMPERATURE = 0.7
TOP_P = 0.95

def get_client(hf_token):
    """
    HuggingFace InferenceClient 생성.
    """
    if not hf_token:
        raise ValueError("HuggingFace API 토큰이 필요합니다.")

    return InferenceClient(COHERE_MODEL, token=hf_token)

def respond_cohere_qna(question, system_message, max_tokens, temperature, top_p, hf_token):
    """
    Cohere 모델을 사용하여 질문에 대한 응답 생성.
    """
    client = get_client(hf_token)

    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": question}
    ]

    response = client.chat_completion(
        messages=messages,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
    )
    return response.choices[0].message.content

# 블로그 생성 함수 정의
def generate_blog(tone, ref1, ref2, ref3):
    """
    블로그 글을 생성하는 함수.
    HuggingFace 토큰은 환경 변수에서 직접 가져옵니다.
    """
    # HuggingFace 토큰 가져오기
    hf_token_value = os.getenv("HF_TOKEN")
    if not hf_token_value:
        return "HuggingFace 토큰이 설정되지 않았습니다."

    # 프롬프트 구성
    prompt = f"말투: {tone}\n참조문 1: {ref1}\n참조문 2: {ref2}\n참조문 3: {ref3}"

    # 응답 생성
    return respond_cohere_qna(
        question=prompt,
        system_message=SYSTEM_MESSAGE,
        max_tokens=MAX_TOKENS,
        temperature=TEMPERATURE,
        top_p=TOP_P,
        hf_token=hf_token_value
    )

# Gradio UI 구성
with gr.Blocks() as demo:
    gr.Markdown("# 블로그 생성기")

    # HuggingFace 토큰 (secrets 사용)
    hf_token = os.getenv("HF_TOKEN")

    # 입력 필드
    with gr.Row():
        tone = gr.Radio(
            choices=["친근하게", "일반적인", "전문적인"],
            label="말투 바꾸기",
            value="일반적인"
        )

        ref1 = gr.Textbox(label="참조문 1", lines=3)
        ref2 = gr.Textbox(label="참조문 2", lines=3)
        ref3 = gr.Textbox(label="참조문 3", lines=3)

    # 결과 출력
    answer_output = gr.Textbox(label="생성된 블로그 글", lines=10, interactive=False)

    # 버튼 배치
    submit_button = gr.Button("생성")

    submit_button.click(
        fn=generate_blog,
        inputs=[tone, ref1, ref2, ref3],
        outputs=answer_output
    )

# 앱 실행
if __name__ == "__main__":
    demo.launch()