import gradio as gr
from huggingface_hub import InferenceClient
import os

# Cohere Command R+ 모델 ID 정의
COHERE_MODEL = "CohereForAI/c4ai-command-r-plus-08-2024"

# Hugging Face API 토큰을 환경 변수에서 가져옴
def get_client():
    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+ 모델을 사용하여 블로그 글을 생성하는 함수
    """
    system_message = f"""반드시 한글로 답변할 것.
    너는 최고의 비서이며 요청에 따라 주어진 말투를 사용하여 블로그를 작성한다.
    말투: {tone}.
    """
    question = f"참조글 1: {ref1}\n참조글 2: {ref2}\n참조글 3: {ref3}\n블로그 글을 생성하라."

    try:
        client = get_client()
        response = client.chat_completion(
            messages=[
                {"role": "system", "content": system_message},
                {"role": "user", "content": question}
            ],
            max_tokens=4000,
            temperature=0.7,
            top_p=0.95
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"오류가 발생했습니다: {str(e)}"

# Gradio 인터페이스 구축
with gr.Blocks() as demo:
    gr.Markdown("# 블로그 생성기")

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

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

    output = gr.Textbox(label="생성된 블로그 글", lines=10, interactive=False)

    generate_button = gr.Button("생성하기")

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

if __name__ == "__main__":
    demo.launch()

# requirements.txt
gradio
huggingface-hub==4.44.1