from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# Hugging Face Hub에 올린 학습된 KoBERT 모델 경로
MODEL_NAME = "yeeeeeeeeeeeeeeeeeeo/Kobert-parrot"

# 1️⃣ 모델과 토크나이저 불러오기
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)

# 2️⃣ 디바이스 설정
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()  # 평가 모드로 변경

# 3️⃣ 감정 예측 함수
def predict_emotion(text):
    """
    입력 텍스트의 감정을 예측하여 반환합니다.
    반환값: 감정 클래스 인덱스 또는 확률 벡터
    """
    # 토큰화 및 텐서 변환
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)

    # 모델 추론
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits

    # 확률 및 클래스 예측
    probs = torch.softmax(logits, dim=-1)
    pred_class = torch.argmax(probs, dim=-1).item()

    return pred_class, probs.cpu().numpy()