File size: 842 Bytes
cd28ab1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
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)
)
|