Update app.py

app.py

@@ -1,35 +1,21 @@
-# app.py
-
 import torch
 import re
+import json
+import gradio as gr
 from konlpy.tag import Okt
 from transformers import AutoTokenizer, BertForSequenceClassification
-import gradio as gr
 
 # --- 1. 설정 및 전처리 함수 ---
 
 # 디바이스 설정
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# 저장된 모델과 토크나이저가 있는 경로
-MODEL_DIR = './kobert_genre_classifier_archive'
+# 각 모델이 저장된 경로
+EMOTION_MODEL_DIR = './kobert_emotion_classifier_archive'
+GENRE_MODEL_DIR = './kobert_genre_classifier_archive'
 
-# 레이블 맵핑 (※※※ 중요: 노트북에서 사용한 레이블 순서에 맞게 실제 장르명으로 수정하세요 ※※※)
-# 예시: labels_ids = {'댄스': 0, '발라드': 1, '록': 2, ...}
-id_to_label = {
-    0: '록/메탈', # 실제 장르 이름으로 변경
-    1: '댄스',
-    2: 'R&B/Soul',
-    3: '발라드',
-    4: '랩/힙합합',
-    5: '트로트'
-}
-# ↑↑↑↑↑ 이 부분을 실제 장르명으로 꼭 수정해주세요! ↑↑↑↑↑
-
-
-# 노트북에서 사용한 전처리 함수 (그대로 복사)
+# 노트북에서 사용한 전처리 함수 (공통 사용)
 okt = Okt()
-
 def remove_english(text):
     return re.sub(r'[A-Za-z]+', '', text)
 
@@ -38,69 +24,82 @@ def extract_pos(text):
     text = remove_english(text)
     return ' '.join([word for word, pos in okt.pos(text) if pos in allowed_pos])
 
-# --- 2. 모델 및 토크나이저 로드 ---
+# --- 2. 두 모델 및 토크나이저 로드 ---
 
 try:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, trust_remote_code=True)
-    model = BertForSequenceClassification.from_pretrained(MODEL_DIR, trust_remote_code=True)
-    model.to(device)
-    model.eval()
-    print("모델과 토크나이저를 성공적으로 로드했습니다.")
+    # 감정 분류 모델 로드
+    emotion_tokenizer = AutoTokenizer.from_pretrained(EMOTION_MODEL_DIR, trust_remote_code=True)
+    emotion_model = BertForSequenceClassification.from_pretrained(EMOTION_MODEL_DIR, trust_remote_code=True)
+    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_tokenizer = AutoTokenizer.from_pretrained(GENRE_MODEL_DIR, trust_remote_code=True)
+    genre_model = BertForSequenceClassification.from_pretrained(GENRE_MODEL_DIR, trust_remote_code=True)
+    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}")
-    # Hugging Face Spaces에서는 오류 발생 시 앱이 실행되지 않도록 처리
-    tokenizer, model = None, None
-
-# --- 3. 예측 함수 ---
+    emotion_model, genre_model = None, None
 
-def predict_genre(text):
-    if not model or not tokenizer:
+# --- 3. 통합 예측 함수 ---
+def predict_emotion_and_genre(text):
+    if not emotion_model or not genre_model:
         raise gr.Error("모델이 로드되지 않았습니다. Space의 로그를 확인해주세요.")
 
-    # 1. 입력된 가사 전처리
+    # 1. 입력된 가사 공통 전처리
     preprocessed_text = extract_pos(text)
 
-    # 2. 토크나이징
-    inputs = tokenizer(
-        preprocessed_text,
-        return_tensors='pt',
-        truncation=True,
-        padding='max_length',
-        max_length=512 # 노트북에서 설정한 MAX_LENGTH와 동일하게
+    # 2. 감정 예측 수행
+    emotion_inputs = emotion_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)}
 
-    # 3. 예측
+    # 3. 장르 예측 수행
+    genre_inputs = genre_tokenizer(
+        preprocessed_text, return_tensors='pt', truncation=True, padding=True, max_length=512
+    ).to(device)
     with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
+        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)}
 
-    # 4. 확률값 계산 및 결과 포매팅
-    probabilities = torch.softmax(logits, dim=1).squeeze().cpu().numpy()
-    
-    # Gradio의 Label 컴포넌트에 맞게 출력 형식 변경
-    confidences = {id_to_label[i]: float(prob) for i, prob in enumerate(probabilities)}
-    
-    return confidences
+    return emotion_confidences, genre_confidences
 
 # --- 4. Gradio 인터페이스 생성 ---
-
-title = "🎤 한국어 노래 가사 장르 분류기 🎶"
-description = "KoBERT를 파인튜닝하여 만든 노래 가사 장르 분류 모델입니다. 아래에 가사를 입력하고 '분류하기' 버튼을 누르면 장르를 예측해줍니다."
+title = "🎤 한국어 가사 감정 및 장르 동시 분석기 🎶"
+description = "KoBERT를 파인튜닝하여 만든 모델입니다. 가사를 입력하면 감정과 장르를 동시에 예측합니다."
 examples = [
-    ["무얼 믿은 걸까 부족했던 내게서 나조차 못 믿던 내게 여태 머문 사람"],
-    ["미치도록 사랑했던 지겹도록 다투었던 네가 먼저 떠나고 여긴 온종일 비가 왔어"],
-    ["우린 멋진 나이야 좀 어리긴 하지만 하고픈 일이나 가고픈 길 해야 할 일들까지"]
+    ["슬픔의 밑바닥에서 난 너를 만나"],
+    ["가슴이 웅장해진다 이건 못 참지"],
+    ["너와 함께라면 어디든 갈 수 있어"],
+    ["오늘 밤 주인공은 나야 나"]
 ]
 
-
-# Gradio 인터페이스 실행
 iface = gr.Interface(
-    fn=predict_genre,
+    fn=predict_emotion_and_genre,
     inputs=gr.Textbox(lines=10, placeholder="여기에 노래 가사를 입력하세요...", label="노래 가사"),
-    outputs=gr.Label(num_top_classes=3, label="예측된 장르"),
+    # ✅ outputs 부분을 수정하여 두 결과 모두 상위 3개를 표시하도록 설정합니다.
+    outputs=[
+        gr.Label(num_top_classes=3, label="감정 예측 결과"),
+        gr.Label(num_top_classes=3, label="장르 예측 결과")
+    ],
     title=title,
     description=description,
     examples=examples
 )
 
+# 앱 실행
 iface.launch()