| import torch |
| import torch.nn as nn |
| from transformers import PreTrainedModel, AutoModel, AutoConfig, PretrainedConfig |
| import transformers |
|
|
|
|
| class DistilBertClassifier(PreTrainedModel): |
|
|
| def __init__(self, bert_config, model_name='distilbert-base-uncased', tokenizer_len=30528, freeze_bert=False): |
|
|
|
|
| super().__init__(bert_config) |
| D_in, H, D_out = 256, 50, 91 |
|
|
| self.bert = AutoModel.from_pretrained(model_name) |
| self.bert.resize_token_embeddings(tokenizer_len) |
| self.classifier = nn.Sequential( |
| nn.GELU(), |
| nn.Linear(self.bert.config.hidden_size, 300), |
| nn.GELU(), |
| nn.Dropout(0.05), |
| nn.Linear(300, 91) |
| ) |
|
|
| if freeze_bert: |
| for param in self.bert.parameters(): |
| param.requires_grad = False |
|
|
| def forward(self, input_ids, attention_mask): |
|
|
| outputs = self.bert(input_ids=input_ids, |
| attention_mask=attention_mask) |
|
|
| last_hidden_state_cls = outputs[0][:, 0, :] |
| logits = self.classifier(last_hidden_state_cls) |
| return logits |
