Update app.py

app.py

@@ -3,12 +3,16 @@ import pandas as pd
 import numpy as np
 import torch
 from transformers import pipeline
+from sentence_transformers import SentenceTransformer, util
 
 # ===============================
-# 모델 로드
+# 설정
 # ===============================
 device = 0 if torch.cuda.is_available() else -1
 
+# ===============================
+# 모델 로드
+# ===============================
 stt_model = pipeline(
     "automatic-speech-recognition",
     model="openai/whisper-large-v3-turbo",
@@ -19,13 +23,22 @@ emotion_model = pipeline(
     "text-classification",
     model="monologg/koelectra-base-v3-goemotions",
     device=device,
-    top_k=1
+    top_k=None
 )
 
-# 데이터 로드
+sbert_model = SentenceTransformer("jhgan/ko-sroberta-multitask")
+
+# ===============================
+# 데이터 로드 + 임베딩 캐싱
+# ===============================
 df = pd.read_csv("book_db_final.csv")
 
+book_texts = df["contents"].fillna(df["title"]).tolist()
+book_embeddings = sbert_model.encode(book_texts, convert_to_tensor=True)
+
+# ===============================
 # 감정 매핑
+# ===============================
 EMOTION_MAP = {
     "joy": "기쁨",
     "sadness": "슬픔",
@@ -37,19 +50,77 @@ EMOTION_MAP = {
     "optimism": "기대"
 }
 
+EMOTION_LABELS = ["기쁨","신뢰","공포","놀람","슬픔","혐오","분노","기대"]
+
 # ===============================
-# 메인 처리 함수
+# 감정 분석
+# ===============================
+def get_emotion_scores(text):
+    results = emotion_model(text)[0]
+
+    scores = {emo: 0.0 for emo in EMOTION_LABELS}
+
+    # 모델 점수
+    for r in results:
+        label = r["label"].lower()
+        mapped = EMOTION_MAP.get(label)
+        if mapped:
+            scores[mapped] += r["score"]
+
+    # 한국어 보정
+    t = text.lower()
+    if "슬프" in t or "우울" in t:
+        scores["슬픔"] += 0.3
+    if "화나" in t or "짜증" in t:
+        scores["분노"] += 0.3
+    if "행복" in t or "좋다" in t:
+        scores["기쁨"] += 0.3
+
+    return scores
+
+# ===============================
+# 추천 (SBERT 최적화)
+# ===============================
+def recommend_books(user_text, emotion):
+    pool = df[df["emotion"] == emotion]
+
+    if pool.empty:
+        return []
+
+    idxs = pool.index.tolist()
+    pool_embs = book_embeddings[idxs]
+
+    user_emb = sbert_model.encode(user_text, convert_to_tensor=True)
+    sims = util.cos_sim(user_emb, pool_embs)[0].cpu().numpy()
+
+    pool = pool.copy()
+    pool["sim"] = sims
+
+    pool = pool.sort_values("sim", ascending=False).head(3)
+
+    books = []
+    for _, row in pool.iterrows():
+        books.append({
+            "title": row["title"],
+            "url": row["url"],
+            "contents": str(row.get("contents", ""))[:120]
+        })
+
+    return books
+
+# ===============================
+# 메인 함수
 # ===============================
 def process_voice_only(audio_input):
     if audio_input is None:
         return {"error": "음성을 녹음해주세요."}
 
     try:
-        # 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
         stt_result = stt_model({"sampling_rate": sr, "raw": y})
         final_text = stt_result["text"]
 
@@ -57,31 +128,28 @@ def process_voice_only(audio_input):
             return {"error": "음성이 인식되지 않았습니다."}
 
         # 감정 분석
-        emo_result = emotion_model(final_text)[0][0]
-        raw_label = emo_result["label"].lower()
-        best_emo = EMOTION_MAP.get(raw_label, "기대")
+        scores = get_emotion_scores(final_text)
 
-        # 추천
-        recs = df[df["emotion"] == best_emo].head(3)
+        best_emo = max(scores, key=scores.get)
+        top3 = sorted(scores.items(), key=lambda x: x[1], reverse=True)[:3]
 
-        books = []
-        for _, row in recs.iterrows():
-            books.append({
-                "title": row["title"],
-                "url": row["url"],
-                "contents": str(row.get("contents", ""))[:120]
-            })
+        # 추천
+        books = recommend_books(final_text, best_emo)
 
         return {
             "text": final_text,
             "emotion": best_emo,
+            "emotion_scores": {k: round(v, 3) for k, v in scores.items()},
+            "top3": [
+                {"emotion": e, "score": round(s, 3)}
+                for e, s in top3
+            ],
             "books": books
         }
 
     except Exception as e:
         return {"error": str(e)}
 
-
 # ===============================
 # UI
 # ===============================
@@ -91,15 +159,15 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             audio_in = gr.Audio(label="마이크 입력", sources=["microphone"])
-            submit_btn = gr.Button("분석", variant="primary")
+            btn = gr.Button("분석", variant="primary")
 
         with gr.Column():
-            output_json = gr.JSON(label="결과")
+            output = gr.JSON(label="결과")
 
-    submit_btn.click(
+    btn.click(
         fn=process_voice_only,
-        inputs=[audio_in],
-        outputs=output_json
+        inputs=audio_in,
+        outputs=output
     )
 
 demo.launch()