File size: 1,020 Bytes
6826247 |
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 |
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
class ToxicityDetector:
def __init__(self):
model_name = "unitary/toxic-bert"
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
self.labels = [
"toxicity", "severe_toxicity", "obscene", "threat",
"insult", "identity_attack", "sexual_explicit"
]
def detect(self, prompt):
inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True)
with torch.no_grad():
outputs = self.model(**inputs)
scores = torch.sigmoid(outputs.logits).squeeze().tolist()
results = [
{"label": label, "score": round(score, 3)}
for label, score in zip(self.labels, scores)
if score > 0.3
]
return {
"label": "Toxic" if results else "Safe",
"details": results
}
|