Spaces:
Sleeping
Sleeping
| import torch.nn as nn | |
| import torchvision.models as models | |
| import torch | |
| from transformers import AutoTokenizer, AutoModel, AutoConfig | |
| class Classifier(nn.Module): | |
| def __init__(self, input_size = 512, output_sizes = [1], dropout_rate = 0.1): | |
| super(Classifier, self).__init__() | |
| self.hs_head = nn.Sequential( | |
| nn.Dropout(dropout_rate), | |
| nn.Linear(input_size, output_sizes[0]) | |
| ) | |
| self.abusive_head = nn.Sequential( | |
| nn.Dropout(dropout_rate), | |
| nn.Linear(input_size, output_sizes[1]) | |
| ) | |
| self.target_head = nn.Sequential( | |
| nn.Dropout(dropout_rate), | |
| nn.Linear(input_size, output_sizes[2]) | |
| ) | |
| self.strength_head = nn.Sequential( | |
| nn.Dropout(dropout_rate), | |
| nn.Linear(input_size, output_sizes[3]) | |
| ) | |
| self.type_head = nn.Sequential( | |
| nn.Dropout(dropout_rate), | |
| nn.Linear(input_size, output_sizes[4]) | |
| ) | |
| def forward(self, input): | |
| return self.hs_head(input), self.abusive_head(input), self.target_head(input), \ | |
| self.strength_head(input), self.type_head(input) |