explainability.py

import shap
import matplotlib.pyplot as plt

class ExplainabilityModule:
    """
    SHAP-based explainability module.
    """
    def __init__(self, model):
        self.model = model
        self.explainer = None
        
    def calculate_shap(self, X_train, X_instance):
        """
        img_path: Path to save the SHAP plot image.
        """
        # Create explainer
        # TreeExplainer is best for XGBoost
        self.explainer = shap.TreeExplainer(self.model)
        
        shap_values = self.explainer.shap_values(X_instance)
        
        return shap_values, self.explainer.expected_value

    def plot_shap_summary(self, X_train):
        """
        Generates SHAP summary plot figure.
        """
        if not self.explainer:
            self.explainer = shap.TreeExplainer(self.model)
            
        shap_values = self.explainer.shap_values(X_train)
        
        # Create a new figure and make it active
        plt.figure()
        shap.summary_plot(shap_values, X_train, show=False)
        plt.tight_layout()
        
        # Get the current figure
        fig = plt.gcf()
        
        # We don't close it here because we return it for Streamlit to render
        return fig