File size: 1,350 Bytes
5de4467 78b8ecc 5de4467 78b8ecc 5de4467 78b8ecc 5de4467 78b8ecc 5de4467 | 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 30 31 32 33 34 35 | import torch
import torch.nn as nn
from transformers import PreTrainedModel, AutoModel, AutoConfig
from .configuration_multitask_toxicity import MultiTaskToxicityConfig
class MultiTaskToxicityEncoder(PreTrainedModel):
config_class = MultiTaskToxicityConfig
def __init__(self, config):
super().__init__(config)
self.config = config
base_config = AutoConfig.from_pretrained(config.base_model_name)
self.encoder = AutoModel.from_config(base_config)
hidden_size = base_config.hidden_size
self.dropout = nn.Dropout(config.dropout_prob)
self.profanity_head = nn.Linear(hidden_size, 1)
self.threat_head = nn.Linear(hidden_size, 1)
self.illegal_head = nn.Linear(hidden_size, 1)
self.post_init()
def forward(self, input_ids, attention_mask=None, **kwargs):
outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
cls_embedding = outputs.last_hidden_state[:, 0, :]
cls_embedding = self.dropout(cls_embedding)
profanity_logit = self.profanity_head(cls_embedding)
threat_logit = self.threat_head(cls_embedding)
illegal_logit = self.illegal_head(cls_embedding)
return profanity_logit, threat_logit, illegal_logit
|