Capstone-lpx/custom-mybertclassifier

1from transformers import BertTokenizer, BertModel2from .configuration_my_bert_classifier import MyBertClassifierConfig3from torch import nn4from transformers.modeling_utils import PreTrainedModel5 6class MyBertClassifier(PreTrainedModel):7    8    config_class = MyBertClassifierConfig9 10    def __init__(self, config):11 12        super(MyBertClassifier, self).__init__(config)13 14        self.bert = BertModel.from_pretrained('bert-base-cased')15        self.dropout = nn.Dropout(0.5)16        self.linear = nn.Linear(768, 5)17        self.relu = nn.ReLU()18 19    def forward(self, input_id, mask):20 21        _, pooled_output = self.bert(input_ids= input_id, attention_mask=mask,return_dict=False)22        dropout_output = self.dropout(pooled_output)23        linear_output = self.linear(dropout_output)24        final_layer = self.relu(linear_output)25 26        return final_layer