Views
No views yet
1from transformers import AutoModelForSequenceClassification, AutoTokenizer
2import torch
3
4model = AutoModelForSequenceClassification.from_pretrained("pretrained_model")
5tokenizer = AutoTokenizer.from_pretrained("model_tokenizer")
6
7X_train = ["Why is Owen's retirement from football not mentioned? He hasn't played a game since 2005."]
8batch = tokenizer(X_train, truncation=True, padding='max_length', return_tensors="pt")
9labels = ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"]
10
11with torch.no_grad():
12 outputs = model(**batch)
13 predictions = torch.sigmoid(outputs.logits)*100
14 probs = predictions[0].tolist()
15 for i in range(len(probs)):
16 print(f"{labels[i]}: {round(probs[i], 3)}%")toxic: 0.676%
severe_toxic: 0.001%
obscene: 0.098%
threat: 0.007%
insult: 0.021%
identity_hate: 0.004%