Views
No views yet
| Класс | Precision | Recall | F1-Score | Threshold |
|---|---|---|---|---|
| profanity | 0.939856 | 0.961855 | 0.950728 | 0.2 |
| threat | 0.956522 | 0.956522 | 0.956522 | 0.4 |
| illegal | 1.000000 | 0.600000 | 0.750000 | 0.5 |
1import torch
2import torch.nn as nn
3from transformers import AutoModel, AutoTokenizer
4
5class MultiTaskToxicityEncoder(nn.Module):
6 def __init__(self, model_name="cointegrated/rubert-tiny2"):
7 super().__init__()
8 self.encoder = AutoModel.from_pretrained(model_name)
9 hidden_size = self.encoder.config.hidden_size
10 self.profanity_head = nn.Linear(hidden_size, 1)
11 self.threat_head = nn.Linear(hidden_size, 1)
12 self.illegal_head = nn.Linear(hidden_size, 1)
13
14 def forward(self, input_ids, attention_mask):
15 outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
16 cls_embedding = outputs.last_hidden_state[:, 0, :]
17 return (
18 self.profanity_head(cls_embedding),
19 self.threat_head(cls_embedding),
20 self.illegal_head(cls_embedding)
21 )
22
23#загрузка модели
24model = MultiTaskToxicityEncoder.from_pretrained("Gloomreach/ru-toxicity-multi-task-encoder")
25tokenizer = AutoTokenizer.from_pretrained("cointegrated/rubert-tiny2")
26
27text = "Пример текста для проверки"
28inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
29
30with torch.no_grad():
31 logits = model(inputs["input_ids"], inputs["attention_mask"])
32 probs = torch.sigmoid(torch.cat(logits, dim=1)).numpy()[0]
33
34print(f"Profanity: {probs[0]:.2%}")
35print(f"Threat: {probs[1]:.2%}")
36print(f"Illegal: {probs[2]:.2%}")