Update src/core_model/predict.py

src/core_model/predict.py

@@ -4,20 +4,26 @@ import torch.nn.functional as F
 from transformers import XLMRobertaTokenizer, XLMRobertaForSequenceClassification
 
 class MindGuardPredictor:
-    # The __init__ function runs automatically the moment you create a MindGuardPredictor object.
-    def __init__(self, model_path=None):
-       self.model_id = "MohitRajput45/mindguard-xlmr"
+    # 1. Identify the Model Hub ID
+        self.model_id = "MohitRajput45/mindguard-xlmr"
         
-        # If your weights are inside 'final_mindguard_model' folder on HF Hub:
-        # Use the subfolder parameter
-        self.tokenizer = XLMRobertaTokenizer.from_pretrained(
-            self.model_id, 
-            subfolder="final_mindguard_model"
-        )
-        self.model = XLMRobertaForSequenceClassification.from_pretrained(
-            self.model_id, 
-            subfolder="final_mindguard_model"
-        )
+        # 2. Load the components
+        # Note: If your files are in the root of the repo, remove subfolder="final_mindguard_model"
+        try:
+            self.tokenizer = XLMRobertaTokenizer.from_pretrained(
+                self.model_id, 
+                subfolder="final_mindguard_model"
+            )
+            self.model = XLMRobertaForSequenceClassification.from_pretrained(
+                self.model_id, 
+                subfolder="final_mindguard_model"
+            )
+            print("✅ Model and Tokenizer loaded successfully from Hub.")
+        except Exception as e:
+            print(f"❌ Error loading model: {e}")
+            # Fallback attempt without subfolder if the above fails
+            self.tokenizer = XLMRobertaTokenizer.from_pretrained(self.model_id)
+            self.model = XLMRobertaForSequenceClassification.from_pretrained(self.model_id)
         
         # --- THE FIX: The English Translation Dictionary ---
         # Paste the exact dictionary that printed in your Colab terminal here!