from transformers import BertForSequenceClassification, AutoConfig, AutoTokenizer
import torch.nn as nn

class CustomBertForSequenceClassification(BertForSequenceClassification):
    def __init__(self, config):
        super().__init__(config)
        # Replace the default classifier (single linear layer) with a Sequential head
        self.classifier = nn.Sequential(
            nn.Linear(config.hidden_size, config.custom_head_hidden_size),
            nn.ReLU(),
            nn.Linear(config.custom_head_hidden_size, config.num_labels)
        )