import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForCausalLM
import torch
import torch.nn.functional as F

# 감정 분석용 모델
emotion_model = AutoModelForSequenceClassification.from_pretrained("beomi/KcELECTRA-base", num_labels=3)
emotion_tokenizer = AutoTokenizer.from_pretrained("beomi/KcELECTRA-base")
emotion_labels = ['부정', '중립', '긍정']

# 텍스트 생성용 GPT 모델
gpt_model = AutoModelForCausalLM.from_pretrained("skt/kogpt2-base-v2")
gpt_tokenizer = AutoTokenizer.from_pretrained("skt/kogpt2-base-v2")

# 감정 분석 함수
def predict_emotion(text):
    inputs = emotion_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = emotion_model(**inputs)
        probs = F.softmax(outputs.logits, dim=1)
        pred = torch.argmax(probs, dim=1).item()
    return emotion_labels[pred]

# GPT 이어쓰기 함수
def emotional_gpt(user_input):
    emotion = predict_emotion(user_input)

    if emotion == "긍정":
        prompt = "기분 좋은 하루였다. "
    elif emotion == "부정":
        prompt = "우울한 기분으로 시작된 하루, "
    else:
        prompt = "평범한 하루가 시작되었다. "
    
    prompt += user_input

    input_ids = gpt_tokenizer.encode(prompt, return_tensors="pt")
    output = gpt_model.generate(input_ids, max_length=150, do_sample=True, temperature=0.8, top_k=50)
    result = gpt_tokenizer.decode(output[0], skip_special_tokens=True)

    return f"🧠 감정 분석 결과: {emotion}\n\n✍️ GPT가 이어 쓴 글:\n{result}"

# Gradio 인터페이스 구성
gr.Interface(
    fn=emotional_gpt,
    inputs=gr.Textbox(lines=3, label="✍️ 감정을 담은 문장을 입력해주세요!", placeholder="예: 오늘 너무 외로웠어"),
    outputs="text",
    title="🎭 감정형 GPT 한글 작문 AI",
    description="🧠 감정을 먼저 파악하고 ✨ 그 감정에 어울리는 문장을 이어서 작성해줍니다!",
    theme="soft",
    examples=[
        ["기분이 너무 좋았어"],
        ["진짜 외롭고 힘든 하루였어"],
        ["회의가 그냥 그랬어"]
    ]
).launch()