models/space_er.py

"""
Arabic Mental Health Named Entity Recognition Module
Model: GLiNER (fine-tuned for Arabic mental health by AhmadDarif)

This module loads the GLiNER model and exposes a function to extract named entities
related to mental health from Arabic input text.

Entity labels include:
- MEDICATION: أسماء الأدوية
- DOSAGE: الجرعات
- DURATION: مدة الاستخدام

Dependencies:
- gliner
- torch
"""

from gliner import GLiNER

# Load the model
print("🚀 Loading GLiNER model for Arabic mental health NER...")
try:
    model = GLiNER.from_pretrained("AhmadDarif/Arabic-Mental-NER")
    print("✅ Loaded fine-tuned model: AhmadDarif/Arabic-Mental-NER")
except Exception as e:
    print(f"⚠️ Could not load fine-tuned model. Falling back. Error: {str(e)}")
    model = GLiNER.from_pretrained("urchade/gliner_multi-v2.1")
    print("✅ Loaded fallback model: urchade/gliner_multi-v2.1")

# Entity labels to extract
LABELS = ["MEDICATION", "DOSAGE", "DURATION"]

# Mapping from English to Arabic labels
label_map = {
    "MEDICATION": "اسم الدواء",
    "DOSAGE": "الجرعة",
    "DURATION": "مدة الاستخدام"
}

def extract_entities(text: str):
    """
    Extract named entities from Arabic text using GLiNER.

    Args:
        text (str): Input Arabic text.

    Returns:
        list[list[str]]: Table of extracted entities [text, arabic_label] or empty list.
    """
    if not text.strip():
        return []

    try:
        entities = model.predict_entities(text, LABELS)
        if not entities:
            return []

        # Return results with Arabic labels
        return [[ent["text"], label_map.get(ent["label"], ent["label"])] for ent in entities]

    except Exception:
        return []  # ✅ لا نرجع رسالة نصية بل نرجع قائمة فاضية فقط