Spaces:
Running
Running
File size: 1,913 Bytes
010ed04 b715cee 010ed04 07840b6 010ed04 39308cb eebda4e 39308cb eebda4e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
class EmotionModel:
def __init__(self):
self.model_name = "joeddav/distilbert-base-uncased-go-emotions-student"
self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
self.model = AutoModelForSequenceClassification.from_pretrained(self.model_name)
self.labels = self.model.config.id2label
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.sigmoid(logits)[0]
return {
self.labels[i]: float(probs[i])
for i in range(len(probs)) if probs[i] > 0.4
}
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 _score_text(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
def predict(self, text, window_size=20, stride=10):
tokens = self.tokenizer.tokenize(text)
if len(tokens) <= window_size:
return self._score_text(text)
scores = []
for i in range(0, len(tokens) - window_size + 1, stride):
window_tokens = tokens[i:i + window_size]
window_text = self.tokenizer.convert_tokens_to_string(window_tokens)
score = self._score_text(window_text)
scores.append(score)
return max(scores, default=0.0) |