Spaces:
Running
Running
| # src/core_model/predict.py | |
| import os | |
| import torch | |
| import torch.nn.functional as F | |
| from transformers import XLMRobertaTokenizer, XLMRobertaForSequenceClassification | |
| class MindGuardPredictor: | |
| def __init__(self): | |
| # 1. Identify the Model Hub ID | |
| self.model_id = "MohitRajput45/mindguard-xlmr" | |
| # 2. Identify the EXACT path inside that Repo | |
| # Based on your HF screenshot, this is the full path to the files | |
| hf_subfolder = "artifacts/xlmr_weights/final_mindguard_model" | |
| # 3. Load the components from the Hub | |
| try: | |
| print(f"DEBUG: Attempting to load model from {self.model_id} in {hf_subfolder}...") | |
| self.tokenizer = XLMRobertaTokenizer.from_pretrained( | |
| self.model_id, | |
| subfolder=hf_subfolder | |
| ) | |
| self.model = XLMRobertaForSequenceClassification.from_pretrained( | |
| self.model_id, | |
| subfolder=hf_subfolder | |
| ) | |
| print("✅ Model and Tokenizer loaded successfully from Hub!") | |
| except Exception as e: | |
| print(f"❌ Error loading model: {e}") | |
| raise e # Stop the app if the model fails to load | |
| # 3. Emotion Mapping (Mathematical ID to English Word) | |
| self.emotion_map = { | |
| 0: 'Anxiety', 1: 'Bipolar', 2: 'Depression', 3: 'Normal', | |
| 4: 'Personality disorder', 5: 'Stress', 6: 'Suicidal', 7: 'admiration', | |
| 8: 'amusement', 9: 'anger', 10: 'annoyance', 11: 'approval', | |
| 12: 'caring', 13: 'confusion', 14: 'curiosity', 15: 'desire', | |
| 16: 'disappointment', 17: 'disapproval', 18: 'disgust', 19: 'embarrassment', | |
| 20: 'excitement', 21: 'fear', 22: 'gratitude', 23: 'grief', | |
| 24: 'joy', 25: 'love', 26: 'nervousness', 27: 'neutral', | |
| 28: 'optimism', 29: 'pride', 30: 'realization', 31: 'relief', | |
| 32: 'remorse', 33: 'sadness', 34: 'surprise' | |
| } | |
| # 4. Device Setup (GPU vs CPU) | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| self.model.to(self.device) | |
| self.model.eval() | |
| def determine_risk_level(self, emotion): | |
| """Categorizes emotions into clinical risk buckets.""" | |
| emotion = emotion.lower() | |
| high_risk = ['panic', 'severe anxiety', 'depression', 'grief', 'suicidal', 'personality disorder'] | |
| medium_risk = ['stress', 'anxiety', 'anger', 'burnout', 'fear', 'nervousness'] | |
| if emotion in high_risk: | |
| return "High" | |
| elif emotion in medium_risk: | |
| return "Medium" | |
| else: | |
| return "Low" | |
| def predict(self, text): | |
| """The core engine: Text -> Tensor -> Prediction.""" | |
| # Tokenize input | |
| inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=128, padding=True) | |
| inputs = {key: val.to(self.device) for key, val in inputs.items()} | |
| # Run inference | |
| with torch.no_grad(): | |
| outputs = self.model(**inputs) | |
| logits = outputs.logits | |
| # Convert math to percentages | |
| probabilities = F.softmax(logits, dim=-1) | |
| confidence_score, predicted_class_id = torch.max(probabilities, dim=-1) | |
| # Map ID back to English | |
| class_id_number = predicted_class_id.item() | |
| predicted_label = self.emotion_map.get(class_id_number, "Unknown") | |
| risk_level = self.determine_risk_level(predicted_label) | |
| return { | |
| "text": text, | |
| "emotion": predicted_label, | |
| "confidence": round(confidence_score.item() * 100, 2), | |
| "risk_level": risk_level | |
| } | |
| # --- Standard Testing Block --- | |
| if __name__ == "__main__": | |
| predictor = MindGuardPredictor() | |
| sample_text = "I have a massive presentation tomorrow and my chest is tight." | |
| result = predictor.predict(sample_text) | |
| print("\n--- Prediction Results ---") | |
| print(f"Emotion: {result['emotion']} ({result['confidence']}%)") | |
| print(f"Risk Level: {result['risk_level']}") |