import gradio as gr
import random
import os
from huggingface_hub import InferenceClient

MODELS = {
    "Zephyr 7B Beta": "HuggingFaceH4/zephyr-7b-beta",
    "DeepSeek Coder V2": "deepseek-ai/DeepSeek-Coder-V2-Instruct",
    "Meta Llama 3.1 8B": "meta-llama/Meta-Llama-3.1-8B-Instruct",
    "Meta-Llama 3.1 70B-Instruct": "meta-llama/Meta-Llama-3.1-70B-Instruct",
    "Microsoft": "microsoft/Phi-3-mini-4k-instruct",
    "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",
    "Aya-23-35B": "CohereForAI/aya-23-35B"
}

def create_client(model_name):
    return InferenceClient(model_name, token=os.getenv("HF_TOKEN"))

def call_api(model, content, system_message, max_tokens, temperature, top_p):
    client = create_client(MODELS[model])
    messages = [{"role": "system", "content": system_message}, {"role": "user", "content": content}]
    random_seed = random.randint(0, 1000000)
    response = client.chat_completion(messages=messages, max_tokens=max_tokens, temperature=temperature, top_p=top_p, seed=random_seed)
    return response.choices[0].message.content

def generate_copywriting(model, topic, system_message, max_tokens, temperature, top_p):
    # 반드시 한글로 생성하도록 프롬프트를 명시
    content = f"'{topic}'에 대한 카피라이팅을 작성하되, 반드시 한국어로 작성하세요. {system_message}를 참고하여 10개의 카피를 작성하세요."
    return call_api(model, content, system_message, max_tokens, temperature, top_p)

title = "SNS 마케팅 카피라이팅 생성기 (한국어 전용)"

with gr.Blocks() as demo:
    gr.Markdown(f"# {title}")
    
    model = gr.Radio(choices=list(MODELS.keys()), label="언어 모델 선택", value="Zephyr 7B Beta")
    topic = gr.Textbox(label="카피라이팅 주제 입력", lines=1, placeholder="예: 여름 세일, 신제품 출시")
    system_message = gr.Textbox(label="카피라이팅 작성 규칙", lines=10, value="""
      1. 관심을 끌 수 있는 강력한 문구로 작성하라.
      2. 필요한 경우 간단한 비유나 은유를 사용하여 문구를 보강하라.
      3. 명확하고 설득력 있는 CTA(Call to Action)를 작성하라.
      4. 긴급성이나 혜택을 강조하는 문구를 사용하라.
    """)
    
    with gr.Accordion("고급 설정", open=False):
        max_tokens = gr.Slider(label="Max Tokens", minimum=0, maximum=4000, value=500, step=100)
        temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, value=0.75, step=0.05)
        top_p = gr.Slider(label="Top P", minimum=0.1, maximum=1.0, value=0.95, step=0.05)
    
    generate_btn = gr.Button("카피라이팅 생성하기")
    output = gr.Textbox(label="생성된 카피라이팅", lines=10)
    
    generate_btn.click(fn=generate_copywriting, 
                       inputs=[model, topic, system_message, max_tokens, temperature, top_p], 
                       outputs=[output])

demo.launch()