# ==========================================
# Hugging Face 모델 사용 - 감정 분석 Gradio
# ==========================================

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import PeftModel

# 모델 로드
print("모델 로드 중...")

BASE_MODEL = "klue/bert-base"
LORA_MODEL = "Cyplus72/nsmc-sentiment-lora" # 여러분의 Model 

tokenizer = AutoTokenizer.from_pretrained(LORA_MODEL)
base_model = AutoModelForSequenceClassification.from_pretrained(
    BASE_MODEL,
    num_labels=2
)
model = PeftModel.from_pretrained(base_model, LORA_MODEL)
model.eval()

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

print(f"완료! (Device: {device})")

# 감정 분석 함수
def analyze_sentiment(text):
    if not text.strip():
        return "텍스트를 입력해주세요", {}
    
    # 토크나이징
    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=128,
        padding=True
    ).to(device)
    
    # 예측
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.softmax(outputs.logits, dim=-1)[0]
    
    # 결과
    pred = torch.argmax(probs).item()
    label = "😊 긍정" if pred == 1 else "😞 부정"
    confidence = probs[pred].item()
    
    result = f"**{label}** (확신도: {confidence*100:.1f}%)"
    
    prob_dict = {
        "😞 부정": float(probs[0]),
        "😊 긍정": float(probs[1])
    }
    
    return result, prob_dict

# Gradio UI
demo = gr.Interface(
    fn=analyze_sentiment,
    inputs=gr.Textbox(
        label="영화 리뷰",
        placeholder="영화에 대한 리뷰를 입력하세요...",
        lines=3
    ),
    outputs=[
        gr.Markdown(label="분석 결과"),
        gr.Label(label="감정 확률", num_top_classes=2)
    ],
    title="영화 리뷰 감정 분석",
    description="LoRA로 파인튜닝된 NSMC 감정 분석 모델입니다.",
    examples=[
        ["정말 재미있는 영화였어요! 강력 추천합니다."],
        ["시간 낭비였습니다. 별로였어요."],
        ["배우들의 연기가 훌륭했습니다."],
        ["스토리가 지루하고 재미없었어요."],
    ],
    theme="soft",
    allow_flagging="never"
)

# 실행
demo.launch(share=True, debug=True)