import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
# === 强制指定 BERT 配置 ===
from transformers import BertConfig, BertTokenizer, BertForSequenceClassification
MODEL_REPO = "robot4/emotion"
MODEL_SUBFOLDER = "checkpoints_finally_end"
print(f"正在加载 BERT 模型: {MODEL_REPO} ...")
try:
    # 1. 显式加载 BERT Tokenizer
    tokenizer = BertTokenizer.from_pretrained(MODEL_REPO, subfolder=MODEL_SUBFOLDER)
    
    # 2. 显式加载 BERT Config 并修正 model_type
    config = BertConfig.from_pretrained(MODEL_REPO, subfolder=MODEL_SUBFOLDER)
    
    # 3. 加载模型
    model = BertForSequenceClassification.from_pretrained(
        MODEL_REPO, 
        config=config, 
        subfolder=MODEL_SUBFOLDER
    )
    
except Exception as e:
    print(f"Error: {e}")
    raise e

# 自动判断是否有 GPU (Space 上通常是 CPU，除非您付费)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

def predict(text):
    if not text: return None, "请输入内容"
    
    # 1. 预处理
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128, padding=True)
    inputs = {k: v.to(device) for k, v in inputs.items()}
    
    # 2. 推理
    with torch.no_grad():
        logits = model(**inputs).logits
        probs = torch.nn.functional.softmax(logits, dim=-1)
        
    # 3. 解析结果
    pred_idx = torch.argmax(probs).item()
    confidence = probs[0][pred_idx].item()
    
    # 标签映射
    id2label = {0: '😡 消极 (Negative)', 1: '😐 中性 (Neutral)', 2: '😊 积极 (Positive)'}
    label = id2label.get(pred_idx, "Unknown")
    
    # 返回给界面的数据
    return {
        '积极': probs[0][2].item(),
        '中性': probs[0][1].item(),
        '消极': probs[0][0].item()
    }, f"预测结果: {label}\n置信度: {confidence:.4f}"

# === 构建界面 ===
with gr.Blocks(title="中文情感分析") as demo:
    gr.Markdown(f"# 🎭 中文情感分析演示 (BERT)")
    gr.Markdown(f"模型加载自: [Hugging Face Hub]({MODEL_REPO})")
    
    with gr.Row():
        with gr.Column():
            inp = gr.Textbox(label="输入中文评论", lines=4, placeholder="比如：这家店真的太好吃了，强烈推荐！")
            btn = gr.Button("开始分析", variant="primary")
            
        with gr.Column():
            out_label = gr.Label(label="情感概率")
            out_text = gr.Textbox(label="详细结果")
    
    btn.click(predict, inputs=inp, outputs=[out_label, out_text])
    
    gr.Examples(
        examples=["这家店太难吃了，避雷！", "还可以，中规中矩。", "超级好评，下次还来！", "物流稍微有点慢，但东西不错。"], 
        inputs=inp
    )

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