app.py

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

# 모델 설정
MODELS = {
    'ei': ('epinfomax/mbti-ei-classifier', ['E', 'I']),
    'sn': ('epinfomax/mbti-sn-classifier', ['S', 'N']),
    'tf': ('epinfomax/mbti-tf-classifier', ['T', 'F']),
    'jp': ('epinfomax/mbti-jp-classifier', ['J', 'P']),
}

DESCRIPTIONS = {
    'E': '외향 (Extraversion)',
    'I': '내향 (Introversion)',
    'S': '감각 (Sensing)',
    'N': '직관 (iNtuition)',
    'T': '사고 (Thinking)',
    'F': '감정 (Feeling)',
    'J': '판단 (Judging)',
    'P': '인식 (Perceiving)',
}

# 모델 로드
print("모델 로딩 중...")
tokenizer = None
models = {}

for key, (model_name, _) in MODELS.items():
    if tokenizer is None:
        tokenizer = AutoTokenizer.from_pretrained(model_name)
    models[key] = AutoModelForSequenceClassification.from_pretrained(model_name)
    models[key].eval()

print("모델 로딩 완료!")


def predict_mbti(text: str):
    """텍스트로부터 MBTI 예측"""
    if not text or len(text.strip()) < 10:
        return "텍스트를 10자 이상 입력해주세요.", "", "", "", ""

    inputs = tokenizer(
        text,
        return_tensors='pt',
        truncation=True,
        max_length=256,
        padding=True
    )

    results = {}
    mbti = ""

    for key, model in models.items():
        with torch.no_grad():
            outputs = model(**inputs)
            probs = torch.softmax(outputs.logits, dim=1)
            pred = torch.argmax(probs, dim=1).item()
            confidence = probs[0][pred].item()

        labels = MODELS[key][1]
        results[key] = {
            'prediction': labels[pred],
            'confidence': confidence,
            'prob_0': probs[0][0].item(),
            'prob_1': probs[0][1].item(),
        }
        mbti += labels[pred]

    # 결과 포맷팅
    mbti_result = f"## 🎯 예측 결과: **{mbti}**"

    def format_axis(key, label0, label1):
        r = results[key]
        p0, p1 = r['prob_0'] * 100, r['prob_1'] * 100
        pred = r['prediction']

        bar0 = "█" * int(p0 / 5) + "░" * (20 - int(p0 / 5))
        bar1 = "█" * int(p1 / 5) + "░" * (20 - int(p1 / 5))

        mark0 = " ✓" if pred == label0 else ""
        mark1 = " ✓" if pred == label1 else ""

        return f"""### {label0}/{label1} 축
| 유형 | 확률 |
|------|------|
| **{label0}** {DESCRIPTIONS[label0]} | `{bar0}` {p0:.1f}%{mark0} |
| **{label1}** {DESCRIPTIONS[label1]} | `{bar1}` {p1:.1f}%{mark1} |
"""

    ei_result = format_axis('ei', 'E', 'I')
    sn_result = format_axis('sn', 'S', 'N')
    tf_result = format_axis('tf', 'T', 'F')
    jp_result = format_axis('jp', 'J', 'P')

    return mbti_result, ei_result, sn_result, tf_result, jp_result


# Gradio 인터페이스
with gr.Blocks(title="MBTI 예측기", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
# 🧠 MBTI 예측기
### 채팅/대화 스타일로 MBTI를 예측합니다

평소 쓰는 말투나 채팅 내용을 입력하면 AI가 MBTI를 분석합니다.
*※ 재미용 모델입니다. 실제 MBTI와 다를 수 있어요!*

---
""")

    with gr.Row():
        with gr.Column(scale=1):
            text_input = gr.Textbox(
                label="💬 텍스트 입력",
                placeholder="평소 말투나 채팅 내용을 입력하세요...\n\n예시:\n- 카톡 대화 내용\n- 일상적으로 쓰는 말투\n- SNS 게시글 등",
                lines=8,
            )

            with gr.Row():
                submit_btn = gr.Button("🔍 분석하기", variant="primary", size="lg")
                clear_btn = gr.Button("🗑️ 지우기", size="lg")

            gr.Examples(
                examples=[
                    ["오늘 친구들이랑 카페 갔는데 완전 분위기 좋았어!! 새로 나온 음료도 맛있고 ㅎㅎ 다음에 또 가야지~"],
                    ["음... 혼자 있는 시간이 좋아요. 책 읽거나 음악 듣는 게 제 취미예요. 사람 많은 곳은 좀 피곤하더라고요."],
                    ["이번 프로젝트 일정 정리했어. 1단계는 다음주까지, 2단계는 그 다음주. 효율적으로 진행하자."],
                    ["와 이거 진짜 감동적이다 ㅠㅠ 눈물 날 것 같아... 너무 좋은 이야기야"],
                ],
                inputs=text_input,
            )

    with gr.Row():
        mbti_output = gr.Markdown(label="결과")

    with gr.Row():
        with gr.Column():
            ei_output = gr.Markdown()
        with gr.Column():
            sn_output = gr.Markdown()

    with gr.Row():
        with gr.Column():
            tf_output = gr.Markdown()
        with gr.Column():
            jp_output = gr.Markdown()

    gr.Markdown("""
---
### ℹ️ 모델 정보
- **베이스 모델**: klue/roberta-large (한국어 특화)
- **학습 데이터**: 한국어 MBTI 대화 데이터셋
- **정확도**: E/I 64%, S/N 68%, T/F 52%, J/P 61%

Made with ❤️ by [epinfomax](https://huggingface.co/epinfomax)
""")

    submit_btn.click(
        fn=predict_mbti,
        inputs=text_input,
        outputs=[mbti_output, ei_output, sn_output, tf_output, jp_output]
    )

    clear_btn.click(
        fn=lambda: ("", "", "", "", "", ""),
        outputs=[text_input, mbti_output, ei_output, sn_output, tf_output, jp_output]
    )


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