Update models/space_b.py

models/space_b.py

@@ -1,10 +1,15 @@
+# MentalQA – Arabic Mental Health Assistant (chat + classifier)
+
 import torch
-from transformers import (AutoTokenizer, AutoModelForCausalLM,
-                          AutoModelForSequenceClassification, pipeline)
+from transformers import (
+    AutoTokenizer, AutoModelForCausalLM,
+    AutoModelForSequenceClassification, pipeline
+)
 
 CHAT_REPO = "yasser-alharbi/MentalQA"
 CLASSIFIER_REPO = "yasser-alharbi/MentalQA-Classification"
 
+# Load chat model
 chat_tok = AutoTokenizer.from_pretrained(CHAT_REPO, use_fast=False)
 chat_model = AutoModelForCausalLM.from_pretrained(
     CHAT_REPO,
@@ -13,6 +18,7 @@ chat_model = AutoModelForCausalLM.from_pretrained(
     low_cpu_mem_usage=True,
 )
 
+# Load classifier
 clf_tok = AutoTokenizer.from_pretrained(CLASSIFIER_REPO)
 clf_model = AutoModelForSequenceClassification.from_pretrained(CLASSIFIER_REPO)
 
@@ -20,8 +26,13 @@ device_idx = 0 if torch.cuda.is_available() else -1
 clf_pipe = pipeline("text-classification", model=clf_model, tokenizer=clf_tok, device=device_idx)
 
 label_map = {
-    "LABEL_0": "A", "LABEL_1": "B", "LABEL_2": "C",
-    "LABEL_3": "D", "LABEL_4": "E", "LABEL_5": "F", "LABEL_6": "G"
+    "LABEL_0": "A",  # تشخيص
+    "LABEL_1": "B",  # علاج
+    "LABEL_2": "C",  # تشريح
+    "LABEL_3": "D",  # وبائيات
+    "LABEL_4": "E",  # نمط حياة
+    "LABEL_5": "F",  # مقدم خدمة
+    "LABEL_6": "G",  # أخرى
 }
 
 SYSTEM_MSG = (
@@ -29,29 +40,32 @@ SYSTEM_MSG = (
     "لا تذكر اسمك أو منصة عملك إلا إذا سُئلت صراحةً عن هويتك."
 )
 
-def classify_question(text: str, thr: float = 0.5) -> str:
+def classify_question(text: str, threshold: float = 0.5) -> str:
     pred = max(clf_pipe(text), key=lambda x: x["score"])
-    return label_map.get(pred["label"], pred["label"]) if pred["score"] >= thr else "G"
+    return label_map.get(pred["label"], "G") if pred["score"] >= threshold else "G"
 
-def build_prompt(question: str, tag: str) -> str:
+def build_prompt(question: str, final_qt: str) -> str:
     return (
-        f"{SYSTEM_MSG}\n\nfinal_QT: {tag}\n\n"
+        f"{SYSTEM_MSG}\n\n"
+        f"final_QT: {final_qt}\n\n"
         f"سؤال المستخدم:\n{question}\n\n"
         "اكتب فقرة واحدة مفصّلة لا تقل عن ثلاث جمل مترابطة، بعد أن تفكّر خطوة بخطوة.\n"
         "الإجابة النهائية:\n"
     )
 
-def generate_mentalqa_answer(question: str) -> str:
-    tag = classify_question(question)
-    prompt = build_prompt(question, tag)
-    chat_ids = chat_tok.apply_chat_template(
-        [{"role": "system", "content": SYSTEM_MSG}, {"role": "user", "content": prompt}],
+def generate_mentalqa_answer(question: str, threshold: float = 0.5) -> str:
+    final_qt = classify_question(question, threshold)
+    prompt = build_prompt(question, final_qt)
+
+    chat_input = chat_tok.apply_chat_template(
+        [{"role": "system", "content": SYSTEM_MSG},
+         {"role": "user", "content": prompt}],
         add_generation_prompt=True,
         return_tensors="pt"
     ).to(chat_model.device)
 
-    gen_ids = chat_model.generate(
-        chat_ids,
+    gen_output = chat_model.generate(
+        chat_input,
         max_new_tokens=128,
         do_sample=True,
         temperature=0.6,
@@ -62,5 +76,5 @@ def generate_mentalqa_answer(question: str) -> str:
         eos_token_id=chat_tok.eos_token_id,
     )[0]
 
-    answer_ids = gen_ids[chat_ids.shape[1]:]
-    return chat_tok.decode(answer_ids, skip_special_tokens=True).strip()
+    answer = chat_tok.decode(gen_output[chat_input.shape[1]:], skip_special_tokens=True)
+    return answer.strip()