from transformers import pipeline
import gradio as gr

# ✅ KoAlpaca-Polyglot 모델은 한국어 대화에 특화된 모델입니다.
generator = pipeline("text-generation", model="beomi/KoAlpaca-Polyglot-12.8B")

# ✅ 감정 톤별 대화 생성 함수
def chat_with_emotion(user_input, tone="자연스러운"):
    if not user_input:
        return "입력을 입력해주세요!"
    prompt = f"({tone} 톤으로 응답해줘)\n사용자: {user_input}\nAI:"
    response = generator(prompt, max_new_tokens=150, do_sample=True, temperature=0.7)
    return response[0]["generated_text"].replace(prompt, "").strip()

# ✅ Gradio 인터페이스 구성
demo = gr.Interface(
    fn=chat_with_emotion,
    inputs=[
        gr.Textbox(label="사용자 입력", placeholder="오늘 하루가 힘들었어..."),
        gr.Radio(["자연스러운", "위로하는", "공감하는", "안심시키는", "활기찬"], label="감정 톤", value="자연스러운"),
    ],
    outputs=gr.Textbox(label="AI 응답"),
    title="감정 기반 한국어 챗봇 (KoAlpaca-Polyglot)",
    description="한국어로 감정 톤에 따라 공감형 답변을 생성합니다."
)

demo.launch()