core_ai.py

import os
import re
import pickle
import warnings
from functools import lru_cache

import numpy as np
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from deep_translator import GoogleTranslator

warnings.filterwarnings("ignore")

CLASSES = ["anxiety", "depression", "stress"]


# ── KEYWORD OVERRIDE ─────────────────────────────────────────────────────────
DEPRESSION_KEYWORDS = [
    # Arabic (Fusha + dialects)
    "اكتئاب", "مكتئب", "مكتئبة", "حزن", "حزين", "حزينة", "يأس", "يائس", "يائسة",
    "فراغ", "إحساس بالفراغ", "بلا معنى", "لا معنى", "مالهاش معنى", "بلا هدف",
    "لا أمل", "مفيش أمل", "تعبت من الحياة", "زهقت من الحياة",
    "مش لاقي معنى", "مش لاقية معنى", "حاسس بالفراغ", "حاسة بالفراغ",
    "مفيش طاقة", "مفيش رغبة", "بكاء", "عايز أبكي", "عايزة أبكي",
    "وحيد", "وحيدة", "عزلة", "منعزل", "منعزلة",
    "إرهاق نفسي", "إرهاق عاطفي", "مش حاسس بحاجة", "مش حاسة بحاجة",
    "زهقت", "تعبت", "مش طايق", "مش طايقة", "نفسيتي وحشة", "نفسيتي في الأرض",
    "مش قادر أكمل", "مش قادرة أكمل", "مش عايش", "مش قادر أعيش",
    "مش عايز أصحى", "مش عايزة أصحى", "دموع", "بدمع", "قلبي تقيل",
    "مش حاسس بنفسي", "مش حاسة بنفسي", "ما بحس بشي", "ما في فايدة",
    "مافي امل", "ما في امل", "حياتي خربت", "خسرت كل حاجة",
    # English
    "depressed", "depression", "hopeless", "hopelessness", "empty", "emptiness",
    "worthless", "meaningless", "no meaning", "no purpose", "cannot go on",
    "cant go on", "no energy", "no motivation", "crying", "feel nothing",
    "numb", "isolated", "lonely", "loneliness", "sad", "sadness",
    "despair", "grief", "miserable", "broken", "lost all hope",
]

ANXIETY_KEYWORDS = [
    # Arabic
    "قلق", "قلقان", "قلقانة", "خوف", "خايف", "خايفة", "توتر", "متوتر", "متوترة",
    "هلع", "مش مرتاح", "مش مرتاحة", "ذعر", "رهاب", "وسواس",
    # English
    "panic", "anxious", "anxiety", "worried", "worry", "fear",
    "scared", "nervous", "restless", "tense", "phobia", "ocd",
]

STRESS_KEYWORDS = [
    # Arabic
    "ضغط", "ضغوط", "مضغوط", "مضغوطة", "إجهاد", "مجهد", "مجهدة",
    # English
    "overwhelmed", "stressed", "stress", "burnout", "exhausted", "overloaded",
]


def keyword_boost(text: str, scores: dict) -> dict:
    text_lower = text.lower()

    dep_hits = sum(1 for kw in DEPRESSION_KEYWORDS if kw.lower() in text_lower)
    anx_hits = sum(1 for kw in ANXIETY_KEYWORDS if kw.lower() in text_lower)
    str_hits = sum(1 for kw in STRESS_KEYWORDS if kw.lower() in text_lower)

    if dep_hits == 0 and anx_hits == 0 and str_hits == 0:
        return scores

    s = dict(scores)

    if dep_hits > 0 and dep_hits >= anx_hits and dep_hits >= str_hits:
        boost = min(0.55 + dep_hits * 0.10, 0.85)
        s["depression"] = boost
        remaining = 1.0 - boost
        total_rest = s["anxiety"] + s["stress"]
        if total_rest > 0:
            s["anxiety"] = round(remaining * s["anxiety"] / total_rest, 4)
            s["stress"] = round(remaining * s["stress"] / total_rest, 4)
        s["depression"] = round(boost, 4)

    elif anx_hits > 0 and anx_hits >= dep_hits and anx_hits >= str_hits:
        boost = min(0.55 + anx_hits * 0.10, 0.85)
        s["anxiety"] = boost
        remaining = 1.0 - boost
        total_rest = s["depression"] + s["stress"]
        if total_rest > 0:
            s["depression"] = round(remaining * s["depression"] / total_rest, 4)
            s["stress"] = round(remaining * s["stress"] / total_rest, 4)
        s["anxiety"] = round(boost, 4)

    elif str_hits > 0 and str_hits >= dep_hits and str_hits >= anx_hits:
        boost = min(0.55 + str_hits * 0.10, 0.85)
        s["stress"] = boost
        remaining = 1.0 - boost
        total_rest = s["depression"] + s["anxiety"]
        if total_rest > 0:
            s["depression"] = round(remaining * s["depression"] / total_rest, 4)
            s["anxiety"] = round(remaining * s["anxiety"] / total_rest, 4)
        s["stress"] = round(boost, 4)

    total = sum(s.values())
    if total > 0:
        s = {k: round(v / total, 4) for k, v in s.items()}

    return s


@lru_cache(maxsize=1)
def load_xlmr():
    model_id = os.getenv("HF_MODEL_ID", "AliSakr9997/Mental-XLMR-Model")
    token = os.getenv("HF_TOKEN")
    kwargs = {"token": token} if token else {}
    local_dir = os.path.join(os.path.dirname(__file__), "mental_xlmr_final")
    local_weights = any(
        os.path.exists(os.path.join(local_dir, fname))
        for fname in ("pytorch_model.bin", "model.safetensors")
    )
    source = local_dir if local_weights else model_id
    tokenizer = AutoTokenizer.from_pretrained(source, **kwargs)
    model = AutoModelForSequenceClassification.from_pretrained(source, **kwargs)
    le_path = os.path.join(os.path.dirname(__file__), "mental_xlmr_final", "label_encoder.pkl")
    with open(le_path, "rb") as f:
        le = pickle.load(f)
    model.eval()
    return tokenizer, model, le


@lru_cache(maxsize=1)
def load_survey():
    scaler = pickle.load(open(os.path.join(os.path.dirname(__file__), "scaler.pkl"), "rb"))
    weights = pickle.load(open(os.path.join(os.path.dirname(__file__), "model_weights.pkl"), "rb"))

    def predict(x):
        for w in weights:
            if len(w) == 2:
                x = np.dot(x, w[0]) + w[1]
                x = np.maximum(0, x)
        x = np.exp(x) / np.sum(np.exp(x))
        return x

    return scaler, predict


def clean_text(text: str) -> str:
    text = re.sub(r"(.)\1{2,}", r"\1\1", text)
    text = re.sub(r"[^\w\s\u0600-\u06FF\[\]]", " ", text)
    return re.sub(r"\s+", " ", text).strip()


def translate_to_en(text: str) -> str:
    try:
        return GoogleTranslator(source="auto", target="en").translate(text)
    except Exception:
        return ""


def predict_text(text: str) -> dict:
    tokenizer, model, le = load_xlmr()
    cleaned = clean_text(text)
    text_en = translate_to_en(cleaned)
    combined = (text_en + " [SEP] " + cleaned) if text_en else cleaned
    inputs = tokenizer(combined, return_tensors="pt", truncation=True, max_length=192, padding=True)
    with torch.no_grad():
        probs = torch.softmax(model(**inputs).logits, dim=-1).squeeze().numpy()
    raw_scores = {c: round(float(p), 4) for c, p in zip(le.classes_, probs)}
    boosted = keyword_boost(text + " " + text_en, raw_scores)
    return boosted


def predict_survey(answers: list) -> dict:
    scaler, survey_predict = load_survey()
    data = scaler.transform(np.array(answers).reshape(1, -1))
    pred = survey_predict(data)[0]
    return {
        "depression": round(float(pred[0]), 4),
        "anxiety": round(float(pred[1]), 4),
        "stress": round(float(pred[2]), 4),
    }


def fuse_scores(text_s, survey_s, w_text=0.4, w_survey=0.6):
    return {c: round(w_text * text_s[c] + w_survey * survey_s[c], 4) for c in CLASSES}