Update models/space_a.py

models/space_a.py

@@ -1,8 +1,12 @@
+# -*- coding: utf-8 -*-
+# ALLaM 7B – Arabic Mental Health Question Summarizer
+
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_id = "ALLaM-AI/ALLaM-7B-Instruct-preview"
+
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
@@ -11,7 +15,27 @@ model = AutoModelForCausalLM.from_pretrained(
     low_cpu_mem_usage=True,
 )
 
-def summarize_question(text: str, max_new_tokens: int = 60) -> str:
+def _generate(system_text: str, user_text: str, max_new_tokens: int = 60) -> str:
+    messages = []
+    if system_text.strip():
+        messages.append({"role": "system", "content": system_text})
+    messages.append({"role": "user", "content": user_text})
+
+    prompt_text = tokenizer.apply_chat_template(messages, tokenize=False)
+    inputs = tokenizer(prompt_text, return_tensors="pt").to(model.device)
+
+    with torch.no_grad():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=False,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+
+    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
+    return decoded.replace("[/INST]", "").strip()
+
+def summarize_question(question_text: str) -> str:
     examples = """
 أمثلة على كيفية التفكير خطوة بخطوة ثم إعطاء السؤال المختصر:
 السؤال الأصلي:
@@ -22,35 +46,30 @@ def summarize_question(text: str, max_new_tokens: int = 60) -> str:
 3. يسأل عن الحصول على وصفة Ativan كبديل أو إضافة للعلاج.
 السؤال المختصر:
 استفسار حول عدم فعالية citalopram والحاجة الى Ativan
+السؤال الأصلي:
+كنت اعاني من قلق وأخذت دواء سبرالكس لمدة 4 شهور وتوقفت عنه في 2014. سأذهب لتحليل بول للعمل، هل سيظهر أثر الدواء؟
+فكر خطوة بخطوة:
+1. المستخدم كان يعاني من قلق وتناول سبرالكس قبل عدة سنوات.
+2. يخشى أن يظهر الدواء القديم في فحص البول المطلوب للعمل.
+3. يريد معرفة إن كان ما زال موجوداً في جسمه.
+السؤال المختصر:
+سؤال حول بقاء أثر الدواء في تحليل البول بعد مدة؟
 """.strip()
 
     new_segment = f"""
 الآن لديك سؤال جديد، فكر خطوة بخطوة بشكل مشابه ثم أعطني السؤال المختصر:
 تأكد من عدم اضافة ملاحظات أو اضافة معلومة غير موجودة في السؤال
 السؤال الأصلي:
-{text}
+{question_text}
 السؤال المختصر:
-"""
+""".strip()
 
-    full_prompt = f"""
+    final_prompt = f"""
 أنت مساعد لغوي مختص بأسئلة الصحة النفسية.
 هدفك هو إعادة كتابة الأسئلة المطوّلة في شكل مختصر ومباشر، يركّز على النقطة الأساسية.
 {examples}
 {new_segment}
 """.strip()
 
-    inputs = tokenizer(
-        tokenizer.apply_chat_template([{"role": "user", "content": full_prompt}], tokenize=False),
-        return_tensors="pt"
-    ).to(model.device)
-
-    with torch.no_grad():
-        output = model.generate(
-            **inputs,
-            max_new_tokens=max_new_tokens,
-            do_sample=False,
-            pad_token_id=tokenizer.eos_token_id
-        )
-
-    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
-    return decoded.split("السؤال المختصر:")[-1].strip() if "السؤال المختصر:" in decoded else decoded.strip()
+    response = _generate("", final_prompt)
+    return response.split("السؤال المختصر:")[-1].strip() if "السؤال المختصر:" in response else response.strip()