MentalQA / models /space_a.py
NLPAlhuzali's picture
Update models/space_a.py
6e92cbd verified
# -*- 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,
torch_dtype="auto",
device_map="auto",
low_cpu_mem_usage=True,
)
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 = """
أمثلة على كيفية التفكير خطوة بخطوة ثم إعطاء السؤال المختصر:
السؤال الأصلي:
انا فيني اكتئاب وقلق ووصف لي دكتور citalopram استخدمتها ثلاث شهور ولا نفع احتاج وصفه Ativan
فكر خطوة بخطوة:
1. المستخدم يعاني من اكتئاب وقلق.
2. تناول دواء citalopram لثلاثة أشهر من دون تحسّن.
3. يسأل عن الحصول على وصفة Ativan كبديل أو إضافة للعلاج.
السؤال المختصر:
استفسار حول عدم فعالية citalopram والحاجة الى Ativan
السؤال الأصلي:
كنت اعاني من قلق وأخذت دواء سبرالكس لمدة 4 شهور وتوقفت عنه في 2014. سأذهب لتحليل بول للعمل، هل سيظهر أثر الدواء؟
فكر خطوة بخطوة:
1. المستخدم كان يعاني من قلق وتناول سبرالكس قبل عدة سنوات.
2. يخشى أن يظهر الدواء القديم في فحص البول المطلوب للعمل.
3. يريد معرفة إن كان ما زال موجوداً في جسمه.
السؤال المختصر:
سؤال حول بقاء أثر الدواء في تحليل البول بعد مدة؟
""".strip()
new_segment = f"""
الآن لديك سؤال جديد، فكر خطوة بخطوة بشكل مشابه ثم أعطني السؤال المختصر:
تأكد من عدم اضافة ملاحظات أو اضافة معلومة غير موجودة في السؤال
السؤال الأصلي:
{question_text}
السؤال المختصر:
""".strip()
final_prompt = f"""
أنت مساعد لغوي مختص بأسئلة الصحة النفسية.
هدفك هو إعادة كتابة الأسئلة المطوّلة في شكل مختصر ومباشر، يركّز على النقطة الأساسية.
{examples}
{new_segment}
""".strip()
response = _generate("", final_prompt)
return response.split("السؤال المختصر:")[-1].strip() if "السؤال المختصر:" in response else response.strip()