app.py

import torch
import re
import json
import gradio as gr
from konlpy.tag import Okt
from transformers import AutoTokenizer, BertForSequenceClassification

# --- 1. 설정 및 전처리 함수 ---
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

BASE_TOKENIZER_DIR = 'base'
EMOTION_MODEL_DIR = 'kobert_emotion_classifier'
GENRE_MODEL_DIR = 'kobert_genre_classifier_archive'

okt = Okt()
def remove_english(text):
    return re.sub(r'[A-Za-z]+', '', text)

def extract_pos(text):
    allowed_pos = ['Noun', 'Verb', 'Adjective']
    text = remove_english(text)
    return ' '.join([word for word, pos in okt.pos(text) if pos in allowed_pos])

# --- 2. 두 모델 및 토크나이저 로드 ---
try:
    tokenizer = AutoTokenizer.from_pretrained(BASE_TOKENIZER_DIR, trust_remote_code=True)
    print("✅ 공용 토크나이저 로드 성공")

    # 감정 분류 모델 로드 (파일에서 레이블 읽기)
    emotion_model = BertForSequenceClassification.from_pretrained(EMOTION_MODEL_DIR)
    emotion_model.to(device)
    emotion_model.eval()
    with open(f"{EMOTION_MODEL_DIR}/labels_ids.json", "r", encoding="utf-8") as f:
        emotion_labels_ids = json.load(f)
    id_to_emotion_label = {v: k for k, v in emotion_labels_ids.items()}
    print("✅ 감정 분류 모델 로드 성공")

    # 장르 분류 모델 로드
    genre_model = BertForSequenceClassification.from_pretrained(GENRE_MODEL_DIR)
    genre_model.to(device)
    genre_model.eval()
    
  
    id_to_genre_label = {
        0: '록/메탈',
        1: '댄스',
        2: 'R&B/Soul',
        3: '발라드',
        4: '랩/힙합', 
        5: '트로트'
    }
    print("✅ 장르 분류 모델 로드 성공 (레이블 직접 정의)")

except Exception as e:
    print(f"모델 또는 토크나이저 로딩 중 오류 발생: {e}")
    emotion_model, genre_model = None, None

# --- 3. 통합 예측 함수 (이하 동일) ---
def predict_emotion_and_genre(text):
    if not emotion_model or not genre_model:
        raise gr.Error("모델이 로드되지 않았습니다. Space의 로그를 확인해주세요.")

    preprocessed_text = extract_pos(text)

    # 감정 예측
    emotion_inputs = tokenizer(preprocessed_text, return_tensors='pt', truncation=True, padding=True, max_length=384).to(device)
    with torch.no_grad():
        emotion_logits = emotion_model(**emotion_inputs).logits
    emotion_probs = torch.softmax(emotion_logits, dim=1).squeeze().cpu().numpy()
    emotion_confidences = {id_to_emotion_label[i]: float(prob) for i, prob in enumerate(emotion_probs)}

    # 장르 예측
    genre_inputs = tokenizer(preprocessed_text, return_tensors='pt', truncation=True, padding=True, max_length=512).to(device)
    with torch.no_grad():
        genre_logits = genre_model(**genre_inputs).logits
    genre_probs = torch.softmax(genre_logits, dim=1).squeeze().cpu().numpy()
    genre_confidences = {id_to_genre_label[i]: float(prob) for i, prob in enumerate(genre_probs)}

    return emotion_confidences, genre_confidences

# --- 4. Gradio 인터페이스 (이하 동일) ---
title = "🎤 한국어 가사 감정 및 장르 동시 분석기 🎶"
description = "KoBERT를 파인튜닝하여 만든 모델입니다. 가사를 입력하면 감정과 장르를 동시에 예측합니다."
examples = [
    ["슬픔의 밑바닥에서 난 너를 만나"],
    ["가슴이 웅장해진다 이건 못 참지"],
    ["너와 함께라면 어디든 갈 수 있어"],
    ["오늘 밤 주인공은 나야 나"]
]
iface = gr.Interface(
    fn=predict_emotion_and_genre,
    inputs=gr.Textbox(lines=10, placeholder="여기에 노래 가사를 입력하세요...", label="노래 가사"),
    outputs=[
        gr.Label(num_top_classes=3, label="감정 예측 결과"),
        gr.Label(num_top_classes=3, label="장르 예측 결과")
    ],
    title=title,
    description=description,
    examples=examples
)

iface.launch()