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)