| import torch | |
| import torch.nn as nn | |
| from transformers import AutoModel | |
| class MultiTaskToxicityEncoder(nn.Module): | |
| def __init__(self, base_model="cointegrated/rubert-tiny2"): | |
| super().__init__() | |
| self.encoder = AutoModel.from_pretrained(base_model) | |
| hidden_size = self.encoder.config.hidden_size | |
| self.profanity_head = nn.Linear(hidden_size, 1) | |
| self.threat_head = nn.Linear(hidden_size, 1) | |
| self.illegal_head = nn.Linear(hidden_size, 1) | |
| def forward(self, input_ids, attention_mask): | |
| outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask) | |
| cls_embedding = outputs.last_hidden_state[:, 0, :] | |
| return ( | |
| self.profanity_head(cls_embedding), | |
| self.threat_head(cls_embedding), | |
| self.illegal_head(cls_embedding) | |
| ) | |