import torch
import numpy as np
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# 모델 로딩
model = AutoModelForSequenceClassification.from_pretrained("M1NJ1/klue-bert-emotion")
tokenizer = AutoTokenizer.from_pretrained("M1NJ1/klue-bert-emotion")
model.eval().to("cuda" if torch.cuda.is_available() else "cpu")

label_map = list(range(44))  # 실제 감정 라벨 이름으로 수정 가능

def predict_multi(text):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=64).to(device)
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.sigmoid(outputs.logits).squeeze().cpu().numpy()
    emotions = []
    for i, p in enumerate(probs):
        if p > 0.5:
            emotions.append(f"감정{i}")
    return emotions

def generate_response(emotions, mode):
    if not emotions:
        return "감정을 파악하지 못했어요. 다시 말해줄 수 있나요?"
    response = ""
    if mode == 'T':
        if '감정1' in emotions:
            response += "불안은 준비 부족에서 올 수 있어요. 차근히 정리해볼까요?
"
    elif mode == 'F':
        if '감정1' in emotions:
            response += "불안한 하루였겠군요... 괜찮아요, 당신 잘하고 있어요.
"
    return response or "당신의 감정에 귀 기울이고 있어요."