| from transformers import BertTokenizer, BertModel | |
| from .configuration_my_bert_classifier import MyBertClassifierConfig | |
| from torch import nn | |
| from transformers.modeling_utils import PreTrainedModel | |
| class MyBertClassifier(PreTrainedModel): | |
| config_class = MyBertClassifierConfig | |
| def __init__(self, config): | |
| super(MyBertClassifier, self).__init__(config) | |
| self.bert = BertModel.from_pretrained('bert-base-cased') | |
| self.dropout = nn.Dropout(0.5) | |
| self.linear = nn.Linear(768, 5) | |
| self.relu = nn.ReLU() | |
| def forward(self, input_id, mask): | |
| _, pooled_output = self.bert(input_ids= input_id, attention_mask=mask,return_dict=False) | |
| dropout_output = self.dropout(pooled_output) | |
| linear_output = self.linear(dropout_output) | |
| final_layer = self.relu(linear_output) | |
| return final_layer |