Update model.py

model.py

@@ -18,3 +18,15 @@ class EmotionModel:
             for i in range(len(probs)) if probs[i] > 0.3
         }
 
+class SuicidalIntentModel:
+    def __init__(self):
+        self.model_name = "sentinet/suicidality"
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        self.model = AutoModelForSequenceClassification.from_pretrained(self.model_name)
+
+    def predict(self, text):
+        inputs = self.tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+        with torch.no_grad():
+            logits = self.model(**inputs).logits
+        probs = torch.nn.functional.softmax(logits, dim=1)
+        return float(probs[0][1])  # Probability of suicidal intent