현진_app.py 우울모델 추가

app.py

@@ -3,11 +3,15 @@ import re
 import time
 import requests
 import numpy as np
+import torch
+import joblib
+import xgboost as xgb
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from sentence_transformers import SentenceTransformer
-import os
 from typing import Optional, List,Dict
+from transformers import BertTokenizer, BertForSequenceClassification
+
 
 
 #####################################
@@ -1147,6 +1151,39 @@ def chat_response(user_input, mode="emotion", max_retries=5):
     return "🚨 모델 로딩이 너무 오래 걸립니다. 잠시 후 다시 시도하세요."
 
 
+#우울분류 모델 추가
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+tokenizer = BertTokenizer.from_pretrained("monologg/kobert")
+bert_model = BertForSequenceClassification.from_pretrained("monologg/kobert", num_labels=2)
+bert_model.load_state_dict(torch.load("emotion_bert_model.pth", map_location=device))
+bert_model.to(device)
+bert_model.eval()
+
+xgb_model = joblib.load("xgboost_model.pkl")
+vectorizer = joblib.load("tfidf_vectorizer.pkl")
+
+def predict_depression(text: str):
+    encoding = tokenizer(text, truncation=True, padding="max_length", max_length=64, return_tensors="pt")
+    input_ids = encoding["input_ids"].to(device)
+    attention_mask = encoding["attention_mask"].to(device)
+    with torch.no_grad():
+        outputs = bert_model(input_ids, attention_mask=attention_mask)
+        probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)
+        kobert_score = probabilities[0][1].item()
+    text_vec = vectorizer.transform([text])
+    xgb_proba = xgb_model.predict_proba(text_vec)[0][1]
+    kobert_score = max(0.35, min(kobert_score, 0.88))
+    xgb_proba = max(0.3, min(xgb_proba, 0.83))
+    combined_score = (kobert_score * 0.55) + (xgb_proba * 0.45)
+    if combined_score > 0.78:
+        label = "상담 권장"
+    elif combined_score > 0.65:
+        label = "관심 필요"
+    else:
+        label = "정상"
+    return combined_score, label
+
 #####################################
 # 6) FastAPI Endpoint
 #####################################
@@ -1202,6 +1239,19 @@ class ChatOrRecommendRequest(BaseModel):
 # (5) 자동 분기 엔드포인트
 @app.post("/chat_or_recommend")
 def chat_or_recommend(req: ChatOrRecommendRequest):
+    depression_score, depression_label = predict_depression(req.user_input)
+    if depression_label == "상담 권장":
+        counseling_response = (
+            "입력하신 메시지에서 심각한 우울 신호가 감지되었습니다.\n"
+            "전문 상담을 받으실 것을 강력히 권장드립니다.\n"
+            "빠른 시일 내에 전문가와 상담하시길 바랍니다."
+        )
+        return {
+            "mode": "counseling",
+            "response": counseling_response,
+            "depression_score": round(depression_score, 4),
+            "depression_label": depression_label
+        }
     user_text = req.user_input
     mode = req.mode.lower()