Update src/core_model/predict.py

src/core_model/predict.py

@@ -6,18 +6,18 @@ from transformers import XLMRobertaTokenizer, XLMRobertaForSequenceClassificatio
 class MindGuardPredictor:
     # The __init__ function runs automatically the moment you create a MindGuardPredictor object.
     def __init__(self, model_path=None):
-        print("🧠 Loading MindGuard Core Engine...")
+       self.model_id = "MohitRajput45/mindguard-xlmr"
         
-        # 1. THE FIX: Point to your new Model Hub ID
-        # Replace 'MohitRajput45/mindguard-xlmr' with your actual Model Repo name
-        self.model_id = "MohitRajput45/mindguard-xlmr"
-        
-        print(f"DEBUG: Loading model from Hugging Face Hub: {self.model_id}")
-
-        # 2. LOAD DIRECTLY FROM THE CLOUD
-        # Streamlit will cache this so it only downloads once!
-        self.tokenizer = XLMRobertaTokenizer.from_pretrained(self.model_id)
-        self.model = XLMRobertaForSequenceClassification.from_pretrained(self.model_id)
+        # 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"
+        )
         
         # --- THE FIX: The English Translation Dictionary ---
         # Paste the exact dictionary that printed in your Colab terminal here!