xai/lime_explainer.py

"""LIME explainer for Myanmar Ghost model.

Uses LIME (Local Interpretable Model-agnostic Explanations)
to explain individual predictions.
"""

import logging
from typing import Any, Callable, Dict, List, Optional, Tuple

import numpy as np
import torch
from lime.lime_text import LimeTextExplainer

logger = logging.getLogger(__name__)


class ThankingLIMEExplainer:
    """LIME-based explainer for Myanmar text classification."""
    
    def __init__(
        self,
        model,
        tokenizer,
        class_names: Optional[List[str]] = None,
        kernel_width: float = 25.0,
    ):
        """
        Args:
            model: PyTorch model
            tokenizer: Tokenizer
            class_names: Names for output classes
            kernel_width: Kernel width for LIME
        """
        self.model = model
        self.tokenizer = tokenizer
        self.class_names = class_names or [
            "negative", "neutral", "positive", "sarcastic"
        ]
        
        self.model.eval()
        
        # Create LIME explainer
        self.explainer = LimeTextExplainer(
            class_names=self.class_names,
            kernel_width=kernel_width,
            verbose=False,
        )
    
    def _predict_proba(self, texts: List[str]) -> np.ndarray:
        """Prediction function for LIME.
        
        Args:
            texts: List of text strings
        
        Returns:
            Probability array (n_samples, n_classes)
        """
        # Tokenize
        encoding = self.tokenizer(
            texts,
            padding=True,
            truncation=True,
            max_length=128,
            return_tensors="pt",
        )
        
        input_ids = encoding["input_ids"]
        attention_mask = encoding["attention_mask"]
        
        with torch.no_grad():
            outputs = self.model(input_ids, attention_mask)
            
            if hasattr(outputs, "logits"):
                logits = outputs.logits
            else:
                logits = outputs
            
            probs = torch.softmax(logits, dim=-1)
        
        return probs.cpu().numpy()
    
    def explain(
        self,
        text: str,
        num_features: int = 10,
        num_samples: int = 5000,
        top_labels: int = 4,
    ) -> Any:
        """Explain a single text prediction.
        
        Args:
            text: Myanmar text to explain
            num_features: Number of features to show
            num_samples: Number of samples for LIME
            top_labels: Number of top labels to explain
        
        Returns:
            LIME Explanation object
        """
        logger.info(f"Explaining with LIME: {text[:50]}...")
        
        explanation = self.explainer.explain_instance(
            text,
            self._predict_proba,
            num_features=num_features,
            num_samples=num_samples,
            top_labels=top_labels,
        )
        
        return explanation
    
    def get_word_importance(
        self,
        text: str,
        class_index: Optional[int] = None,
        num_features: int = 10,
    ) -> List[Tuple[str, float]]:
        """Get word importance for a specific class.
        
        Args:
            text: Myanmar text
            class_index: Class index (None = predicted class)
            num_features: Number of top features
        
        Returns:
            List of (word, importance) tuples
        """
        explanation = self.explain(text, num_features=num_features)
        
        if class_index is None:
            # Use predicted class
            class_index = explanation.available_labels()[0]
        
        exp_list = explanation.as_list(label=class_index)
        
        return exp_list
    
    def visualize(
        self,
        explanation: Any,
        output_path: Optional[str] = None,
    ) -> str:
        """Generate text visualization of explanation.
        
        Args:
            explanation: LIME Explanation object
            output_path: Optional path to save
        
        Returns:
            Visualization text
        """
        output = "\n" + "=" * 60 + "\n"
        output += "LIME EXPLANATION\n"
        output += "=" * 60 + "\n"
        
        for label in explanation.available_labels()[:3]:
            label_name = self.class_names[label]
            output += f"\n{label_name.upper()}:\n"
            output += "-" * 40 + "\n"
            
            for word, weight in explanation.as_list(label=label):
                sign = "+" if weight > 0 else ""
                output += f"  {word}: {sign}{weight:.4f}\n"
        
        print(output)
        
        if output_path:
            with open(output_path, "w", encoding="utf-8") as f:
                f.write(output)
            logger.info(f"Visualization saved to {output_path}")
        
        return output
    
    def batch_explain(
        self,
        texts: List[str],
        num_features: int = 10,
    ) -> List[Dict[str, Any]]:
        """Explain multiple texts.
        
        Args:
            texts: List of texts
            num_features: Number of features per explanation
        
        Returns:
            List of explanation dictionaries
        """
        results = []
        
        for text in texts:
            explanation = self.explain(text, num_features=num_features)
            
            result = {
                "text": text,
                "predicted_class": self.class_names[
                    explanation.available_labels()[0]
                ],
                "explanations": {},
            }
            
            for label in explanation.available_labels():
                result["explanations"][self.class_names[label]] = {
                    word: weight
                    for word, weight in explanation.as_list(label=label)
                }
            
            results.append(result)
        
        return results


class SegmentLevelLIME:
    """LIME with Myanmar-specific segmentation."""
    
    def __init__(
        self,
        model,
        tokenizer,
        segment_syllables: bool = True,
    ):
        """
        Args:
            model: PyTorch model
            tokenizer: Tokenizer
            segment_syllables: Segment by Myanmar syllables
        """
        self.model = model
        self.tokenizer = tokenizer
        self.segment_syllables = segment_syllables
    
    def _segment_text(self, text: str) -> List[str]:
        """Segment text into interpretable units.
        
        For Myanmar, this can be syllable or word level.
        """
        if self.segment_syllables:
            # Simple syllable segmentation
            # Myanmar syllables end with vowel markers or consonants
            segments = []
            current = ""
            
            for char in text:
                current += char
                # Check for syllable boundary (simplified)
                if char in "း့်ှင်း":
                    segments.append(current)
                    current = ""
            
            if current:
                segments.append(current)
            
            return segments if segments else [text]
        else:
            return text.split()
    
    def _join_segments(self, segments: List[str]) -> str:
        """Join segments back to text."""
        return "".join(segments)
    
    def explain(
        self,
        text: str,
        num_samples: int = 1000,
    ) -> Dict[str, Any]:
        """Explain text with syllable-level segmentation.
        
        Args:
            text: Myanmar text
            num_samples: Number of LIME samples
        
        Returns:
            Explanation dictionary
        """
        segments = self._segment_text(text)
        
        logger.info(f"Segmented into {len(segments)} units")
        
        # Use standard LIME with segmented text
        base_explainer = ThankingLIMEExplainer(
            self.model,
            self.tokenizer,
        )
        
        explanation = base_explainer.explain(
            text,
            num_features=len(segments),
            num_samples=num_samples,
        )
        
        return {
            "text": text,
            "segments": segments,
            "explanation": explanation,
            "word_importance": base_explainer.get_word_importance(text),
        }


def create_lime_explainer(
    model,
    tokenizer,
    class_names: Optional[List[str]] = None,
) -> ThankingLIMEExplainer:
    """Factory function to create LIME explainer."""
    return ThankingLIMEExplainer(
        model=model,
        tokenizer=tokenizer,
        class_names=class_names,
    )


if __name__ == "__main__":
    print("ThankingLIMEExplainer loaded")
    print("Use create_lime_explainer() to create an explainer")