import json import os import joblib import numpy as np import pandas as pd import torch from transformers import AutoModelForSequenceClassification, AutoTokenizer def _prob_to_level(p: float, thr_high: float, thr_med: float) -> str: if p >= thr_high: return "High" if p >= thr_med: return "Medium" return "Low" def _fuse_level(chat_level: str, prof_level: str, s: float, thr_high: float, thr_med: float) -> str: # Safety-first overrides (reduce false negatives) if chat_level == "High": return "High" if chat_level == "Medium" and prof_level == "High": return "High" if s >= thr_high: return "High" if s >= thr_med: return "Medium" return "Low" class SuicideRiskPredictor: """ Real-time predictor for ONE user input: - text (string) - profile (dict of structured features matching your BD columns) """ def __init__( self, model_name: str = "banglabert", models_dir: str = "outputs/models", artifacts_dir: str = "outputs/artifacts", fusion_meta_path: str = None, profile_meta_path: str = "outputs/artifacts/profile_risk_meta.joblib", ): self.model_name = model_name # ---- Load fusion config (weights/thresholds) ---- if not fusion_meta_path: fusion_meta_path = os.path.join(artifacts_dir, f"fusion_meta_{model_name}.json") if os.path.exists(fusion_meta_path): with open(fusion_meta_path, "r", encoding="utf-8") as f: meta = json.load(f) else: meta = {} # Safely extract weights (with fallback) weights = meta.get("weights", {}) self.w_chat = float(weights.get("w_chat", weights.get("chat", 0.60))) self.w_prof = float(weights.get("w_prof", weights.get("profile", 0.40))) # Safely extract thresholds chat_thrs = meta.get("chat_thresholds", {"high": 0.70, "medium": 0.40}) self.chat_thr_high = float(chat_thrs.get("high", 0.70)) self.chat_thr_med = float(chat_thrs.get("medium", 0.40)) final_thrs = meta.get("final_thresholds", {"high": 0.65, "medium": 0.40}) self.final_thr_high = float(final_thrs.get("high", 0.65)) self.final_thr_med = float(final_thrs.get("medium", 0.40)) self.max_len = int(meta.get("max_len", 64)) # ---- Chat model ---- self.chat_model_dir = os.path.join(models_dir, f"chat_brain_{model_name}") self.tokenizer = AutoTokenizer.from_pretrained(self.chat_model_dir) self.chat_model = AutoModelForSequenceClassification.from_pretrained(self.chat_model_dir) self.chat_model.eval() self.device = "cuda" if torch.cuda.is_available() else "cpu" self.chat_model.to(self.device) # ---- Profile artifacts ---- self.pre = joblib.load(os.path.join(artifacts_dir, "profile_preprocessor.joblib")) self.gmm = joblib.load(os.path.join(artifacts_dir, "gmm.joblib")) self.kde = joblib.load(os.path.join(artifacts_dir, "kde.joblib")) self.profile_meta = joblib.load(profile_meta_path) if os.path.exists(profile_meta_path) else None # calibration percentiles for risk conversion if self.profile_meta: self.low_gmm, self.high_gmm = self.profile_meta["calibration_ll_percentiles_gmm"] self.low_kde, self.high_kde = self.profile_meta["calibration_ll_percentiles_kde"] self.p70 = self.profile_meta["p70"] self.p90 = self.profile_meta["p90"] self.kde_bandwidth = self.profile_meta["kde_bandwidth"] self.w_gmm = self.profile_meta["w_gmm"] self.w_kde = self.profile_meta["w_kde"] else: # fallback if meta not found self.low_gmm, self.high_gmm = (-50.0, -5.0) self.low_kde, self.high_kde = (-50.0, -5.0) self.p70, self.p90 = (0.70, 0.90) self.w_gmm, self.w_kde = (0.6, 0.4) def _chat_prob(self, text: str) -> float: enc = self.tokenizer( [str(text)], truncation=True, padding=True, max_length=self.max_len, return_tensors="pt", ).to(self.device) with torch.no_grad(): logits = self.chat_model(**enc).logits p = torch.softmax(logits, dim=1)[:, 1].detach().cpu().numpy()[0] return float(p) def _ll_to_risk(self, ll: np.ndarray, low: float, high: float) -> np.ndarray: norm = np.clip((ll - low) / (high - low + 1e-9), 0, 1) return 1 - norm def _profile_risk(self, profile_dict: dict) -> tuple[float, str]: """ profile_dict keys should match your BD structured columns AFTER preprocessing. Missing keys are allowed (will become NaN and get imputed). """ expected_cols = [ "age_group", "age", "gender", "profession_group", "religion", "hometown", "reason", "reason_description", "time", "temperature", "feels_like", "temp_min", "temp_max", "air_pressure", "air_humidity", "wind_speed", "wind_deg", "clouds_sky", "weather_main", "weather_description" ] # Populate missing columns with NaN filled_dict = {} for col in expected_cols: filled_dict[col] = profile_dict.get(col, np.nan) df = pd.DataFrame([filled_dict]) X = self.pre.transform(df) ll_gmm = self.gmm.score_samples(X) ll_kde = self.kde.score_samples(X) risk_gmm = self._ll_to_risk(ll_gmm, self.low_gmm, self.high_gmm)[0] risk_kde = self._ll_to_risk(ll_kde, self.low_kde, self.high_kde)[0] risk_profile = float(self.w_gmm * risk_gmm + self.w_kde * risk_kde) # level using stored percentiles if risk_profile >= self.p90: level = "High" elif risk_profile >= self.p70: level = "Medium" else: level = "Low" return risk_profile, level def predict_one(self, text: str, profile: dict) -> dict: chat_prob = self._chat_prob(text) chat_level = _prob_to_level(chat_prob, self.chat_thr_high, self.chat_thr_med) risk_profile, risk_level_profile = self._profile_risk(profile) final_risk_score = float(self.w_chat * chat_prob + self.w_prof * risk_profile) final_risk_level = _fuse_level( chat_level, risk_level_profile, final_risk_score, self.final_thr_high, self.final_thr_med ) return { "chat_prob": chat_prob, "chat_level": chat_level, "risk_profile": risk_profile, "risk_level_profile": risk_level_profile, "final_risk_score": final_risk_score, "final_risk_level": final_risk_level, }