Update src/explainability/shap_explainer.py

src/explainability/shap_explainer.py

@@ -14,21 +14,20 @@ class MindGuardSHAPExplainer:
     def __init__(self):
         print("🔍 Initializing MindGuard SHAP Explainability Engine...")
         
-        # --- STRICT ARCHITECTURE PATHING ---
-        # 1. Locate the current script (src/explainability/shap_explainer.py)
-        self.script_dir = os.path.dirname(os.path.abspath(__file__))
+        # --- HUB ARCHITECTURE PATHING ---
+        # 1. Point to your Model Hub ID (NOT a local path)
+        self.model_id = "MohitRajput45/mindguard-xlmr"
         
-        # 2. Traverse up TWO directories to hit the MINDGUARD_AI_PROJECT root
-        self.project_root = os.path.abspath(os.path.join(self.script_dir, "../../"))
+        # 2. Use the exact path to where the weights live inside that Hub repo
+        self.hf_subfolder = "artifacts/xlmr_weights/final_mindguard_model"
         
-        # 3. Define the exact path to where we saved the trained brain in Day 3
-        self.model_path = os.path.join(self.project_root, "artifacts", "xlmr_weights", "final_mindguard_model")
-        
-        # 4. Define where the visual HTML reports will be saved
-        self.artifacts_dir = os.path.join(self.project_root, "artifacts")
-        
-        # --- THE FIX 1: The Translation Dictionary ---
-        # We must define the 35 English emotions so SHAP doesn't output "LABEL_X"
+        # 3. Define where the visual HTML reports will be saved (relative to root)
+        # On Hugging Face, /app is the root.
+        self.artifacts_dir = os.path.join(os.getcwd(), "artifacts")
+        if not os.path.exists(self.artifacts_dir):
+            os.makedirs(self.artifacts_dir)
+
+        # --- THE TRANSLATION DICTIONARY ---
         self.emotion_map = {
             0: 'Anxiety', 1: 'Bipolar', 2: 'Depression', 3: 'Normal', 
             4: 'Personality disorder', 5: 'Stress', 6: 'Suicidal', 7: 'admiration', 
@@ -41,73 +40,72 @@ class MindGuardSHAPExplainer:
             32: 'remorse', 33: 'sadness', 34: 'surprise'
         }
 
-        print(f"Loading Core Brain from: {self.model_path}...")
+        print(f"Loading Core Brain from Hub: {self.model_id}...")
         
-        # --- LOAD THE AI CORE ---
-        # Initialize the tokenizer (translates text to numbers) and the model (the brain)
-        self.tokenizer = XLMRobertaTokenizer.from_pretrained(self.model_path)
-        self.model = XLMRobertaForSequenceClassification.from_pretrained(self.model_path)
-        
-        # --- THE FIX 2: Inject the Dictionary into the Model's Brain ---
-        # This permanently forces the Hugging Face model to speak English instead of Math
-        self.model.config.id2label = self.emotion_map
-        self.model.config.label2id = {v: k for k, v in self.emotion_map.items()}
+        # --- LOAD THE AI CORE FROM HUB ---
+        try:
+            self.tokenizer = XLMRobertaTokenizer.from_pretrained(
+                self.model_id, 
+                subfolder=self.hf_subfolder
+            )
+            self.model = XLMRobertaForSequenceClassification.from_pretrained(
+                self.model_id, 
+                subfolder=self.hf_subfolder
+            )
+            
+            # Inject the Dictionary into the Model's Brain
+            self.model.config.id2label = self.emotion_map
+            self.model.config.label2id = {v: k for k, v in self.emotion_map.items()}
 
-        # Use Hugging Face's pipeline to wrap the model for easy SHAP integration
-        # Set device to 0 if a GPU is detected, otherwise fallback to CPU (-1)
-        self.device = 0 if torch.cuda.is_available() else -1
-        self.classifier = pipeline(
-            "text-classification", 
-            model=self.model, 
-            tokenizer=self.tokenizer, 
-            device=self.device, 
-            top_k=None # top_k=None forces the AI to output scores for ALL 35 emotions, not just the top guess
-        )
-        
-        # --- WARM UP SHAP ---
-        print("⚙️ Warming up Game Theory Math (SHAP)...")
-        # Pass our classifier pipeline into the SHAP Explainer engine
-        self.explainer = shap.Explainer(self.classifier)
-        print("✅ SHAP Explainer ready!")
+            # Set device: 0 for GPU, -1 for CPU
+            self.device = 0 if torch.cuda.is_available() else -1
+            
+            # Hugging Face pipeline for SHAP integration
+            self.classifier = pipeline(
+                "text-classification", 
+                model=self.model, 
+                tokenizer=self.tokenizer, 
+                device=self.device, 
+                top_k=None # Get scores for all 35 emotions
+            )
+            
+            print("⚙️ Warming up Game Theory Math (SHAP)...")
+            self.explainer = shap.Explainer(self.classifier)
+            print("✅ SHAP Explainer ready!")
+            
+        except Exception as e:
+            print(f"❌ SHAP Initialization failed: {e}")
+            raise e
 
     def generate_visual_report(self, text):
         """
-        Takes a raw string of text, runs it through the model, 
-        calculates SHAP values, and outputs an interactive HTML file.
+        Takes raw text, calculates SHAP values, and outputs an interactive HTML file.
         """
         print(f"\n🧠 Analyzing: '{text}'")
         
         # 1. Run the Game Theory calculations
-        # This isolates the impact of every single word on the final prediction
         shap_values = self.explainer([text])
         
-        # 2. Define the exact save location for the HTML report
+        # 2. Define save location
         html_path = os.path.join(self.artifacts_dir, "shap_report.html")
         
-        # --- THE FIX 3: Targeted Slicing ---
-        # Instead of drawing 35 overlapping arrows, find the emotion the AI was MOST confident in.
+        # 3. Targeted Slicing: Find the most confident emotion to visualize
+        # shap_values[0].values is (tokens, 35). We sum across tokens to find the top class.
         best_class_index = shap_values[0].values.sum(axis=0).argmax()
         
-        # 3. Generate the visualization ONLY for the winning emotion
-        # display=False ensures it generates the raw HTML instead of trying to open a Jupyter widget
+        # 4. Generate visualization for the winning emotion
+        # Note: best_class_index must be passed as the index for the 3rd dimension
         shap_html = shap.plots.text(shap_values[0, :, best_class_index], display=False)
         
-        # 4. Save the HTML string to a physical file in the artifacts folder
+        # 5. Save the HTML
         with open(html_path, "w", encoding="utf-8") as f:
             f.write(shap_html)
             
-        print(f"✅ Diagnostic Complete!")
-        print(f"Visual Report saved to: {html_path}")
-        print("Go to your 'artifacts' folder and open 'shap_report.html' in your browser.")
+        print(f"✅ Diagnostic Complete! Visual Report saved to: {html_path}")
+        return html_path
 
 # --- EXECUTION BLOCK ---
-# This block only runs if this specific file is executed directly from the terminal
 if __name__ == "__main__":
-    # Instantiate our explainer class
     explainer = MindGuardSHAPExplainer()
-    
-    # Define a test patient input
     sample_text = "I have a massive presentation tomorrow and my chest is tight."
-    
-    # Generate the explanation report
     explainer.generate_visual_report(sample_text)