Update models/space_a.py

models/space_a.py

@@ -1,170 +1,56 @@
-# -*- coding: utf-8 -*-
-# app.py – Arabic Questions Summarization in Mental Healthcare with ALLaM
-
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import gradio as gr
 
-# ------------------------------------------------------------------
-# 1. Device setup
-# ------------------------------------------------------------------
 device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Using device: {device}")
-
-# ------------------------------------------------------------------
-# 2. Load ALLaM – choose params by hardware
-# ------------------------------------------------------------------
 model_id = "ALLaM-AI/ALLaM-7B-Instruct-preview"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-
-                                 # 🖥️  CPU-only path
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     torch_dtype="auto",
-    device_map="auto",        # works for CPU or GPU Space
+    device_map="auto",
     low_cpu_mem_usage=True,
 )
 
-
-# ------------------------------------------------------------------
-# 3. Generation helpers  
-# ------------------------------------------------------------------
-def generate_chat_response(
-    system_text: str,
-    user_text: str,
-    max_new_tokens: int = 40,
-    temperature: float = 0.2,
-):
-    messages = []
-    if system_text.strip():
-        messages.append({"role": "system", "content": system_text})
-    messages.append({"role": "user", "content": user_text})
-
-    chat_input_text = tokenizer.apply_chat_template(messages, tokenize=False)
-    inputs = tokenizer(
-        chat_input_text,
-        return_tensors="pt",
-        return_token_type_ids=False,
-    )
-    inputs = {k: v.to(model.device) for k, v in inputs.items()}
-
-    with torch.inference_mode():
-        output_tokens = model.generate(
-            **inputs,
-            max_new_tokens=max_new_tokens,
-            do_sample=False,
-            pad_token_id=tokenizer.eos_token_id,
-        )
-
-    output_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
-    output_text = output_text.replace("[/INST]", "")
-    return output_text.strip()
-
-
-def generate_text(prompt: str, max_new_tokens: int = 60):
-    system_text = ""  # keep system empty exactly as in notebook
-    return generate_chat_response(
-        system_text=system_text,
-        user_text=prompt,
-        max_new_tokens=max_new_tokens,
-    )
-
-# ------------------------------------------------------------------
-# 4. Prompt construction 
-# ------------------------------------------------------------------
-def prompt_short_question_cot_few_shots(question_text: str) -> str:
+def summarize_question(text: str, max_new_tokens: int = 60) -> str:
     examples = """
 أمثلة على كيفية التفكير خطوة بخطوة ثم إعطاء السؤال المختصر:
-
 السؤال الأصلي:
 انا فيني اكتئاب وقلق ووصف لي دكتور citalopram استخدمتها ثلاث شهور ولا نفع احتاج وصفه Ativan
-
 فكر خطوة بخطوة:
 1. المستخدم يعاني من اكتئاب وقلق.
 2. تناول دواء citalopram لثلاثة أشهر من دون تحسّن.
 3. يسأل عن الحصول على وصفة Ativan كبديل أو إضافة للعلاج.
-
 السؤال المختصر:
 استفسار حول عدم فعالية citalopram والحاجة الى Ativan
-
-
-السؤال الأصلي:
-كنت اعاني من قلق وأخذت دواء سبرالكس لمدة 4 شهور وتوقفت عنه في 2014. سأذهب لتحليل بول للعمل، هل سيظهر أثر الدواء؟
-
-فكر خطوة بخطوة:
-1. المستخدم كان يعاني من قلق وتناول سبرالكس قبل عدة سنوات.
-2. يخشى أن يظهر الدواء القديم في فحص البول المطلوب للعمل.
-3. يريد معرفة إن كان ما زال موجوداً في جسمه.
-
-السؤال المختصر:
-سؤال حول بقاء أثر الدواء في تحليل البول بعد مدة؟
-
-
-السؤال الأصلي:
-أصبت بالاضطراب الوجداني منذ 2008، تكررت النوبات عدة مرات. أتناول تيجرتول وأرايبرزول وأولابكس. أرغب بالعلاج النفسي دون أدوية.
-
-فكر خطوة بخطوة:
-1. المستخدم لديه اضطراب وجداني منذ سنوات طويلة.
-2. لديه نوبات متكررة أعوام 2008، 2009، 2013، 2017، و2018.
-3. يأخذ عدة أدوية (تيجرتول وأرايبرزول وأولابكس).
-4. يريد الآن علاجاً غير دوائي، ربما علاجاً نفسياً بديلاً.
-
-السؤال المختصر:
-استفسار حول علاج الاضطراب الوجداني بدون دواء
 """.strip()
 
-    new_question_segment = f"""
+    new_segment = f"""
 الآن لديك سؤال جديد، فكر خطوة بخطوة بشكل مشابه ثم أعطني السؤال المختصر:
 تأكد من عدم اضافة ملاحظات أو اضافة معلومة غير موجودة في السؤال
 السؤال الأصلي:
-{question_text}
-
+{text}
 السؤال المختصر:
 """
 
-    prompt = f"""
+    full_prompt = f"""
 أنت مساعد لغوي مختص بأسئلة الصحة النفسية.
 هدفك هو إعادة كتابة الأسئلة المطوّلة في شكل مختصر ومباشر، يركّز على النقطة الأساسية.
 {examples}
-{new_question_segment}
+{new_segment}
 """.strip()
-    return prompt
-
-
-def summarize_question_cot_few_shots(question_text: str) -> str:
-    prompt = prompt_short_question_cot_few_shots(question_text)
-    raw_output = generate_text(prompt, max_new_tokens=60)
-
-    if "السؤال المختصر:" in raw_output:
-        short_summary = raw_output.split("السؤال المختصر:")[-1].strip()
-    else:
-        short_summary = raw_output.strip()
-
-    return short_summary
-
-# ------------------------------------------------------------------
-# 5. Gradio interface
-# ------------------------------------------------------------------
-def gradio_predict(question: str) -> str:
-    return summarize_question_cot_few_shots(question)
 
+    inputs = tokenizer(
+        tokenizer.apply_chat_template([{"role": "user", "content": full_prompt}], tokenize=False),
+        return_tensors="pt"
+    ).to(model.device)
 
-demo = gr.Interface(
-    fn=gradio_predict,
-    inputs=gr.Textbox(
-        lines=7,
-        label="السؤال الأصلي",
-        placeholder="اكتب سؤالك المطوّل هنا...",
-    ),
-    outputs=gr.Textbox(label="السؤال المختصر"),
-    title="🧠 Arabic Mental-Health Question Summarizer (ALLaM-7B)",
-    description=(
-        "يعيد هذا التطبيق صياغة الأسئلة العربية الطويلة حول الصحة النفسية إلى سؤال مختصر ومباشر."
-    ),
-    examples=[
-        "كنت أعاني من قلق وأخذت دواء سبرالكس لمدة 4 شهور وتوقفت عنه عام 2014..."
-    ],
-)
+    with torch.no_grad():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=False,
+            pad_token_id=tokenizer.eos_token_id
+        )
 
-if __name__ == "__main__":
-    demo.launch()
+    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
+    return decoded.split("السؤال المختصر:")[-1].strip() if "السؤال المختصر:" in decoded else decoded.strip()