Update models/space_b.py

models/space_b.py

@@ -1,123 +1,58 @@
-# app.py ─────────────────────────────────────────────────────────────
-"""
-MentalQA demo Space
-Loads:
-  • yasser-alharbi/MentalQA               (ALLaM-7B-based chat model)
-  • yasser-alharbi/MentalQA-Classification  (final_QT intent classifier)
-and exposes an Arabic RTL Gradio interface.
-"""
+import torch
+from transformers import (AutoTokenizer, AutoModelForCausalLM,
+                          AutoModelForSequenceClassification, pipeline)
 
-
-import torch, gradio as gr
-from transformers import (AutoTokenizer,
-                          AutoModelForCausalLM,
-                          AutoModelForSequenceClassification,
-                          pipeline)
-
-# ─────────────────── HF repos ──────────────────────────────────────
-CHAT_REPO       = "yasser-alharbi/MentalQA"
+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,
     torch_dtype="auto",
-    device_map="auto",        # works for CPU or GPU Space
+    device_map="auto",
     low_cpu_mem_usage=True,
 )
 
-# ─────────────────── Load classifier ───────────────────────────────
 clf_tok = AutoTokenizer.from_pretrained(CLASSIFIER_REPO)
 clf_model = AutoModelForSequenceClassification.from_pretrained(CLASSIFIER_REPO)
 
 device_idx = 0 if torch.cuda.is_available() else -1
-clf_pipe = pipeline("text-classification",
-                    model=clf_model,
-                    tokenizer=clf_tok,
-                    device=device_idx)
-
+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"
 }
 
-# ─────────────────── Prompt helpers ────────────────────────────────
 SYSTEM_MSG = (
-    "أنت مساعد ذكي للصحة النفسية اسمه MentalQA"
+    "أنت مساعد ذكي للصحة النفسية اسمه MentalQA. "
     "لا تذكر اسمك أو منصة عملك إلا إذا سُئلت صراحةً عن هويتك."
-    "بالإضافة إلى ذلك:\n"
-    "عندما يحييك أحد بتحية عربية:\n"
-    "     - السلام عليكم => وعليكم السلام\n"
-    "     - صباح الخير   => صباح النور\n"
-    "     - مساء الخير   => مساء النور\n\n"
 )
 
-def build_prompt_arabic(question, final_qt_list):
-    qt_str = ", ".join(final_qt_list)
-
-    prompt = (
-        # ── Core rules ─────────────────────────────────────────
-        "أجب باللغة العربية استنادًا إلى القواعد التالية:\n"
-        "1) هذه ليست استشارة طبية بديلة؛ قدّم إرشادات عامة وتمهيدية.\n"
-        "2) لا تستخدم أسماء شخصية أو تدّعي ملكية.\n"
-        "3) إذا كان السؤال خارج الصحة النفسية، قل: 'عذراً، ولكن هذا السؤال خارج نطاق قدرتي.'\n"
-        "4) استرشد بقيم final_QT (A تشخيص، B علاج، C تشريح، D وبائيات، "
-        "E نمط حياة، F خيارات مقدم الخدمة، G أخرى).\n"
-        "5) إذا كانت حالة المريض حرجة، أبدِ تعاطفك أولاً ثم وجّه النصيحة.\n"
-        "6) إذا احتاج المريض لتوجيه مباشر، قل: 'قد يفيد التواصل مع مختص نفسي أو مستشار موثوق.'\n\n"
-
-
-
-        # ── Few‑shot exemplar WITH reasoning ─────────────────────────────
-        "مثال توضيحي للإجابة المفصّلة مع خطوات التفكير:\n"
-        "سؤال: أشعر بإرهاقٍ مستمر ولا أستطيع التركيز، ماذا أفعل؟\n"
-        "التفكير خطوة بخطوة:\n"
-        "1) تحديد ما إذا كان الإرهاق جسدياً أم نفسياً.\n"
-        "2) فحص نمط النوم والعادات اليومية.\n"
-        "3) التفكير في عوامل الضغط والرعاية الذاتية.\n"
-        "4) وضع خطة من نصائح تدريجية سهلة التطبيق.\n"
-        "الإجابة النهائية:\n"
-        "قد يرتبط الإرهاق بعدم انتظام النوم أو بضغوطٍ نفسية متراكمة. "
-        "من المهم أولاً مراجعة نمط حياتك: اضبط مواعيد نوم ثابتة، وابتعد عن المنبّهات قبل النوم بساعتين. "
-        "مارس المشي الخفيف أو تمارين الاسترخاء يوميّاً لتخفيف التوتر. "
-        "إذا استمر الإرهاق أكثر من أسبوعين رغم هذه التغييرات، فكر في زيارة طبيب لفحص فيتامين د ووظائف الغدة الدرقية. "
-        "دوّن مشاعرك في مفكرة يومية لتفريغ القلق وتشخيص الأسباب بدقة.\n"
-        "—\n\n"
-
-
-        # ── User section ─────────────────────────────────────────────────
-        f"final_QT: {qt_str}\n\n"
-        "سؤال المستخدم:\n"
-        f"{question}\n\n"
+def classify_question(text: str, thr: 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"
 
-        # ── Final directive ──────────────────────────────────────────────
-        "اكتب فقرة واحدة مفصّلة لا تقل عن ثلاث جمل مترابطة، بعد أن تفكّر خطوة بخطوة، \n"
+def build_prompt(question: str, tag: str) -> str:
+    return (
+        f"{SYSTEM_MSG}\n\nfinal_QT: {tag}\n\n"
+        f"سؤال المستخدم:\n{question}\n\n"
+        "اكتب فقرة واحدة مفصّلة لا تقل عن ثلاث جمل مترابطة، بعد أن تفكّر خطوة بخطوة.\n"
         "الإجابة النهائية:\n"
     )
-    return prompt
-
-def classify_question(text: str, thr: float = 0.5):
-    pred = max(clf_pipe(text), key=lambda x: x["score"])
-    return label_map.get(pred["label"], pred["label"]) if pred["score"] >= thr else "G"
 
-def chat_generate(prompt: str, max_new_tokens: int = 128):
+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}],
+        [{"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,
-        max_new_tokens=max_new_tokens,
+        max_new_tokens=128,
         do_sample=True,
         temperature=0.6,
         top_p=0.95,
@@ -128,33 +63,4 @@ def chat_generate(prompt: str, max_new_tokens: int = 128):
     )[0]
 
     answer_ids = gen_ids[chat_ids.shape[1]:]
-    return chat_tok.decode(answer_ids,
-                           skip_special_tokens=True,
-                           clean_up_tokenization_spaces=True).strip()
-
-def get_mentalqa_answer(question: str, thr: float = 0.5):
-    tag   = classify_question(question, thr)
-    prompt= build_prompt_arabic(question, tag)
-    return chat_generate(prompt)
-
-# ─────────────────── Gradio UI ─────────────────────────────────────
-CSS = """
-#container{max-width:640px;margin:1.5rem auto;}
-#question_box label,#answer_box label,
-#question_box textarea,#answer_box textarea{
- direction:rtl;text-align:right;
-}
-"""
-
-with gr.Blocks(css=CSS, theme="soft") as demo:
-    gr.Markdown("<h2 style='text-align:center;'>🧠 MentalQA – مساعد الصحة النفسية</h2>"
-                "<p style='text-align:center;'>اكتب سؤالك النفسي باللغة العربية وسيجيبك النموذج.</p>")
-    with gr.Group(elem_id="container"):
-        q = gr.Textbox(lines=3, placeholder="اكتب سؤالك هنا...", label="سؤال:", elem_id="question_box")
-        a = gr.Textbox(lines=5, label="الإجابة:", elem_id="answer_box")
-        btn = gr.Button("إرسال")
-        btn.click(get_mentalqa_answer, inputs=q, outputs=a)
-        q.submit(get_mentalqa_answer, inputs=q, outputs=a)
-
-if __name__ == "__main__":
-    demo.launch()
+    return chat_tok.decode(answer_ids, skip_special_tokens=True).strip()