| 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"] |
|
|
|
|
| @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() |
| return {c: round(float(p), 4) for c, p in zip(le.classes_, probs)} |
|
|
|
|
| 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} |
|
|
