File size: 405 Bytes
5c5b473
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
from transformers import pipeline

# REAL multi-label model
classifier = pipeline(
    "text-classification",
    model="unitary/unbiased-toxic-roberta",
    top_k=None
)

def predict_toxicity(text: str):
    results = classifier(text)[0]

    scores = {}

    for item in results:
        label = item["label"].lower()
        score = float(item["score"])
        scores[label] = score

    return scores