core/explainer.py

import re
from typing import Any, Dict, List, Tuple

import numpy as np
import torch
from transformers import DistilBertForSequenceClassification, DistilBertTokenizer


class SHAPExplainer:
    def __init__(self, model_path: str = "."):
        print("🚀 Initializing SHAP Explainer...")
        # Import shap lazily to avoid triggering optional plotting imports
        # during app startup (e.g., Spaces reload scanner).
        try:
            import shap  # type: ignore

            self._shap = shap
        except Exception as e:
            print(f"⚠️ SHAP import failed; falling back to heuristic explainer: {e}")
            self._shap = None

        self.model = DistilBertForSequenceClassification.from_pretrained(model_path)
        self.tokenizer = DistilBertTokenizer.from_pretrained(model_path)

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device).eval()
        print(f"✅ Model loaded on: {self.device}")

        def model_predict(texts):
            if isinstance(texts, str):
                texts = [texts]
            elif hasattr(texts, "__iter__") and not isinstance(texts, str):
                texts = list(texts)

            texts = [str(t) for t in texts]

            inputs = self.tokenizer(
                texts,
                padding=True,
                truncation=True,
                max_length=256,
                return_tensors="pt",
            )
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            with torch.no_grad():
                outputs = self.model(**inputs)
                probs = torch.softmax(outputs.logits, dim=-1)
                return probs[:, 1].cpu().numpy()

        self.model_predict = model_predict

        if self._shap is None:
            self.explainer = None
            return

        try:
            masker = self._shap.maskers.Text(tokenizer=self.tokenizer)
            self.explainer = self._shap.Explainer(self.model_predict, masker=masker)
            print("✅ SHAP Explainer initialized successfully!")
        except Exception as e:
            print(f"❌ Error initializing SHAP: {e}")
            self.explainer = None

    def _get_bias_probability(self, text: str) -> float:
        """Get the model's bias probability for text (class 1)."""
        inputs = self.tokenizer(
            text, return_tensors="pt", truncation=True, max_length=256
        )
        inputs = {k: v.to(self.device) for k, v in inputs.items()}

        with torch.no_grad():
            outputs = self.model(**inputs)
            probs = torch.softmax(outputs.logits, dim=-1)
            return probs[0, 1].item()

    def get_shap_values(
        self, text: str, max_evals: int = 500
    ) -> List[Tuple[str, float]]:
        """Compute SHAP contributions for words in `text`.

        If the SHAP explainer could not be created during init this will run a
        lightweight keyword-based fallback. The returned list contains tuples
        (word, score) sorted by absolute impact.

        Args:
            text: Input text to explain.
            max_evals: Maximum number of SHAP evaluations/samples.

        Returns:
            List of (word, score) tuples ordered by descending absolute impact.
        """
        if self.explainer is None:
            return self._fallback_analysis(text)

        try:
            shap_values = self.explainer([text], max_evals=max_evals)

            # DistilBERT SHAP output may be (1, tokens, outputs) or (1, tokens).
            if len(shap_values.values.shape) == 3:
                biased_values = shap_values.values[0, :, 0]
            else:
                biased_values = shap_values.values[0, :]

            tokens = shap_values.data[0]

            # Combine subword tokens (e.g., BERT-style `##` tokens) into full words
            combined_scores = self._combine_subword_scores(tokens, biased_values)
            # Filter out special tokens and very short tokens
            filtered_scores = [
                (w, float(s))
                for w, s in combined_scores
                if w.strip() and w not in ["[cls]", "[sep]", "[pad]"] and len(w) > 1
            ]
            filtered_scores.sort(key=lambda x: abs(x[1]), reverse=True)
            return filtered_scores[:10]

        except Exception as e:
            print(f"❌ SHAP calculation failed: {e}")
            return self._fallback_analysis(text)

    def _combine_subword_scores(
        self, tokens: List[str], scores: np.ndarray
    ) -> List[Tuple[str, float]]:
        """Merge BPE/subword tokens into human-readable words with averaged scores.

        Many tokenizers split words into subwords prefixed with '##'. This
        function recombines those tokens and averages their SHAP scores so the
        UI shows per-word importance instead of per-subword noise.

        Args:
            tokens: List of token strings from the tokenizer/SHAP.
            scores: Array of SHAP scores aligned with `tokens`.

        Returns:
            List of (word, averaged_score) tuples.
        """
        combined = []
        current_word, current_score, token_count = "", 0.0, 0

        for token, score in zip(tokens, scores):
            # If token is a continuation subword (BERT-style) start with '##'
            if token.startswith("##"):
                # append the subword (without the prefix) to the current word
                current_word += token[2:]
                current_score += score
                token_count += 1
            else:
                # push the previously accumulated word if present
                if current_word:
                    combined.append((current_word, current_score / token_count))
                current_word, current_score, token_count = token, score, 1

        if current_word:
            combined.append((current_word, current_score / token_count))

        # Final cleanup: remove any lingering '##' markers and strip whitespace
        return [(word.replace("##", "").strip(), score) for word, score in combined]

    def _fallback_analysis(self, text: str) -> List[Tuple[str, float]]:
        """Fallback analysis when SHAP fails.

        Very rough heuristic: use token length + presence of sensitive words.
        """
        bias_prob = self._get_bias_probability(text)

        inputs = self.tokenizer(
            text, return_tensors="pt", truncation=True, max_length=256
        )
        tokens = self.tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
        word_scores: List[Tuple[str, float]] = []

        for token in tokens:
            if token in ["[cls]", "[sep]", "[pad]"]:
                continue
            clean_token = token.replace("##", "").strip()
            if len(clean_token) <= 1:
                continue

            score = len(clean_token) * 0.1
            if clean_token.lower() in [
                "women",
                "men",
                "female",
                "male",
                "she",
                "he",
            ]:
                score += 0.5
            if clean_token.lower() in [
                "nurse",
                "engineer",
                "teacher",
                "secretary",
            ]:
                score += 0.3
            if clean_token.lower() in [
                "emotional",
                "aggressive",
                "decisive",
                "leadership",
            ]:
                score += 0.2

            word_scores.append((clean_token, score * bias_prob))

        word_scores.sort(key=lambda x: x[1], reverse=True)
        return word_scores[:8]

    def classify_bias(
        self,
        text: str,
        bias_prob: float,
        shap_results: List[Tuple[str, float]],
    ) -> Tuple[str, List[str]]:
        """Rule-based bias type tagging + bias label banding."""
        text_lower = text.lower()
        bias_types: List[str] = []
        rule_score = 0.0

        group_words = ["women", "men", "woman", "man", "she", "he", "female", "male"]
        role_words = [
            "nurse",
            "engineer",
            "secretary",
            "teacher",
            "leader",
            "leadership",
            "manager",
        ]
        aggression_phrases = ["more aggressive", "be more aggressive"]

        has_group_word = any(g in text_lower for g in group_words)

        if has_group_word and any(r in text_lower for r in role_words):
            if any(p in text_lower for p in ["should be", "supposed to be", "meant to be"]):
                bias_types.append("gender_role_stereotype")
                rule_score = max(rule_score, 0.8)

        if has_group_word and any(
            p in text_lower
            for p in ["are better at", "naturally better", "inherently better", "better at"]
        ):
            bias_types.append("gender_superiority")
            rule_score = max(rule_score, 0.8)

        if "too emotional" in text_lower or "less emotional" in text_lower:
            bias_types.append("emotionality_judgment")
            rule_score = max(rule_score, 0.7)
        elif "emotional" in text_lower and has_group_word:
            bias_types.append("emotionality_judgment")
            rule_score = max(rule_score, 0.6)

        if any(phrase in text_lower for phrase in aggression_phrases):
            if "to succeed" in text_lower or "for success" in text_lower:
                bias_types.append("aggression_norm")
                rule_score = max(rule_score, 0.7)

        skill_based = (
            "people with" in text_lower
            and "skills" in text_lower
            and any(f in text_lower for f in ["role", "roles", "position", "positions", "job", "jobs"])
        )

        if (not has_group_word) and skill_based and not bias_types:
            effective_bias_prob = min(bias_prob, 0.2)
        else:
            effective_bias_prob = max(bias_prob, rule_score)

        if effective_bias_prob >= 0.7:
            label = "BIASED"
        elif effective_bias_prob >= 0.3:
            label = "POSSIBLY_BIASED"
        else:
            label = "LIKELY_NEUTRAL"

        return label, sorted(list(set(bias_types)))