services/local_nlp_classifier.py

"""
Tier 2 Local NLP Classifier — Pre-trained Transformer Models.

This module provides near-LLM-quality enrichment using locally-downloaded
HuggingFace transformer models. It sits BETWEEN the HuggingFace API (Tier 1)
and the lexicon-based heuristics (Tier 3) in the fallback chain.

Models used:
  - Category:  facebook/bart-large-mnli (zero-shot classification, multilingual)
  - Sentiment: CAMeL-Lab/bert-base-arabic-camelbert-msa-sentiment (Arabic)
               distilbert-base-uncased-finetuned-sst-2-english (English)
  - NER:       CAMeL-Lab/bert-base-arabic-camelbert-msa-ner (Arabic)
               dslim/bert-base-NER (English)
  - Keywords:  KeyBERT with paraphrase-multilingual-MiniLM-L12-v2
"""

import os
import re
import logging
from typing import Optional

logger = logging.getLogger(__name__)

# Lazy-loaded singletons — models are heavy, we only load them once and only when needed.
_category_pipeline = None
_sentiment_ar_pipeline = None
_sentiment_en_pipeline = None
_ner_ar_pipeline = None
_ner_en_pipeline = None
_keyword_model = None
_models_available = None  # None = not checked yet, True/False after first attempt


# ---------------------------------------------------------------------------
# Our 13 categories (must match CATEGORIES_MAP keys in categories_lexicon.py)
# ---------------------------------------------------------------------------
CATEGORY_LABELS = [
    "Sports", "Technology", "Gaming", "Business", "Health", "Science",
    "Food", "Entertainment", "Politics", "Education", "Crime", "Weather", "Travel"
]

# ---------------------------------------------------------------------------
# Region inference from NER-detected locations
# ---------------------------------------------------------------------------
from services.lexicons import MIDDLE_EAST_COUNTRIES, LOCATIONS_MAP


def _check_dependencies() -> bool:
    """Check if transformers & torch are installed."""
    global _models_available
    if _models_available is not None:
        return _models_available
    try:
        import transformers  # noqa: F401
        import torch  # noqa: F401
        _models_available = True
        logger.info("NLP model dependencies (transformers, torch) are available.")
    except ImportError as e:
        _models_available = False
        logger.warning(f"NLP model dependencies not installed ({e}). Tier 2 fallback disabled.")
    return _models_available


# ========================= MODEL LOADERS =========================

def _get_category_pipeline():
    """Load custom fine-tuned category classifier from ./models/category_classifier if available."""
    global _category_pipeline
    if _category_pipeline is not None:
        return _category_pipeline
    try:
        from transformers import pipeline
        custom_model_path = "./models/category_classifier"
        if os.path.exists(os.path.join(custom_model_path, "model.safetensors")) or os.path.exists(os.path.join(custom_model_path, "pytorch_model.bin")):
            _category_pipeline = pipeline(
                "text-classification",
                model=custom_model_path,
                tokenizer=custom_model_path,
                device=-1
            )
            logger.info("Loaded CUSTOM fine-tuned category classifier from ./models/category_classifier")
        else:
            _category_pipeline = pipeline(
                "zero-shot-classification",
                model="facebook/bart-large-mnli",
                device=-1
            )
            logger.info("Loaded zero-shot category classifier (facebook/bart-large-mnli)")
    except Exception as e:
        logger.warning(f"Failed to load category classifier: {e}")
        _category_pipeline = None
    return _category_pipeline


def _get_sentiment_pipeline(language: str):
    """Language-specific sentiment pipeline with custom model support."""
    global _sentiment_ar_pipeline, _sentiment_en_pipeline

    custom_sent_dir = "./models/sentiment_analyzer"
    checkpoint_dir = os.path.join(custom_sent_dir, "checkpoint-5823")
    active_path = checkpoint_dir if os.path.exists(os.path.join(checkpoint_dir, "model.safetensors")) else custom_sent_dir

    if os.path.exists(os.path.join(active_path, "model.safetensors")) or os.path.exists(os.path.join(active_path, "pytorch_model.bin")):
        if _sentiment_ar_pipeline is not None:
            return _sentiment_ar_pipeline
        try:
            from transformers import pipeline
            pipe = pipeline(
                "text-classification",
                model=active_path,
                tokenizer="xlm-roberta-base",
                device=-1
            )
            logger.info(f"Loaded CUSTOM fine-tuned sentiment model from {active_path}")
            _sentiment_ar_pipeline = pipe
            _sentiment_en_pipeline = pipe
            return pipe
        except Exception as e:
            logger.warning(f"Failed loading custom sentiment model: {e}")

    if language == "ar":
        if _sentiment_ar_pipeline is not None:
            return _sentiment_ar_pipeline
        try:
            from transformers import pipeline
            _sentiment_ar_pipeline = pipeline(
                "text-classification",
                model="CAMeL-Lab/bert-base-arabic-camelbert-msa-sentiment",
                device=-1
            )
            logger.info("Loaded Arabic sentiment model (CAMeL-Lab/camelbert-msa-sentiment)")
        except Exception as e:
            logger.warning(f"Failed to load Arabic sentiment model: {e}")
            _sentiment_ar_pipeline = None
        return _sentiment_ar_pipeline
    else:
        if _sentiment_en_pipeline is not None:
            return _sentiment_en_pipeline
        try:
            from transformers import pipeline
            _sentiment_en_pipeline = pipeline(
                "text-classification",
                model="distilbert-base-uncased-finetuned-sst-2-english",
                device=-1
            )
            logger.info("Loaded English sentiment model (distilbert-sst-2)")
        except Exception as e:
            logger.warning(f"Failed to load English sentiment model: {e}")
            _sentiment_en_pipeline = None
        return _sentiment_en_pipeline


def _get_ner_pipeline(language: str):
    """Language-specific NER pipeline with custom model support."""
    global _ner_ar_pipeline, _ner_en_pipeline

    custom_ner_dir = "./models/ner_model"
    if os.path.exists(os.path.join(custom_ner_dir, "model.safetensors")) or os.path.exists(os.path.join(custom_ner_dir, "pytorch_model.bin")):
        if _ner_ar_pipeline is not None:
            return _ner_ar_pipeline
        try:
            from transformers import pipeline
            pipe = pipeline(
                "ner",
                model=custom_ner_dir,
                tokenizer=custom_ner_dir,
                aggregation_strategy="simple",
                device=-1
            )
            logger.info(f"Loaded CUSTOM fine-tuned NER model from {custom_ner_dir}")
            _ner_ar_pipeline = pipe
            _ner_en_pipeline = pipe
            return pipe
        except Exception as e:
            logger.warning(f"Failed to load custom NER model: {e}")

    if language == "ar":
        if _ner_ar_pipeline is not None:
            return _ner_ar_pipeline
        try:
            from transformers import pipeline
            _ner_ar_pipeline = pipeline(
                "ner",
                model="CAMeL-Lab/bert-base-arabic-camelbert-msa-ner",
                aggregation_strategy="simple",
                device=-1
            )
            logger.info("Loaded Arabic NER model (CAMeL-Lab/camelbert-msa-ner)")
        except Exception as e:
            logger.warning(f"Failed to load Arabic NER model: {e}")
            _ner_ar_pipeline = None
        return _ner_ar_pipeline
    else:
        if _ner_en_pipeline is not None:
            return _ner_en_pipeline
        try:
            from transformers import pipeline
            _ner_en_pipeline = pipeline(
                "ner",
                model="dslim/bert-base-NER",
                aggregation_strategy="simple",
                device=-1
            )
            logger.info("Loaded English NER model (dslim/bert-base-NER)")
        except Exception as e:
            logger.warning(f"Failed to load English NER model: {e}")
            _ner_en_pipeline = None
        return _ner_en_pipeline


def _get_keyword_model():
    """KeyBERT with multilingual sentence-transformer embeddings."""
    global _keyword_model
    if _keyword_model is not None:
        return _keyword_model
    try:
        from keybert import KeyBERT
        _keyword_model = KeyBERT(model="paraphrase-multilingual-MiniLM-L12-v2")
        logger.info("Loaded KeyBERT keyword extractor (multilingual-MiniLM)")
    except Exception as e:
        logger.warning(f"Failed to load KeyBERT: {e}")
        _keyword_model = None
    return _keyword_model


# ========================= ENRICHMENT =========================

class LocalNLPClassifier:
    """
    Tier 2 NLP enrichment engine using pre-trained transformer models.

    Usage:
        classifier = LocalNLPClassifier()
        if classifier.is_available():
            enriched_article = classifier.enrich_article(article)
    """

    def is_available(self) -> bool:
        """Check if the required ML libraries are installed."""
        return _check_dependencies()

    def enrich_article(self, article) -> "Article":
        """
        Enrich a single article using local transformer models.
        Each sub-task is wrapped in try/except so a failure in one
        (e.g. sentiment) does not block the others (e.g. NER).
        """
        text = f"{article.title} {article.summary or ''}"
        lang = getattr(article, 'language', 'en')

        # 1. Category (zero-shot classification)
        try:
            self._classify_category(article, text)
        except Exception as e:
            logger.warning(f"Tier 2 category classification failed: {e}")

        # 2. Sentiment
        try:
            self._classify_sentiment(article, text, lang)
        except Exception as e:
            logger.warning(f"Tier 2 sentiment analysis failed: {e}")

        # 3. Named Entity Recognition -> people, organizations, locations, country, region
        try:
            self._extract_entities(article, text, lang)
        except Exception as e:
            logger.warning(f"Tier 2 NER failed: {e}")

        # 4. Keywords & Tags
        try:
            self._extract_keywords(article, text)
        except Exception as e:
            logger.warning(f"Tier 2 keyword extraction failed: {e}")

        # 5. Importance scoring (derived from NER entity density + category)
        try:
            self._score_importance(article, text)
        except Exception as e:
            logger.warning(f"Tier 2 importance scoring failed: {e}")

        # 6. Reading time estimation
        if not article.readingTime:
            article.readingTime = self._estimate_reading_time(text)

        return article

    # ---------------------------------------------------------------
    # Sub-task implementations
    # ---------------------------------------------------------------

    def _classify_category(self, article, text: str):
        """Classify into our 13 news categories (custom model or zero-shot)."""
        if article.category and article.category not in ("General", "Middle East", "Egypt", "World", "North Africa"):
            return  # Already classified by the RSS source metadata

        pipe = _get_category_pipeline()
        if pipe is None:
            return

        if pipe.task == "text-classification":
            result = pipe(text[:512])
            if result:
                label = result[0]['label']
                article.category = label
        else:
            result = pipe(text[:1024], candidate_labels=CATEGORY_LABELS, multi_label=False)
            if result and result['scores'][0] > 0.25:
                article.category = result['labels'][0]
                if len(result['labels']) > 1 and result['scores'][1] > 0.20:
                    article.subcategory = result['labels'][1]
            else:
                article.category = article.category or "General"

    def _classify_sentiment(self, article, text: str, lang: str):
        """Sentiment analysis using language-specific pre-trained models."""
        pipe = _get_sentiment_pipeline(lang)
        if pipe is None:
            return

        result = pipe(text[:512])
        if not result:
            return

        label = result[0]['label'].upper()
        score = result[0]['score']

        if lang == "ar":
            # CAMeLBERT sentiment outputs: positive, negative, neutral
            label_map = {"POSITIVE": "Positive", "NEGATIVE": "Negative", "NEUTRAL": "Neutral"}
            article.sentiment = label_map.get(label, "Neutral")
        else:
            # DistilBERT SST-2 outputs: POSITIVE, NEGATIVE
            if label == "POSITIVE":
                article.sentiment = "Positive"
            elif label == "NEGATIVE":
                article.sentiment = "Negative"
            else:
                article.sentiment = "Neutral"

            # Mixed sentiment: if confidence is low (close to 50/50)
            if 0.45 < score < 0.60:
                article.sentiment = "Mixed"

    def _extract_entities(self, article, text: str, lang: str):
        """NER extraction: people, organizations, locations."""
        pipe = _get_ner_pipeline(lang)
        if pipe is None:
            return

        entities = pipe(text[:512])
        if not entities:
            return

        people = []
        organizations = []
        locations = []

        for ent in entities:
            entity_group = ent.get('entity_group', '')
            word = ent.get('word', '').strip()
            score = ent.get('score', 0)

            # Filter low-confidence entities
            if score < 0.5 or len(word) < 2:
                continue

            # Clean up subword tokens (##prefix)
            word = word.replace('##', '')

            if entity_group in ('PER', 'B-PER', 'I-PER'):
                if word not in people:
                    people.append(word)
            elif entity_group in ('ORG', 'B-ORG', 'I-ORG'):
                if word not in organizations:
                    organizations.append(word)
            elif entity_group in ('LOC', 'B-LOC', 'I-LOC', 'GPE', 'B-GPE', 'I-GPE'):
                if word not in locations:
                    locations.append(word)

        article.people = people[:10]       # Cap at 10
        article.organizations = organizations[:10]
        article.locations = locations[:10] if locations else article.locations

        # Derive country and region from detected locations
        if locations:
            # Try to map NER location names to our gazetteer
            for loc in locations:
                loc_lower = loc.lower()
                if loc_lower in LOCATIONS_MAP:
                    article.country = LOCATIONS_MAP[loc_lower]
                    article.region = "Middle East" if article.country in MIDDLE_EAST_COUNTRIES else "Global"
                    break
            if not article.country:
                article.country = locations[0]
                article.region = "Global"

    def _extract_keywords(self, article, text: str):
        """KeyBERT keyword extraction with multilingual embeddings."""
        kw_model = _get_keyword_model()
        if kw_model is None:
            return

        try:
            keywords = kw_model.extract_keywords(
                text,
                keyphrase_ngram_range=(1, 2),
                stop_words=None,  # KeyBERT handles multilingual stop words
                top_n=12,
                use_maxsum=True,
                nr_candidates=20
            )
        except Exception:
            return

        if keywords:
            article.keywords = [kw[0] for kw in keywords[:6]]
            article.tags = [kw[0] for kw in keywords[:10]]

    def _score_importance(self, article, text: str):
        """
        Importance scoring (1-10) based on:
        - Number of named entities (more entities = more important)
        - High-impact trigger words
        - Category weight
        """
        score = 5  # Base importance

        # Named entity density bonus
        entity_count = len(article.people or []) + len(article.organizations or []) + len(article.locations or [])
        if entity_count >= 5:
            score += 2
        elif entity_count >= 2:
            score += 1

        # High-impact trigger words
        high_impact_terms = [
            "عاجل", "قمة", "حرب", "وفاة", "رئيس", "كوارث", "اغتيال", "انقلاب", "زلزال",
            "breaking", "war", "president", "disaster", "crisis", "urgent", "killed",
            "earthquake", "assassination", "coup", "summit", "death"
        ]
        text_lower = text.lower()
        if any(term in text_lower for term in high_impact_terms):
            score += 3

        # Category bonus
        high_importance_categories = {"Politics", "Crime", "Health", "Science"}
        if article.category in high_importance_categories:
            score += 1

        article.importance = min(10, max(1, score))

    def _estimate_reading_time(self, text: str) -> int:
        """Estimate reading time in minutes from word count."""
        clean_text = re.sub(r'<[^>]+>', '', text)
        words = len(clean_text.split())
        return max(1, round(words / 180))