import gradio as gr
from transformers import pipeline

# 1. 加载模型
# 第一次启动时会下载模型，可能需要几秒钟
classifier = pipeline("text-classification", model="WJL110/emotion-classifier")

# 2. 定义标签映射
label_map = {
    "LABEL_0": "快乐",
    "LABEL_1": "愤怒",
    "LABEL_2": "悲伤"
}

def respond(message, history):
    """
    message: 用户输入的当前文本
    history: 之前的对话历史 (分类模型通常不需要上下文，所以这里我们只处理当前message)
    """
    if not message:
        return "请输入内容"
    
    try:
        # 进行推理
        result = classifier(message)[0]
        
        # 获取结果
        raw_label = result['label']
        score = result['score']
        
        # 映射标签
        emotion = label_map.get(raw_label, raw_label) # 如果找不到key，就显示原始label
        
        # 格式化回复内容
        response_text = (
            f"🤖 分析结果：\n"
            f"------------------\n"
            f"预测情感：**{emotion}**\n"
            f"置信度：{score:.2%}"
        )
        return response_text
        
    except Exception as e:
        return f"发生错误: {str(e)}"

# 3. 创建聊天界面
# 虽然这是分类任务，但使用 ChatInterface 可以给用户一种交互的感觉
demo = gr.ChatInterface(
    fn=respond,
    title="情感分析机器人 (Emotion Classifier)",
    description="输入一段文字，我会分析其中包含的情感（快乐、愤怒、悲伤）。",
    examples=["今天真是太开心了！", "这件事让我很生气。", "听到这个消息很难过。"],
    retry_btn=None,
    undo_btn=None,
    clear_btn="清除",
)

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