Update app.py

app.py

@@ -1,77 +1,89 @@
 import gradio as gr
 import pandas as pd
 import numpy as np
-from sentence_transformers import SentenceTransformer, util
 from transformers import pipeline
 
-# 1. 모델 및 데이터 로드
-# 빠른 속도를 위해 whisper-tiny 사용
-stt_model = pipeline("automatic-speech-recognition", model="openai/whisper-tiny") 
-sbert_model = SentenceTransformer("jhgan/ko-sroberta-multitask")
+# ===============================
+# 모델 로드
+# ===============================
+stt_model = pipeline(
+    "automatic-speech-recognition",
+    model="openai/whisper-large-v3-turbo",
+    device=0
+)
 
-# 데이터베이스 로드 (파일명 확인 필요)
+emotion_model = pipeline(
+    "text-classification",
+    model="monologg/koelectra-base-v3-goemotions",
+    device=0,
+    top_k=1
+)
+
+# 데이터 로드
 df = pd.read_csv("book_db_final.csv")
 
-_EMOTION_DESCS = {
-    "기쁨": "행복하고 즐겁고 유쾌한 기분",
-    "신뢰": "따뜻하고 안정적이며 가족과 우정 같은 유대감",
-    "공포": "무섭고 긴장되며 스릴 있는 공포와 불안",
-    "놀람": "반전과 충격, 예상치 못한 경이로움",
-    "슬픔": "슬프고 외롭고 이별과 상실의 감정",
-    "혐오": "부조리와 불평등, 위선에 대한 비판과 풍자",
-    "분노": "분노와 저항, 투쟁과 갈등",
-    "기대": "성장과 도전, 모험과 희망",
+# 감정 매핑 (Plutchik 유지)
+EMOTION_MAP = {
+    "joy": "기쁨",
+    "sadness": "슬픔",
+    "anger": "분노",
+    "fear": "공포",
+    "surprise": "놀람",
+    "disgust": "혐오",
+    "love": "신뢰",
+    "optimism": "기대"
 }
 
-_LABEL_EMBS = sbert_model.encode(list(_EMOTION_DESCS.values()), convert_to_tensor=True)
-
+# ===============================
+# 메인 처리 함수
+# ===============================
 def process_voice_only(audio_input):
-    # 1. 음성 입력 확인
     if audio_input is None:
         return "음성을 녹음해주세요.", ""
 
-    # 2. STT (Speech to Text) 변환
+    # STT
     sr, y = audio_input
     y = y.astype(np.float32)
     y /= np.max(np.abs(y)) if np.max(np.abs(y)) > 0 else 1
-    
+
     stt_result = stt_model({"sampling_rate": sr, "raw": y})
     final_text = stt_result["text"]
 
     if not final_text.strip():
-        return "음성이 인식되지 않았습니다. 다시 말씀해주세요.", ""
-
-    # 3. 감정 분석
-    user_emb = sbert_model.encode(final_text, convert_to_tensor=True)
-    scores = util.cos_sim(user_emb, _LABEL_EMBS)[0]
-    best_emo = list(_EMOTION_DESCS.keys())[scores.argmax()]
-    
-    # 4. 도서 추천
+        return "음성이 인식되지 않았습니다.", ""
+
+    # 감정 분석 (classifier 사용)
+    emo_result = emotion_model(final_text)[0][0]
+    raw_label = emo_result["label"].lower()
+    best_emo = EMOTION_MAP.get(raw_label, "기대")
+
+    # 추천
     recs = df[df["emotion"] == best_emo].head(3)
-    
+
     result_text = f"인식된 문장: \"{final_text}\"\n분석된 감정: {best_emo}\n\n"
-    
+
     book_list = ""
     for _, row in recs.iterrows():
-        book_list += f"{row['title']}\n {row['url']}\n\n"
-        
+        book_list += f"{row['title']}\n{row['url']}\n\n"
+
     return result_text, book_list
 
-# --- Gradio UI 구성 ---
+
+# ===============================
+# UI
+# ===============================
 with gr.Blocks() as demo:
-    gr.Markdown("# Boolook: 음성 전용 마음 분석 책 추천")
-    
+    gr.Markdown("# Boolook: 음성 기반 감정 분석 책 추천")
+
     with gr.Row():
         with gr.Column():
-            # 오디오 입력만 유지
-            audio_in = gr.Audio(label="마이크로 마음을 들려주세요", sources=["microphone"])
-            submit_btn = gr.Button("분석 및 추천받기", variant="primary")
-            
+            audio_in = gr.Audio(label="마이크 입력", sources=["microphone"])
+            submit_btn = gr.Button("분석", variant="primary")
+
         with gr.Column():
             analysis_out = gr.Textbox(label="분석 결과")
-            books_out = gr.Textbox(label="추천 도서 리스트")
+            books_out = gr.Textbox(label="추천 도서")
 
-    # 버튼 클릭 시 음성 처리 함수 연결
     submit_btn.click(
         fn=process_voice_only,
         inputs=[audio_in],