| import torch |
| import torch.nn as nn |
| from transformers import BertForTokenClassification |
|
|
|
|
| class FiveOClassifier(nn.Module): |
| def __init__(self, clf_hidden_size, clf_num_labels): |
| super(FiveOClassifier, self).__init__() |
| self.dense1 = nn.Linear(clf_hidden_size, clf_hidden_size // 2) |
| self.activation1 = nn.ReLU() |
| self.dropout1 = nn.Dropout(p=0.1) |
| self.dense2 = nn.Linear(clf_hidden_size // 2, clf_hidden_size // 4) |
| self.activation2 = nn.ReLU() |
| self.dropout2 = nn.Dropout(p=0.1) |
| self.output_layer = nn.Linear(clf_hidden_size // 4, clf_num_labels) |
|
|
| def forward(self, clf_input): |
| x = self.dense1(clf_input) |
| x = self.activation1(x) |
| x = self.dropout1(x) |
| x = self.dense2(x) |
| x = self.activation2(x) |
| x = self.dropout2(x) |
| x = self.output_layer(x) |
| return x |
|
|
|
|
| class BertForTokenClassificationWithFiveO(BertForTokenClassification): |
| def __init__(self, config): |
| super().__init__(config) |
| self.num_labels = config.num_labels |
| self.classifier = FiveOClassifier(config.hidden_size, config.num_labels) |
| self.init_weights() |
|
|
| @classmethod |
| def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): |
| model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) |
| model.check_classifier_initialization() |
| return model |
|
|
| def check_classifier_initialization(self): |
| |
| def is_randomly_initialized(tensor): |
| return torch.abs(tensor.mean()) < 1e-3 < tensor.std() < 1e-1 |
|
|
| classifier_weights = [ |
| self.classifier.dense1.weight, |
| self.classifier.dense1.bias, |
| self.classifier.dense2.weight, |
| self.classifier.dense2.bias, |
| self.classifier.output_layer.weight, |
| self.classifier.output_layer.bias |
| ] |
|
|
|
|
| def freeze_bert(self): |
| """Freezes the BERT layers to prevent their parameters from being updated during training.""" |
| for param in self.bert.parameters(): |
| param.requires_grad = False |
| print("BERT layers frozen.") |
|
|
| def unfreeze_bert(self): |
| """Unfreezes the BERT layers to allow their parameters to be updated during training.""" |
| for param in self.bert.parameters(): |
| param.requires_grad = True |
| print("BERT layers unfrozen.") |
|
|
