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"]

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

    Args:
        text (str): Input Arabic text.

    Returns:
        str: Formatted entity extraction result or error message.
    """
    if not text.strip():
        return "يرجى إدخال نص للتحليل / Please enter text to analyze"

    try:
        entities = model.predict_entities(text, LABELS)
        if not entities:
            return "لم يتم العثور على أي كيانات / No entities found"

        result = "الكيانات المستخرجة:\n\n"
        for ent in entities:
            result += f"{ent['text']} => {ent['label']} (confidence: {ent.get('score', 0):.3f})\n"
        return result

    except Exception as e:
        return f"خطأ في التحليل / Analysis Error: {str(e)}"